Merge branch 'main' into alex/google-genai-prices

Author: DouweM

pydantic_ai_slim/pydantic_ai/_json_schema.py

@@ -25,7 +25,7 @@ def __init__(
         *,
         strict: bool | None = None,
         prefer_inlined_defs: bool = False,
-        simplify_nullable_unions: bool = False,
+        simplify_nullable_unions: bool = False,  # TODO (v2): Remove this, no longer used
     ):
         self.schema = schema
 
@@ -146,10 +146,9 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
 
         handled = [self._handle(member) for member in members]
 
-        # convert nullable unions to nullable types
+        # TODO (v2): Remove this feature, no longer used
         if self.simplify_nullable_unions:
             handled = self._simplify_nullable_union(handled)
-
         if len(handled) == 1:
             # In this case, no need to retain the union
             return handled[0] | schema
@@ -161,7 +160,7 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
 
     @staticmethod
     def _simplify_nullable_union(cases: list[JsonSchema]) -> list[JsonSchema]:
-        # TODO: Should we move this to relevant subclasses? Or is it worth keeping here to make reuse easier?
+        # TODO (v2): Remove this method, no longer used
         if len(cases) == 2 and {'type': 'null'} in cases:
             # Find the non-null schema
             non_null_schema = next(

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -292,7 +292,7 @@ async def count_tokens(
                     thinking_config=generation_config.get('thinking_config'),
                     media_resolution=generation_config.get('media_resolution'),
                     response_mime_type=generation_config.get('response_mime_type'),
-                    response_schema=generation_config.get('response_schema'),
+                    response_json_schema=generation_config.get('response_json_schema'),
                 ),
             )
 
@@ -456,7 +456,7 @@ async def _build_content_and_config(
             tools=cast(ToolListUnionDict, tools),
             tool_config=tool_config,
             response_mime_type=response_mime_type,
-            response_schema=response_schema,
+            response_json_schema=response_schema,
             response_modalities=modalities,
         )
         return contents, config

pydantic_ai_slim/pydantic_ai/profiles/google.py

@@ -1,9 +1,5 @@
 from __future__ import annotations as _annotations
 
-import warnings
-
-from pydantic_ai.exceptions import UserError
-
 from .._json_schema import JsonSchema, JsonSchemaTransformer
 from . import ModelProfile
 
@@ -23,84 +19,28 @@ def google_model_profile(model_name: str) -> ModelProfile | None:
 class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
     """Transforms the JSON Schema from Pydantic to be suitable for Gemini.
 
-    Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations)
-    a subset of OpenAPI v3.0.3.
-
-    Specifically:
-    * gemini doesn't allow the `title` keyword to be set
-    * gemini doesn't allow `$defs` — we need to inline the definitions where possible
+    Gemini supports [a subset of OpenAPI v3.0.3](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations).
     """
 
-    def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
-        super().__init__(schema, strict=strict, 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:
-            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 Pydantic AI GitHub"
-                ' and we will fix this behavior.',
-                UserWarning,
-            )
-
-        schema.pop('title', None)
+        # Remove properties not supported by Gemini
         schema.pop('$schema', None)
         if (const := schema.pop('const', None)) is not None:
             # Gemini doesn't support const, but it does support enum with a single value
             schema['enum'] = [const]
         schema.pop('discriminator', None)
         schema.pop('examples', None)
 
-        # TODO: Should we use the trick from pydantic_ai.models.openai._OpenAIJsonSchema
-        #   where we add notes about these properties to the field description?
-        schema.pop('exclusiveMaximum', None)
-        schema.pop('exclusiveMinimum', None)
-
-        # Gemini only supports string enums, so we need to convert any enum values to strings.
-        # Pydantic will take care of transforming the transformed string values to the correct type.
-        if enum := schema.get('enum'):
-            schema['type'] = 'string'
-            schema['enum'] = [str(val) for val in enum]
-
         type_ = schema.get('type')
-        if 'oneOf' in schema and 'type' not in schema:  # pragma: no cover
-            # This gets hit when we have a discriminated union
-            # Gemini returns an API error in this case even though it says in its error message it shouldn't...
-            # Changing the oneOf to an anyOf prevents the API error and I think is functionally equivalent
-            schema['anyOf'] = schema.pop('oneOf')
-
         if type_ == 'string' and (fmt := schema.pop('format', None)):
             description = schema.get('description')
             if description:
                 schema['description'] = f'{description} (format: {fmt})'
             else:
                 schema['description'] = f'Format: {fmt}'
 
