llm-gemini 0.14__tar.gz → 0.14.1__tar.gz

{llm_gemini-0.14 → llm_gemini-0.14.1}/PKG-INFO

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

{llm_gemini-0.14 → llm_gemini-0.14.1}/llm_gemini.egg-info/PKG-INFO

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

{llm_gemini-0.14 → llm_gemini-0.14.1}/llm_gemini.py

@@ -88,18 +88,24 @@ def resolve_type(attachment):
     return mime_type
 
 
-def cleanup_schema(schema):
+def cleanup_schema(schema, in_properties=False):
     "Gemini supports only a subset of JSON schema"
     keys_to_remove = ("$schema", "additionalProperties", "title")
-    # Recursively remove them
+
     if isinstance(schema, dict):
-        for key in keys_to_remove:
-            schema.pop(key, None)
-        for value in schema.values():
-            cleanup_schema(value)
+        # Only remove keys if we're not inside a 'properties' block.
+        if not in_properties:
+            for key in keys_to_remove:
+                schema.pop(key, None)
+        for key, value in list(schema.items()):
+            # If the key is 'properties', set the flag for its value.
+            if key == "properties" and isinstance(value, dict):
+                cleanup_schema(value, in_properties=True)
+            else:
+                cleanup_schema(value, in_properties=in_properties)
     elif isinstance(schema, list):
-        for value in schema:
-            cleanup_schema(value)
+        for item in schema:
+            cleanup_schema(item, in_properties=in_properties)
     return schema
 
 

{llm_gemini-0.14 → llm_gemini-0.14.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "llm-gemini"
-version = "0.14"
+version = "0.14.1"
 description = "LLM plugin to access Google's Gemini family of models"
 readme = "README.md"
 authors = [{name = "Simon Willison"}]
@@ -25,3 +25,7 @@ gemini = "llm_gemini"
 
 [project.optional-dependencies]
 test = ["pytest", "pytest-recording", "pytest-asyncio", "nest-asyncio"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"

{llm_gemini-0.14 → llm_gemini-0.14.1}/tests/test_gemini.py

@@ -4,6 +4,7 @@ import json
 import os
 import pytest
 import pydantic
+from llm_gemini import cleanup_schema
 
 nest_asyncio.apply()
 
@@ -123,3 +124,89 @@ def test_embedding(model_id, monkeypatch):
     elif model_id.endswith("-512"):
         expected_length = 512
     assert len(response) == expected_length
+
+
+@pytest.mark.parametrize(
+    "schema,expected",
+    [
+        # Test 1: Top-level keys removal
+        (
+            {
+                "$schema": "http://json-schema.org/draft-07/schema#",
+                "title": "Example Schema",
+                "additionalProperties": False,
+                "type": "object",
+            },
+            {"type": "object"},
+        ),
+        # Test 2: Preserve keys within a "properties" block
+        (
+            {
+                "type": "object",
+                "properties": {
+                    "authors": {"type": "string"},
+                    "title": {"type": "string"},
+                    "reference": {"type": "string"},
+                    "year": {"type": "string"},
+                },
+                "title": "This should be removed from the top-level",
+            },
+            {
+                "type": "object",
+                "properties": {
+                    "authors": {"type": "string"},
+                    "title": {"type": "string"},
+                    "reference": {"type": "string"},
+                    "year": {"type": "string"},
+                },
+            },
+        ),
+        # Test 3: Nested keys outside and inside properties block
+        (
+            {
+                "definitions": {
+                    "info": {
+                        "title": "Info title",  # should be removed because it's not inside a "properties" block
+                        "description": "A description",
+                        "properties": {
+                            "name": {
+                                "title": "Name Title",
+                                "type": "string",
+                            },  # title here should be preserved
+                            "$schema": {
+                                "type": "string"
+                            },  # should be preserved as it's within properties
+                        },
+                    }
+                },
+                "$schema": "http://example.com/schema",
+            },
+            {
+                "definitions": {
+                    "info": {
+                        "description": "A description",
+                        "properties": {
+                            "name": {"title": "Name Title", "type": "string"},
+                            "$schema": {"type": "string"},
+                        },
+                    }
+                }
+            },
+        ),
+        # Test 4: List of schemas
+        (
+            [
+                {
+                    "$schema": "http://json-schema.org/draft-07/schema#",
+                    "type": "object",
+                },
+                {"title": "Should be removed", "type": "array"},
+            ],
+            [{"type": "object"}, {"type": "array"}],
+        ),
+    ],
+)
+def test_cleanup_schema(schema, expected):
+    # Use a deep copy so the original test data remains unchanged.
+    result = cleanup_schema(schema)
+    assert result == expected