llm-gemini 0.13a0__py3-none-any.whl → 0.13.1__py3-none-any.whl

{llm_gemini-0.13a0.dist-info → llm_gemini-0.13.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: llm-gemini
-Version: 0.13a0
+Version: 0.13.1
 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.23a0
+Requires-Dist: llm>=0.23
 Requires-Dist: httpx
 Requires-Dist: ijson
 Provides-Extra: test

llm_gemini-0.13.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+llm_gemini.py,sha256=JWaRtT8vJzdxrRbCAgoNrKqes2df_T4gqqeDmiw0oRI,14370
+llm_gemini-0.13.1.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+llm_gemini-0.13.1.dist-info/METADATA,sha256=S87OFYnm9K0EVvJIRrJoZT-J0gxhPA4BrnIGkVD8q8c,7016
+llm_gemini-0.13.1.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+llm_gemini-0.13.1.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
+llm_gemini-0.13.1.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
+llm_gemini-0.13.1.dist-info/RECORD,,

llm_gemini.py

@@ -65,8 +65,16 @@ def register_models(register):
     ]:
         can_google_search = model_id in GOOGLE_SEARCH_MODELS
         register(
-            GeminiPro(model_id, can_google_search=can_google_search),
-            AsyncGeminiPro(model_id, can_google_search=can_google_search),
+            GeminiPro(
+                model_id,
+                can_google_search=can_google_search,
+                can_schema="flash-thinking" not in model_id,
+            ),
+            AsyncGeminiPro(
+                model_id,
+                can_google_search=can_google_search,
+                can_schema="flash-thinking" not in model_id,
+            ),
         )
 
 
@@ -82,7 +90,7 @@ def resolve_type(attachment):
 
 def cleanup_schema(schema):
     "Gemini supports only a subset of JSON schema"
-    keys_to_remove = ("$schema", "additionalProperties")
+    keys_to_remove = ("$schema", "additionalProperties", "title")
     # Recursively remove them
     if isinstance(schema, dict):
         for key in keys_to_remove:
@@ -186,9 +194,10 @@ class _SharedGemini:
             default=None,
         )
 
-    def __init__(self, model_id, can_google_search=False):
+    def __init__(self, model_id, can_google_search=False, can_schema=False):
         self.model_id = model_id
         self.can_google_search = can_google_search
+        self.supports_schema = can_schema
         if can_google_search:
             self.Options = self.OptionsWithGoogleSearch
 
@@ -279,9 +288,17 @@ class _SharedGemini:
             return f'```\n{part["codeExecutionResult"]["output"].strip()}\n```\n'
         return ""
 
+    def process_candidates(self, candidates):
+        # We only use the first candidate
+        for part in candidates[0]["content"]["parts"]:
+            yield self.process_part(part)
+
     def set_usage(self, response):
         try:
-            usage = response.response_json[-1].pop("usageMetadata")
+            # Don't record the "content" key from that last candidate
+            for candidate in response.response_json["candidates"]:
+                candidate.pop("content", None)
+            usage = response.response_json.pop("usageMetadata")
             input_tokens = usage.pop("promptTokenCount", None)
             output_tokens = usage.pop("candidatesTokenCount", None)
             usage.pop("totalTokenCount", None)
@@ -311,17 +328,16 @@ class GeminiPro(_SharedGemini, llm.KeyModel):
             for chunk in http_response.iter_bytes():
                 coro.send(chunk)
                 if events:
-                    event = events[0]
-                    if isinstance(event, dict) and "error" in event:
-                        raise llm.ModelError(event["error"]["message"])
-                    try:
-                        part = event["candidates"][0]["content"]["parts"][0]
-                        yield self.process_part(part)
-                    except KeyError:
-                        yield ""
-                    gathered.append(event)
+                    for event in events:
+                        if isinstance(event, dict) and "error" in event:
+                            raise llm.ModelError(event["error"]["message"])
+                        try:
+                            yield from self.process_candidates(event["candidates"])
+                        except KeyError:
+                            yield ""
+                        gathered.append(event)
                     events.clear()
-        response.response_json = gathered
+        response.response_json = gathered[-1]
         self.set_usage(response)
 
 
@@ -344,17 +360,19 @@ class AsyncGeminiPro(_SharedGemini, llm.AsyncKeyModel):
                 async for chunk in http_response.aiter_bytes():
                     coro.send(chunk)
                     if events:
-                        event = events[0]
-                        if isinstance(event, dict) and "error" in event:
-                            raise llm.ModelError(event["error"]["message"])
-                        try:
-                            part = event["candidates"][0]["content"]["parts"][0]
-                            yield self.process_part(part)
-                        except KeyError:
-                            yield ""
-                        gathered.append(event)
+                        for event in events:
+                            if isinstance(event, dict) and "error" in event:
+                                raise llm.ModelError(event["error"]["message"])
+                            try:
+                                for chunk in self.process_candidates(
+                                    event["candidates"]
+                                ):
+                                    yield chunk
+                            except KeyError:
+                                yield ""
+                            gathered.append(event)
                         events.clear()
-        response.response_json = gathered
+        response.response_json = gathered[-1]
         self.set_usage(response)
 
 

llm_gemini-0.13a0.dist-info/RECORD

@@ -1,7 +0,0 @@
-llm_gemini.py,sha256=M4_OIzaF3ytfrCdQ_md9sS6ViDN38JlP9aaykY5Ct0E,13634
-llm_gemini-0.13a0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-llm_gemini-0.13a0.dist-info/METADATA,sha256=_SbfjPGKz4M4Wfvd6p8zyx37qOF99R-aTBxGQ0lnGHY,7018
-llm_gemini-0.13a0.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-llm_gemini-0.13a0.dist-info/entry_points.txt,sha256=n544bpgUPIBc5l_cnwsTxPc3gMGJHPtAyqBNp-CkMWk,26
-llm_gemini-0.13a0.dist-info/top_level.txt,sha256=WUQmG6_2QKbT_8W4HH93qyKl_0SUteL4Ra6_PhyNGKU,11
-llm_gemini-0.13a0.dist-info/RECORD,,