mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-21 05:16:54 +08:00
Refactor: improve cohere calculate total counts (#12007)
### What problem does this PR solve? improve cohere calculate total counts ### Type of change - [x] Refactoring
This commit is contained in:
@ -56,6 +56,12 @@ def total_token_count_from_response(resp):
|
||||
except Exception:
|
||||
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']:
|
||||
try:
|
||||
return resp["usage"]["total_tokens"]
|
||||
|
||||
@ -639,7 +639,7 @@ class CoHereEmbed(Base):
|
||||
)
|
||||
try:
|
||||
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:
|
||||
log_exception(_e, res)
|
||||
raise Exception(f"Error: {res}")
|
||||
@ -653,7 +653,7 @@ class CoHereEmbed(Base):
|
||||
embedding_types=["float"],
|
||||
)
|
||||
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:
|
||||
log_exception(_e, res)
|
||||
raise Exception(f"Error: {res}")
|
||||
|
||||
Reference in New Issue
Block a user