Selaa lähdekoodia

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
tags/v0.20.1
Stephen Hu 2 kuukautta sitten
vanhempi
commit
0a303d9ae1
No account linked to committer's email address
1 muutettua tiedostoa jossa 4 lisäystä ja 1 poistoa
  1. 4
    1
      rag/llm/cv_model.py

+ 4
- 1
rag/llm/cv_model.py Näytä tiedosto

@@ -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…
Peruuta
Tallenna