pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: use the same implement for total token count from res (#10197)

Browse Source 
### What problem does this PR solve?
use the same implement for total token count from res

### Type of change

- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-09-22 17:17:06 +08:00
committed by GitHub
parent ca9f30e1a1
commit 94dbd4aac9
4 changed files with 20 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
rag/llm/chat_model.py
View File

@ -36,7 +36,7 @@ from zhipuai import ZhipuAI
from rag.llm import FACTORY_DEFAULT_BASE_URL, LITELLM_PROVIDER_PREFIX, SupportedLiteLLMProvider
from rag.nlp import is_chinese, is_english
from rag.utils import num_tokens_from_string
from rag.utils import num_tokens_from_string, total_token_count_from_response
# Error message constants
@ -445,15 +445,7 @@ class Base(ABC):
yield total_tokens
def total_token_count(self, resp):
try:
return resp.usage.total_tokens
except Exception:
pass
try:
return resp["usage"]["total_tokens"]
except Exception:
pass
return 0
return total_token_count_from_response(resp)
def _calculate_dynamic_ctx(self, history):
"""Calculate dynamic context window size"""

12
rag/llm/embedding_model.py
View File

@ -33,7 +33,7 @@ from zhipuai import ZhipuAI
from api import settings
from api.utils.file_utils import get_home_cache_dir
from api.utils.log_utils import log_exception
from rag.utils import num_tokens_from_string, truncate
from rag.utils import num_tokens_from_string, truncate, total_token_count_from_response
class Base(ABC):
@ -52,15 +52,7 @@ class Base(ABC):
raise NotImplementedError("Please implement encode method!")
def total_token_count(self, resp):
try:
return resp.usage.total_tokens
except Exception:
pass
try:
return resp["usage"]["total_tokens"]
except Exception:
pass
return 0
return total_token_count_from_response(resp)
class DefaultEmbedding(Base):

15
rag/llm/rerank_model.py
View File

@ -30,7 +30,7 @@ from yarl import URL
from api import settings
from api.utils.file_utils import get_home_cache_dir
from api.utils.log_utils import log_exception
from rag.utils import num_tokens_from_string, truncate
from rag.utils import num_tokens_from_string, truncate, total_token_count_from_response
class Base(ABC):
def __init__(self, key, model_name, **kwargs):
@ -44,18 +44,7 @@ class Base(ABC):
raise NotImplementedError("Please implement encode method!")
def total_token_count(self, resp):
if hasattr(resp, "usage") and hasattr(resp.usage, "total_tokens"):
try:
return resp.usage.total_tokens
except Exception:
pass
if 'usage' in resp and 'total_tokens' in resp['usage']:
try:
return resp["usage"]["total_tokens"]
except Exception:
pass
return 0
return total_token_count_from_response(resp)
class DefaultRerank(Base):