[Bug]: When use OpenAI chat model , raise ERROR: 'CompletionUsage' ob…

…ject has no attribute 'get' #2948
infiniflow · Oct 22, 2024 · f27c7b6 · f27c7b6
1 parent 6a4858a
commit f27c7b6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
@@ -73,7 +73,7 @@ def chat_streamly(self, system, history, gen_conf):
                             + num_tokens_from_string(resp.choices[0].delta.content)
                     )
                     if not hasattr(resp, "usage") or not resp.usage
-                    else resp.usage.get("total_tokens", total_tokens)
+                    else dict(resp.usage).get('total_tokens', total_tokens)
                 )
                 if resp.choices[0].finish_reason == "length":
                     ans += "...\nFor the content length reason, it stopped, continue?" if is_english(