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

Refactor: improve cohere calculate total counts (#12007)

Browse Source 
### What problem does this PR solve?

improve cohere calculate total counts

### Type of change


- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-12-18 10:04:28 +08:00
committed by GitHub
parent 4dd8cdc38b
commit a63dcfed6f
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
common/token_utils.py
View File

@ -56,6 +56,12 @@ def total_token_count_from_response(resp):
except Exception: except Exception:
pass pass
if hasattr(resp, "meta") and hasattr(resp.meta, "billed_units") and hasattr(resp.meta.billed_units, "input_tokens"):
try:
return resp.meta.billed_units.input_tokens
except Exception:
pass
if isinstance(resp, dict) and 'usage' in resp and 'total_tokens' in resp['usage']: if isinstance(resp, dict) and 'usage' in resp and 'total_tokens' in resp['usage']:
try: try:
return resp["usage"]["total_tokens"] return resp["usage"]["total_tokens"]

4
rag/llm/embedding_model.py
View File

@ -639,7 +639,7 @@ class CoHereEmbed(Base):
) )
try: try:
ress.extend([d for d in res.embeddings.float]) ress.extend([d for d in res.embeddings.float])
token_count += res.meta.billed_units.input_tokens token_count += total_token_count_from_response(res)
except Exception as _e: except Exception as _e:
log_exception(_e, res) log_exception(_e, res)
raise Exception(f"Error: {res}") raise Exception(f"Error: {res}")
@ -653,7 +653,7 @@ class CoHereEmbed(Base):
embedding_types=["float"], embedding_types=["float"],
) )
try: try:
return np.array(res.embeddings.float[0]), int(res.meta.billed_units.input_tokens) return np.array(res.embeddings.float[0]), int(total_token_count_from_response(res))
except Exception as _e: except Exception as _e:
log_exception(_e, res) log_exception(_e, res)
raise Exception(f"Error: {res}") raise Exception(f"Error: {res}")