From 0d1273913fd254031ab899379944cd8bfd7bf8ba Mon Sep 17 00:00:00 2001
From: Yannic Spreen-Ledebur <35889034+spreeni@users.noreply.github.com>
Date: Wed, 2 Oct 2024 19:36:58 +0200
Subject: [PATCH] Output token usage in `raw` data for Google Gemini LLMs
 (#16313)

---
 .../llms/llama-index-llms-gemini/CHANGELOG.md             | 5 +++++
 .../llama_index/llms/gemini/utils.py                      | 8 ++++++++
 .../llms/llama-index-llms-gemini/pyproject.toml           | 2 +-
 3 files changed, 14 insertions(+), 1 deletion(-)
 create mode 100644 llama-index-integrations/llms/llama-index-llms-gemini/CHANGELOG.md

diff --git a/llama-index-integrations/llms/llama-index-llms-gemini/CHANGELOG.md b/llama-index-integrations/llms/llama-index-llms-gemini/CHANGELOG.md
new file mode 100644
index 0000000000000..4865565d12602
--- /dev/null
+++ 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
index d7c5ae0ca9026..1dd9fa655d2a9 100644
--- 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
index f2c8ed2bae3d6..92e4cfb6caaa1 100644
--- 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"