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/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):