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 the chat stream logic for NvidiaCV (#9242)

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

Improve the chat stream logic for NvidiaCV

### Type of change


- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-08-05 17:47:00 +08:00
committed by GitHub
parent 98a83543e8
commit 0a303d9ae1
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
rag/llm/cv_model.py
View File

@ -623,15 +623,18 @@ class NvidiaCV(Base):
return "**ERROR**: " + str(e), 0 return "**ERROR**: " + str(e), 0
def chat_streamly(self, system, history, gen_conf, images=[], **kwargs): def chat_streamly(self, system, history, gen_conf, images=[], **kwargs):
total_tokens = 0
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 += response["usage"]["total_tokens"]
for resp in cnt: for resp in cnt:
yield resp yield resp
except Exception as e: except Exception as e:
yield "\n**ERROR**: " + str(e) yield "\n**ERROR**: " + str(e)
yield response["usage"]["total_tokens"] yield total_tokens
class AnthropicCV(Base): class AnthropicCV(Base):