Output token usage in raw data for Google Gemini LLMs (#16313)

run-llama · Oct 2, 2024 · 0d12739 · 0d12739
1 parent 4cd774a
commit 0d12739
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/CHANGELOG.md b/llama-index-integrations/llms/llama-index-llms-gemini/CHANGELOG.md
@@ -0,0 +1,5 @@
+# CHANGELOG — llama-index-llms-gemini
+
+## [0.3.6]
+
+- Output token usage in `raw` data.
diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/utils.py b/llama-index-integrations/llms/llama-index-llms-gemini/llama_index/llms/gemini/utils.py
@@ -42,6 +42,10 @@ def completion_from_gemini_response(
         **(type(top_candidate).to_dict(top_candidate)),
         **(type(response.prompt_feedback).to_dict(response.prompt_feedback)),
     }
+    if response.usage_metadata:
+        raw["usage_metadata"] = type(response.usage_metadata).to_dict(
+            response.usage_metadata
+        )
     return CompletionResponse(text=response.text, raw=raw)
 
 
@@ -58,6 +62,10 @@ def chat_from_gemini_response(
         **(type(top_candidate).to_dict(top_candidate)),
         **(type(response.prompt_feedback).to_dict(response.prompt_feedback)),
     }
+    if response.usage_metadata:
+        raw["usage_metadata"] = type(response.usage_metadata).to_dict(
+            response.usage_metadata
+        )
     role = ROLES_FROM_GEMINI[top_candidate.content.role]
     return ChatResponse(message=ChatMessage(role=role, content=response.text), raw=raw)
 

diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-gemini/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-gemini"
 readme = "README.md"
-version = "0.3.5"
+version = "0.3.6"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"