llm-gemini 0.4.1__py3-none-any.whl → 0.5a0__py3-none-any.whl

{llm_gemini-0.4.1.dist-info → llm_gemini-0.5a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-gemini
-Version: 0.4.1
+Version: 0.5a0
 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.19a0
 Requires-Dist: httpx
 Requires-Dist: ijson
 Provides-Extra: test
@@ -157,3 +157,16 @@ To run the tests:
 ```bash
 pytest
 ```
+
+This project uses [pytest-recording](https://github.com/kiwicom/pytest-recording) to record Gemini API responses for the tests.
+
+If you add a new test that calls the API you can capture the API response like this:
+```bash
+PYTEST_GEMINI_API_KEY="$(llm keys get gemini)" pytest --record-mode once
+```
+You will need to have stored a valid Gemini API key using this command first:
+```bash
+llm keys set gemini
+# Paste key here
+```
+

llm_gemini-0.5a0.dist-info/RECORD

@@ -0,0 +1,7 @@
+llm_gemini.py,sha256=bALsCHsTTFaj2mdfZe2hOIc7aAxDFwNMYb7keyy4-20,11268
+llm_gemini-0.5a0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+llm_gemini-0.5a0.dist-info/METADATA,sha256=ON3jk8T7hK9Zp8IF9Z4TxCL-CvN94Eko8Xnb0X-7HhA,5294
+llm_gemini-0.5a0.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+llm_gemini-0.5a0.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
+llm_gemini-0.5a0.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
+llm_gemini-0.5a0.dist-info/RECORD,,

llm_gemini.py

@@ -210,6 +210,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):
@@ -241,6 +254,7 @@ class GeminiPro(_SharedGemini, llm.Model):
                     gathered.append(event)
                     events.clear()
         response.response_json = gathered
+        self.set_usage(response)
 
 
 class AsyncGeminiPro(_SharedGemini, llm.AsyncModel):
@@ -274,6 +288,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.1.dist-info/RECORD

@@ -1,7 +0,0 @@
-llm_gemini.py,sha256=_7yQ14ffRpyK3ChAOc6M2ufylg2kUKTQK4C3VKOiUgM,10672
-llm_gemini-0.4.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-llm_gemini-0.4.1.dist-info/METADATA,sha256=OhySg0uA_de-JRZrYIJFZu_uSYVKXEGs4W_x3LkQUvE,4864
-llm_gemini-0.4.1.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-llm_gemini-0.4.1.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
-llm_gemini-0.4.1.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
-llm_gemini-0.4.1.dist-info/RECORD,,