-        if '$ref' in schema:
-            raise UserError(f'Recursive `$ref`s in JSON Schema are not supported by Gemini: {schema["$ref"]}')
-
-        if 'prefixItems' in schema:
-            # prefixItems is not currently supported in Gemini, so we convert it to items for best compatibility
-            prefix_items = schema.pop('prefixItems')
-            items = schema.get('items')
-            unique_items = [items] if items is not None else []
-            for item in prefix_items:
-                if item not in unique_items:
-                    unique_items.append(item)
-            if len(unique_items) > 1:  # pragma: no cover
-                schema['items'] = {'anyOf': unique_items}
-            elif len(unique_items) == 1:  # pragma: no branch
-                schema['items'] = unique_items[0]
-            schema.setdefault('minItems', len(prefix_items))
-            if items is None:  # pragma: no branch
-                schema.setdefault('maxItems', len(prefix_items))
+        # Note: exclusiveMinimum/exclusiveMaximum are NOT yet supported
+        schema.pop('exclusiveMinimum', None)
+        schema.pop('exclusiveMaximum', None)
 
         return schema

pydantic_ai_slim/pydantic_ai/ui/ag_ui/_adapter.py

@@ -2,11 +2,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from functools import cached_property
 from typing import (
     TYPE_CHECKING,
     Any,
+    cast,
 )
 
 from ... import ExternalToolset, ToolDefinition
@@ -107,7 +108,14 @@ def toolset(self) -> AbstractToolset[AgentDepsT] | None:
     @cached_property
     def state(self) -> dict[str, Any] | None:
         """Frontend state from the AG-UI run input."""
-        return self.run_input.state
+        state = self.run_input.state
+        if state is None:
+            return None
+
+        if isinstance(state, Mapping) and not state:
+            return None
+
+        return cast('dict[str, Any]', state)
 
     @classmethod
     def load_messages(cls, messages: Sequence[Message]) -> list[ModelMessage]:

tests/json_body_serializer.py

@@ -76,7 +76,7 @@ def serialize(cassette_dict: Any):  # pragma: lax no cover
                     del data['body']
             if content_type == ['application/x-www-form-urlencoded']:
                 query_params = urllib.parse.parse_qs(data['body'])
-                for key in ['client_secret', 'refresh_token']:  # pragma: no cover
+                for key in ['client_id', 'client_secret', 'refresh_token']:  # pragma: no cover
                     if key in query_params:
                         query_params[key] = ['scrubbed']
                         data['body'] = urllib.parse.urlencode(query_params)

tests/models/cassettes/test_google/test_google_dict_with_additional_properties_native_output.yaml

@@ -0,0 +1,78 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '519'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: Create a config named "api-config" with metadata author="Alice" and version="1.0"
+        role: user
+      generationConfig:
+        responseJsonSchema:
+          description: A response with configuration metadata.
+          properties:
+            metadata:
+              additionalProperties:
+                type: string
+              type: object
+            name:
+              type: string
+          required:
+          - name
+          - metadata
+          title: ConfigResponse
+          type: object
+        responseMimeType: application/json
+        responseModalities:
+        - TEXT
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '631'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=1379
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - content:
+          parts:
+          - text: '{"name": "api-config", "metadata": {"author": "Alice", "version": "1.0"}}'
+          role: model
+        finishReason: STOP
+        index: 0
+      modelVersion: gemini-2.5-flash
+      responseId: CZMUacOtKv2SxN8Pi7TrsAs
+      usageMetadata:
+        candidatesTokenCount: 25
+        promptTokenCount: 23
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 23
+        thoughtsTokenCount: 158
+        totalTokenCount: 206
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/cassettes/test_google/test_google_dict_with_additional_properties_native_output_gemini_2_0.yaml

@@ -0,0 +1,87 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '519'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: Create a config named "api-config" with metadata author="Alice" and version="1.0"
+        role: user
+      generationConfig:
+        responseJsonSchema:
+          description: A response with configuration metadata.
+          properties:
+            metadata:
+              additionalProperties:
+                type: string
+              type: object
+            name:
+              type: string
+          required:
+          - name
+          - metadata
+          title: ConfigResponse
+          type: object
+        responseMimeType: application/json
+        responseModalities:
+        - TEXT
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '760'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=818
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - avgLogprobs: -1.4634492981713265e-05
+        content:
+          parts:
+          - text: |-
+              {
+                "name": "api-config",
+                "metadata": {
+                  "author": "Alice",
+                  "version": "1.0"
+                }
+              }
+          role: model
+        finishReason: STOP
+      modelVersion: gemini-2.0-flash
+      responseId: 8pMcab_EMqWd28oP46bOiAk
+      usageMetadata:
+        candidatesTokenCount: 40
+        candidatesTokensDetails:
+        - modality: TEXT
+          tokenCount: 40
+        promptTokenCount: 22
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 22
+        totalTokenCount: 62
+    status:
+      code: 200
+      message: OK
+version: 1