Fix handling of additionalProperties by gemini (#1506)

Co-authored-by: David Montague <35119617+dmontagu@users.noreply.github.com>

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/_json_schema.py

@@ -20,11 +20,11 @@ class WalkJsonSchema(ABC):
     def __init__(
         self, schema: JsonSchema, *, prefer_inlined_defs: bool = False, simplify_nullable_unions: bool = False
     ):
-        self.schema = deepcopy(schema)
+        self.schema = schema
         self.prefer_inlined_defs = prefer_inlined_defs
         self.simplify_nullable_unions = simplify_nullable_unions
 
-        self.defs: dict[str, JsonSchema] = self.schema.pop('$defs', {})
+        self.defs: dict[str, JsonSchema] = self.schema.get('$defs', {})
         self.refs_stack = tuple[str, ...]()
         self.recursive_refs = set[str]()
 
@@ -34,7 +34,11 @@ def transform(self, schema: JsonSchema) -> JsonSchema:
         return schema
 
     def walk(self) -> JsonSchema:
-        handled = self._handle(deepcopy(self.schema))
+        schema = deepcopy(self.schema)
+
+        # First, handle everything but $defs:
+        schema.pop('$defs', None)
+        handled = self._handle(schema)
 
         if not self.prefer_inlined_defs and self.defs:
             handled['$defs'] = {k: self._handle(v) for k, v in self.defs.items()}

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 import base64
+import warnings
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field, replace
@@ -776,6 +777,22 @@ def __init__(self, schema: JsonSchema):
         super().__init__(schema, prefer_inlined_defs=True, simplify_nullable_unions=True)
 
     def transform(self, schema: JsonSchema) -> JsonSchema:
+        # Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini
+        additional_properties = schema.pop(
+            'additionalProperties', None
+        )  # don't pop yet so it's included in the warning
+        if additional_properties:  # pragma: no cover
+            original_schema = {**schema, 'additionalProperties': additional_properties}
+            warnings.warn(
+                '`additionalProperties` is not supported by Gemini; it will be removed from the tool JSON schema.'
+                f' Full schema: {self.schema}\n\n'
+                f'Source of additionalProperties within the full schema: {original_schema}\n\n'
+                'If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n'
+                "If Google's APIs are updated to support this properly, please create an issue on the PydanticAI GitHub"
+                ' and we will fix this behavior.',
+                UserWarning,
+            )
+
         schema.pop('title', None)
         schema.pop('default', None)
         schema.pop('$schema', None)

tests/models/cassettes/test_gemini/test_gemini_additional_properties_is_false.yaml

@@ -0,0 +1,76 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '296'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: What is the temperature in Tokyo?
+        role: user
+      tools:
+        function_declarations:
+        - description: null
+          name: get_temperature
+          parameters:
+            properties:
+              city:
+                type: string
+              country:
+                type: string
+            required:
+            - city
+            - country
+            type: object
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '748'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=523
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - avgLogprobs: -0.12538558465463143
+        content:
+          parts:
+          - text: |
+              The available tools lack the ability to access real-time information, including current temperature.  Therefore, I cannot answer your question.
+          role: model
+        finishReason: STOP
+      modelVersion: gemini-1.5-flash
+      usageMetadata:
+        candidatesTokenCount: 27
+        candidatesTokensDetails:
+        - modality: TEXT
+          tokenCount: 27
+        promptTokenCount: 14
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 14
+        totalTokenCount: 41
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/cassettes/test_gemini/test_gemini_additional_properties_is_true.yaml

@@ -0,0 +1,73 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '264'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: What is the temperature in Tokyo?
+        role: user
+      tools:
+        function_declarations:
+        - description: ''
+          name: get_temperature
+          parameters:
+            properties:
+              location:
+                type: object
+            required:
+            - location
+            type: object
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '741'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=534
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - avgLogprobs: -0.15060695580073766
+        content:
+          parts:
+          - text: |
+              I need a location dictionary to use the `get_temperature` function.  I cannot provide the temperature in Tokyo without more information.
+          role: model
+        finishReason: STOP
+      modelVersion: gemini-1.5-flash
+      usageMetadata:
+        candidatesTokenCount: 28
+        candidatesTokensDetails:
+        - modality: TEXT
+          tokenCount: 28
+        promptTokenCount: 12
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 12
+        totalTokenCount: 40
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/test_gemini.py

@@ -1029,3 +1029,40 @@ async def test_gemini_model_instructions(allow_model_requests: None, gemini_api_
             ),
         ]
     )
+
+
+class CurrentLocation(BaseModel, extra='forbid'):
+    city: str
+    country: str
+
+
+@pytest.mark.vcr()
+async def test_gemini_additional_properties_is_false(allow_model_requests: None, gemini_api_key: str):
+    m = GeminiModel('gemini-1.5-flash', provider=GoogleGLAProvider(api_key=gemini_api_key))
+    agent = Agent(m)
+
+    @agent.tool_plain
+    async def get_temperature(location: CurrentLocation) -> float:  # pragma: no cover
+        return 20.0
+
+    result = await agent.run('What is the temperature in Tokyo?')
+    assert result.output == snapshot(
+        'The available tools lack the ability to access real-time information, including current temperature.  Therefore, I cannot answer your question.\n'
+    )
+
+
+@pytest.mark.vcr()
+async def test_gemini_additional_properties_is_true(allow_model_requests: None, gemini_api_key: str):
+    m = GeminiModel('gemini-1.5-flash', provider=GoogleGLAProvider(api_key=gemini_api_key))
+    agent = Agent(m)
+
+    with pytest.warns(UserWarning, match='.*additionalProperties.*'):
+
+        @agent.tool_plain
+        async def get_temperature(location: dict[str, CurrentLocation]) -> float:  # pragma: no cover
+            return 20.0
+
+        result = await agent.run('What is the temperature in Tokyo?')
+        assert result.output == snapshot(
+            'I need a location dictionary to use the `get_temperature` function.  I cannot provide the temperature in Tokyo without more information.\n'
+        )