llm-gemini 0.4.2__tar.gz → 0.5__tar.gz

{llm_gemini-0.4.2 → llm_gemini-0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-gemini
-Version: 0.4.2
+Version: 0.5
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0
@@ -11,7 +11,7 @@ Project-URL: CI, https://github.com/simonw/llm-gemini/actions
 Classifier: License :: OSI Approved :: Apache Software License
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: llm>=0.18
+Requires-Dist: llm>=0.19
 Requires-Dist: httpx
 Requires-Dist: ijson
 Provides-Extra: test

{llm_gemini-0.4.2 → llm_gemini-0.5}/llm_gemini.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-gemini
-Version: 0.4.2
+Version: 0.5
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0
@@ -11,7 +11,7 @@ Project-URL: CI, https://github.com/simonw/llm-gemini/actions
 Classifier: License :: OSI Approved :: Apache Software License
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: llm>=0.18
+Requires-Dist: llm>=0.19
 Requires-Dist: httpx
 Requires-Dist: ijson
 Provides-Extra: test

{llm_gemini-0.4.2 → llm_gemini-0.5}/llm_gemini.egg-info/requires.txt

@@ -1,4 +1,4 @@
-llm>=0.18
+llm>=0.19
 httpx
 ijson
 

{llm_gemini-0.4.2 → llm_gemini-0.5}/llm_gemini.py

@@ -211,6 +211,19 @@ class _SharedGemini:
             return f'```\n{part["codeExecutionResult"]["output"].strip()}\n```\n'
         return ""
 
+    def set_usage(self, response):
+        try:
+            usage = response.response_json[-1].pop("usageMetadata")
+            input_tokens = usage.pop("promptTokenCount", None)
+            output_tokens = usage.pop("candidatesTokenCount", None)
+            usage.pop("totalTokenCount", None)
+            if input_tokens is not None:
+                response.set_usage(
+                    input=input_tokens, output=output_tokens, details=usage or None
+                )
+        except (IndexError, KeyError):
+            pass
+
 
 class GeminiPro(_SharedGemini, llm.Model):
     def execute(self, prompt, stream, response, conversation):
@@ -242,6 +255,7 @@ class GeminiPro(_SharedGemini, llm.Model):
                     gathered.append(event)
                     events.clear()
         response.response_json = gathered
+        self.set_usage(response)
 
 
 class AsyncGeminiPro(_SharedGemini, llm.AsyncModel):
@@ -275,6 +289,7 @@ class AsyncGeminiPro(_SharedGemini, llm.AsyncModel):
                         gathered.append(event)
                         events.clear()
         response.response_json = gathered
+        self.set_usage(response)
 
 
 @llm.hookimpl

{llm_gemini-0.4.2 → llm_gemini-0.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "llm-gemini"
-version = "0.4.2"
+version = "0.5"
 description = "LLM plugin to access Google's Gemini family of models"
 readme = "README.md"
 authors = [{name = "Simon Willison"}]
@@ -9,7 +9,7 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License"
 ]
 dependencies = [
-    "llm>=0.18",
+    "llm>=0.19",
     "httpx",
     "ijson"
 ]

{llm_gemini-0.4.2 → llm_gemini-0.5}/tests/test_gemini.py

@@ -17,10 +17,15 @@ async def test_prompt():
             "candidates": [
                 {"content": {"parts": [{"text": "Percy"}], "role": "model"}}
             ],
-            "usageMetadata": {"promptTokenCount": 10, "totalTokenCount": 10},
             "modelVersion": "gemini-1.5-flash-002",
         }
     ]
+    assert response.token_details is None
+    assert response.input_tokens == 10
+    # Not sure why our pytest-recording setup doesn't report output tokens
+    # https://github.com/simonw/llm-gemini/issues/25#issuecomment-2487464339
+    assert response.output_tokens is None
+
     # And try it async too
     async_model = llm.get_async_model("gemini-1.5-flash-latest")
     async_model.key = async_model.key or GEMINI_API_KEY