Fix: Token counter expecting response.raw as dict, got ChatCompletion…

…Chunk (#14937)
run-llama · Jul 24, 2024 · 3be95a2 · 3be95a2
1 parent 3a35a6e
commit 3be95a2
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/llama-index-core/llama_index/core/callbacks/token_counting.py b/llama-index-core/llama_index/core/callbacks/token_counting.py
@@ -53,11 +53,14 @@ def get_llm_token_counts(
             response_tokens = 0
 
             if response is not None and response.raw is not None:
-                usage = response.raw.get("usage", None)
+                if isinstance(response.raw, dict):
+                    raw_dict = response.raw
+                else:
+                    raw_dict = response.raw.model_dump()
+
+                usage = raw_dict.get("usage", None)
 
                 if usage is not None:
-                    if not isinstance(usage, dict):
-                        usage = dict(usage)
                     messages_tokens = usage.get("prompt_tokens", 0)
                     response_tokens = usage.get("completion_tokens", 0)