Refactor:Improve the chat stream logic for NvidiaCV (#9242)
### What problem does this PR solve? Improve the chat stream logic for NvidiaCV ### Type of change - [x] Refactoring
This commit is contained in:
parent
98a83543e8
commit
0a303d9ae1
1 changed files with 4 additions and 1 deletions
|
|
@ -623,15 +623,18 @@ class NvidiaCV(Base):
|
|||
return "**ERROR**: " + str(e), 0
|
||||
|
||||
def chat_streamly(self, system, history, gen_conf, images=[], **kwargs):
|
||||
total_tokens = 0
|
||||
try:
|
||||
response = self._request(self._form_history(system, history, images), gen_conf)
|
||||
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:
|
||||
yield resp
|
||||
except Exception as e:
|
||||
yield "\n**ERROR**: " + str(e)
|
||||
|
||||
yield response["usage"]["total_tokens"]
|
||||
yield total_tokens
|
||||
|
||||
|
||||
class AnthropicCV(Base):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue