llm-gemini 0.13__tar.gz → 0.13.1__tar.gz

{llm_gemini-0.13 → llm_gemini-0.13.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: llm-gemini
-Version: 0.13
+Version: 0.13.1
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0

{llm_gemini-0.13 → llm_gemini-0.13.1}/llm_gemini.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: llm-gemini
-Version: 0.13
+Version: 0.13.1
 Summary: LLM plugin to access Google's Gemini family of models
 Author: Simon Willison
 License: Apache-2.0

{llm_gemini-0.13 → llm_gemini-0.13.1}/llm_gemini.py

@@ -90,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:
@@ -288,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)
@@ -320,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)
 
 
@@ -353,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.13 → llm_gemini-0.13.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "llm-gemini"
-version = "0.13"
+version = "0.13.1"
 description = "LLM plugin to access Google's Gemini family of models"
 readme = "README.md"
 authors = [{name = "Simon Willison"}]

llm_gemini-0.13.1/tests/test_gemini.py

@@ -0,0 +1,104 @@
+import llm
+import nest_asyncio
+import json
+import os
+import pytest
+import pydantic
+
+nest_asyncio.apply()
+
+GEMINI_API_KEY = os.environ.get("PYTEST_GEMINI_API_KEY", None) or "gm-..."
+
+
+@pytest.mark.vcr
+@pytest.mark.asyncio
+async def test_prompt():
+    model = llm.get_model("gemini-1.5-flash-latest")
+    response = model.prompt("Name for a pet pelican, just the name", key=GEMINI_API_KEY)
+    assert str(response) == "Percy\n"
+    assert response.response_json == {
+        "candidates": [
+            {
+                "finishReason": "STOP",
+                "safetyRatings": [
+                    {
+                        "category": "HARM_CATEGORY_HATE_SPEECH",
+                        "probability": "NEGLIGIBLE",
+                    },
+                    {
+                        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                        "probability": "NEGLIGIBLE",
+                    },
+                    {
+                        "category": "HARM_CATEGORY_HARASSMENT",
+                        "probability": "NEGLIGIBLE",
+                    },
+                    {
+                        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                        "probability": "NEGLIGIBLE",
+                    },
+                ],
+            }
+        ],
+        "modelVersion": "gemini-1.5-flash-latest",
+    }
+    assert response.token_details == {
+        "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 9}],
+        "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 2}],
+    }
+    assert response.input_tokens == 9
+    assert response.output_tokens == 2
+
+    # And try it async too
+    async_model = llm.get_async_model("gemini-1.5-flash-latest")
+    response = await async_model.prompt(
+        "Name for a pet pelican, just the name", key=GEMINI_API_KEY
+    )
+    text = await response.text()
+    assert text == "Percy\n"
+
+
+@pytest.mark.vcr
+@pytest.mark.asyncio
+async def test_prompt_with_pydantic_schema():
+    class Dog(pydantic.BaseModel):
+        name: str
+        age: int
+        bio: str
+
+    model = llm.get_model("gemini-1.5-flash-latest")
+    response = model.prompt(
+        "Invent a cool dog", key=GEMINI_API_KEY, schema=Dog, stream=False
+    )
+    assert json.loads(response.text()) == {
+        "age": 3,
+        "bio": "A fluffy Samoyed with exceptional intelligence and a love for belly rubs. He's mastered several tricks, including fetching the newspaper and opening doors.",
+        "name": "Cloud",
+    }
+    assert response.response_json == {
+        "candidates": [
+            {
+                "finishReason": "STOP",
+                "safetyRatings": [
+                    {
+                        "category": "HARM_CATEGORY_HATE_SPEECH",
+                        "probability": "NEGLIGIBLE",
+                    },
+                    {
+                        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                        "probability": "NEGLIGIBLE",
+                    },
+                    {
+                        "category": "HARM_CATEGORY_HARASSMENT",
+                        "probability": "NEGLIGIBLE",
+                    },
+                    {
+                        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                        "probability": "NEGLIGIBLE",
+                    },
+                ],
+            }
+        ],
+        "modelVersion": "gemini-1.5-flash-latest",
+    }
+    assert response.input_tokens == 10

llm_gemini-0.13/tests/test_gemini.py

@@ -1,37 +0,0 @@
-import llm
-import nest_asyncio
-import os
-import pytest
-
-nest_asyncio.apply()
-
-GEMINI_API_KEY = os.environ.get("PYTEST_GEMINI_API_KEY", None) or "gm-..."
-
-
-@pytest.mark.vcr
-@pytest.mark.asyncio
-async def test_prompt():
-    model = llm.get_model("gemini-1.5-flash-latest")
-    response = model.prompt("Name for a pet pelican, just the name", key=GEMINI_API_KEY)
-    assert str(response) == "Percy"
-    assert response.response_json == [
-        {
-            "candidates": [
-                {"content": {"parts": [{"text": "Percy"}], "role": "model"}}
-            ],
-            "modelVersion": "gemini-1.5-flash-latest",
-        }
-    ]
-    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")
-    response = await async_model.prompt(
-        "Name for a pet pelican, just the name", key=GEMINI_API_KEY
-    )
-    text = await response.text()
-    assert text == "Percy"