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: Improve how to get total token count for AnthropicCV (#10658)

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

 Improve how to get total token count for AnthropicCV

### Type of change

- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-10-29 09:41:15 +08:00
committed by GitHub
parent e86bd723d1
commit d86d7061ea
2 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
rag/llm/cv_model.py
View File

@ -797,8 +797,7 @@ class NvidiaCV(Base):
try: try:
response = self._request(self._form_history(system, history, images), gen_conf) response = self._request(self._form_history(system, history, images), gen_conf)
cnt = response["choices"][0]["message"]["content"] cnt = response["choices"][0]["message"]["content"]
if "usage" in response and "total_tokens" in response["usage"]: total_tokens += total_token_count_from_response(response)
total_tokens += total_token_count_from_response(response)
for resp in cnt: for resp in cnt:
yield resp yield resp
except Exception as e: except Exception as e:
@ -847,7 +846,7 @@ class AnthropicCV(Base):
prompt = self.prompt(b64, prompt if prompt else vision_llm_describe_prompt()) prompt = self.prompt(b64, prompt if prompt else vision_llm_describe_prompt())
response = self.client.messages.create(model=self.model_name, max_tokens=self.max_tokens, messages=prompt) response = self.client.messages.create(model=self.model_name, max_tokens=self.max_tokens, messages=prompt)
return response["content"][0]["text"].strip(), response["usage"]["input_tokens"] + response["usage"]["output_tokens"] return response["content"][0]["text"].strip(), total_token_count_from_response(response)
def _clean_conf(self, gen_conf): def _clean_conf(self, gen_conf):
if "presence_penalty" in gen_conf: if "presence_penalty" in gen_conf:
@ -874,7 +873,7 @@ class AnthropicCV(Base):
ans += "...\nFor the content length reason, it stopped, continue?" if is_english([ans]) else "······\n由于长度的原因，回答被截断了，要继续吗？" ans += "...\nFor the content length reason, it stopped, continue?" if is_english([ans]) else "······\n由于长度的原因，回答被截断了，要继续吗？"
return ( return (
ans, ans,
response["usage"]["input_tokens"] + response["usage"]["output_tokens"], total_token_count_from_response(response),
) )
except Exception as e: except Exception as e:
return ans + "\n**ERROR**: " + str(e), 0 return ans + "\n**ERROR**: " + str(e), 0

6
rag/utils/__init__.py
View File

@ -63,6 +63,12 @@ def total_token_count_from_response(resp):
return resp["usage"]["total_tokens"] return resp["usage"]["total_tokens"]
except Exception: except Exception:
pass pass
if 'usage' in resp and 'input_tokens' in resp['usage'] and 'output_tokens' in resp['usage']:
try:
return resp["usage"]["input_tokens"] + resp["usage"]["output_tokens"]
except Exception:
pass
return 0 return 0