google: add more information about schema on union (#2426)

Author: Kludex

pydantic_ai_slim/pydantic_ai/profiles/_json_schema.py

@@ -137,8 +137,9 @@ def _handle_array(self, schema: JsonSchema) -> JsonSchema:
         return schema
 
     def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf']) -> JsonSchema:
-        members = schema.get(union_kind)
-        if not members:
+        try:
+            members = schema.pop(union_kind)
+        except KeyError:
             return schema
 
         handled = [self._handle(member) for member in members]
@@ -149,7 +150,7 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
 
         if len(handled) == 1:
             # In this case, no need to retain the union
-            return handled[0]
+            return handled[0] | schema
 
         # If we have keys besides the union kind (such as title or discriminator), keep them without modifications
         schema = schema.copy()

tests/models/test_gemini.py

@@ -360,24 +360,25 @@ class QueryDetails(BaseModel):
 
     # This tests that the enum values are properly converted to strings for Gemini
     assert m._get_tools(mrp) == snapshot(
-        _GeminiTools(
-            function_declarations=[
-                _GeminiFunction(
-                    name='result',
-                    description='This is the tool for the final Result',
-                    parameters={
+        {
+            'function_declarations': [
+                {
+                    'name': 'result',
+                    'description': 'This is the tool for the final Result',
+                    'parameters': {
                         'properties': {
                             'progress': {
                                 'items': {'enum': ['100', '80', '60', '40', '20'], 'type': 'string'},
                                 'type': 'array',
                                 'nullable': True,
+                                'default': None,
                             }
                         },
                         'type': 'object',
                     },
-                )
+                }
             ]
-        )
+        }
     )
 
 
@@ -406,12 +407,12 @@ class Locations(BaseModel):
     )
     mrp = m.customize_request_parameters(mrp)
     assert m._get_tools(mrp) == snapshot(
-        _GeminiTools(
-            function_declarations=[
-                _GeminiFunction(
-                    name='result',
-                    description='This is the tool for the final Result',
-                    parameters={
+        {
+            'function_declarations': [
+                {
+                    'name': 'result',
+                    'description': 'This is the tool for the final Result',
+                    'parameters': {
                         'properties': {
                             'op_location': {
                                 'properties': {
@@ -421,13 +422,14 @@ class Locations(BaseModel):
                                 'required': ['lat', 'lng'],
                                 'nullable': True,
                                 'type': 'object',
+                                'default': None,
                             }
                         },
                         'type': 'object',
                     },
-                )
+                }
             ]
-        )
+        }
     )