make gemini robust

infiniflow · Oct 25, 2024 · f49382a · f49382a
1 parent 648f8e8
commit f49382a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
@@ -780,10 +780,11 @@ def chat_streamly(self, system, history, gen_conf):
                 ans += resp.text
                 yield ans
 
+            yield response._chunks[-1].usage_metadata.total_token_count
         except Exception as e:
             yield ans + "\n**ERROR**: " + str(e)
 
-        yield response._chunks[-1].usage_metadata.total_token_count
+        yield 0
 
 
 class GroqChat: