Flatten allOf properties for OpenAI compatibility

Author: pulphix

docs/models/openai.md

@@ -99,6 +99,10 @@ agent = Agent(model)
 ```
 
 ## OpenAI Responses API
+### Tips for JSON Schema compatibility
+
+Strict Structured Outputs prefer flat object schemas without combinators. If your tool or output schema contains `allOf`/`oneOf`, consider flattening `allOf` ahead of time. See the Tools documentation section "Flattening allOf for provider compatibility" for a sample prepare hook.
+
 
 Pydantic AI also supports OpenAI's [Responses API](https://platform.openai.com/docs/api-reference/responses) through the
 

docs/tools.md

@@ -357,6 +357,30 @@ print(test_model.last_model_request_parameters.function_tools)
 _(This example is complete, it can be run "as is")_
 
 
+### Flattening allOf for provider compatibility
+
+Some providers (especially when using strict Structured Outputs) reject JSON Schema combinators like `allOf/oneOf`.
+If your tool parameter schema includes `allOf`, you can flatten it before sending to the model using a prepare hook.
+
+Example prepare hook that flattens `allOf`:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.tools import RunContext, ToolDefinition
+from pydantic_ai._json_schema import flatten_allof
+
+async def flatten_prepare(ctx: RunContext[None], tool: ToolDefinition) -> ToolDefinition:
+    tool.parameters_json_schema = flatten_allof(tool.parameters_json_schema)
+    return tool
+
+# Register your tools normally, then pass `prepare_tools=flatten_prepare` to Agent if you want to apply globally.
+agent = Agent('openai:gpt-4o', prepare_tools=lambda ctx, tools: [
+    await flatten_prepare(ctx, t) or t for t in tools
+])
+```
+
+Alternatively, you can construct tools with a flattened schema at source (e.g., for MCP-exposed tools) using `Tool.from_schema`.
+
 ## See Also
 
 For more tool features and integrations, see:

pydantic_ai_slim/pydantic_ai/_json_schema.py

@@ -4,12 +4,14 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from .exceptions import UserError
 
 JsonSchema = dict[str, Any]
 
+__all__ = ['JsonSchemaTransformer', 'InlineDefsJsonSchemaTransformer', 'flatten_allof']
+
 
 @dataclass(init=False)
 class JsonSchemaTransformer(ABC):
@@ -30,7 +32,9 @@ def __init__(
         self.schema = schema
 
         self.strict = strict
-        self.is_strict_compatible = True  # Can be set to False by subclasses to set `strict` on `ToolDefinition` when set not set by user explicitly
+        # Can be set to False by subclasses to set `strict` on `ToolDefinition`
+        # when not set explicitly by the user.
+        self.is_strict_compatible = True
 
         self.prefer_inlined_defs = prefer_inlined_defs
         self.simplify_nullable_unions = simplify_nullable_unions
@@ -188,3 +192,108 @@ def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
 
     def transform(self, schema: JsonSchema) -> JsonSchema:
         return schema
+
+
+def _allof_is_object_like(member: JsonSchema) -> bool:
+    member_type = member.get('type')
+    if member_type is None:
+        keys = ('properties', 'additionalProperties', 'patternProperties')
+        return bool(any(k in member for k in keys))
+    return member_type == 'object'
+
+
+def _merge_additional_properties_values(values: list[Any]) -> bool | JsonSchema:
+    if any(isinstance(v, dict) for v in values):
+        return True
+    return False if values and all(v is False for v in values) else True
+
+
+def _flatten_current_level(s: JsonSchema) -> JsonSchema:
+    raw_members = s.get('allOf')
+    if not isinstance(raw_members, list) or not raw_members:
+        return s
+
+    members = cast(list[JsonSchema], raw_members)
+    for raw in members:
+        if not isinstance(raw, dict):
+            return s
+    if not all(_allof_is_object_like(member) for member in members):
+        return s
+
+    processed_members = [_recurse_flatten_allof(member) for member in members]
+    merged: JsonSchema = {k: v for k, v in s.items() if k != 'allOf'}
+    merged['type'] = 'object'
+
+    properties: dict[str, JsonSchema] = {}
+    if isinstance(merged.get('properties'), dict):
+        properties.update(merged['properties'])
+
+    required: set[str] = set(merged.get('required', []) or [])
+    pattern_properties: dict[str, JsonSchema] = dict(merged.get('patternProperties', {}) or {})
+    additional_values: list[Any] = []
+
+    for m in processed_members:
+        if isinstance(m.get('properties'), dict):
+            properties.update(m['properties'])
+        if isinstance(m.get('required'), list):
+            required.update(m['required'])
+        if isinstance(m.get('patternProperties'), dict):
+            pattern_properties.update(m['patternProperties'])
+        if 'additionalProperties' in m:
+            additional_values.append(m['additionalProperties'])
+
+    if properties:
+        merged['properties'] = {k: _recurse_flatten_allof(v) for k, v in properties.items()}
+    if required:
+        merged['required'] = sorted(required)
+    if pattern_properties:
+        merged['patternProperties'] = {k: _recurse_flatten_allof(v) for k, v in pattern_properties.items()}
+
+    if additional_values:
+        merged['additionalProperties'] = _merge_additional_properties_values(additional_values)
+
+    return merged
+
+
+def _recurse_children(s: JsonSchema) -> JsonSchema:
+    t = s.get('type')
+    if t == 'object':
+        if isinstance(s.get('properties'), dict):
+            s['properties'] = {
+                k: _recurse_flatten_allof(cast(JsonSchema, v))
+                for k, v in s['properties'].items()
+                if isinstance(v, dict)
+            }
+        ap = s.get('additionalProperties')
+        if isinstance(ap, dict):
+            ap_schema = cast(JsonSchema, ap)
+            s['additionalProperties'] = _recurse_flatten_allof(ap_schema)
+        if isinstance(s.get('patternProperties'), dict):
+            s['patternProperties'] = {
+                k: _recurse_flatten_allof(cast(JsonSchema, v))
+                for k, v in s['patternProperties'].items()
+                if isinstance(v, dict)
+            }
+    elif t == 'array':
+        items = s.get('items')
+        if isinstance(items, dict):
+            s['items'] = _recurse_flatten_allof(cast(JsonSchema, items))
+    return s
+
+
+def _recurse_flatten_allof(schema: JsonSchema) -> JsonSchema:
+    s = deepcopy(schema)
+    s = _flatten_current_level(s)
+    s = _recurse_children(s)
+    return s
+
+
+def flatten_allof(schema: JsonSchema) -> JsonSchema:
+    """Flatten simple object-only allOf combinations by merging object members.
+
+    - Merges properties and unions required lists.
+    - Combines additionalProperties conservatively: only False if all are False; otherwise True.
+    - Recurses into nested object/array members.
+    - Leaves non-object allOfs untouched.
+    """
+    return _recurse_flatten_allof(schema)

pydantic_ai_slim/pydantic_ai/profiles/google.py

@@ -4,7 +4,7 @@
 
 from pydantic_ai.exceptions import UserError
 
-from .._json_schema import JsonSchema, JsonSchemaTransformer
+from .._json_schema import JsonSchema, JsonSchemaTransformer, flatten_allof
 from . import ModelProfile
 
 
@@ -35,6 +35,8 @@ 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:
+        # Flatten object-only allOf to improve compatibility
+        schema = flatten_allof(schema)
         # Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini
         additional_properties = schema.pop(
             'additionalProperties', None

pydantic_ai_slim/pydantic_ai/profiles/openai.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from typing import Any, Literal
 
-from .._json_schema import JsonSchema, JsonSchemaTransformer
+from .._json_schema import JsonSchema, JsonSchemaTransformer, flatten_allof
 from . import ModelProfile
 
 OpenAISystemPromptRole = Literal['system', 'developer', 'user']
@@ -162,6 +162,8 @@ def walk(self) -> JsonSchema:
         return result
 
     def transform(self, schema: JsonSchema) -> JsonSchema:  # noqa C901
+        # First, flatten object-only allOf constructs to avoid unsupported combinators in strict mode
+        schema = flatten_allof(schema)
         # Remove unnecessary keys
         schema.pop('title', None)
         schema.pop('$schema', None)

tests/models/test_openai_like_providers.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import pytest
+
+from pydantic_ai._json_schema import JsonSchema
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer
+
+
+def _openai_transformer_factory(schema: JsonSchema) -> OpenAIJsonSchemaTransformer:
+    return OpenAIJsonSchemaTransformer(schema, strict=True)
+
+
+TransformerFactory = Callable[[JsonSchema], OpenAIJsonSchemaTransformer]
+
+
+@pytest.mark.parametrize('transformer_factory', [_openai_transformer_factory])
+def test_openai_compatible_transformers_flatten_allof(
+    transformer_factory: TransformerFactory,
+) -> None:
+    schema: JsonSchema = {
+        'type': 'object',
+        'allOf': [
+            {
+                'type': 'object',
+                'properties': {'foo': {'type': 'string'}},
+                'required': ['foo'],
+            },
+            {
+                'type': 'object',
+                'properties': {'bar': {'type': 'integer'}},
+                'required': ['bar'],
+            },
+        ],
+    }
+
+    transformer = transformer_factory(schema)
+    transformed = transformer.walk()
+
+    # allOf should have been flattened by the transformer
+    assert 'allOf' not in transformed
+    assert transformed['type'] == 'object'
+    assert set(transformed.get('required', [])) == {'foo', 'bar'}
+    assert transformed['properties']['foo']['type'] == 'string'
+    assert transformed['properties']['bar']['type'] == 'integer'

tests/test_json_schema_flattener.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+
+def test_flatten_allof_simple_merge() -> None:
+    # Import inline to avoid import errors before implementation exists in editors
+    from pydantic_ai._json_schema import flatten_allof
+
+    schema: dict[str, Any] = {
+        'type': 'object',
+        'allOf': [
+            {
+                'type': 'object',
+                'properties': {'a': {'type': 'string'}},
+                'required': ['a'],
+                'additionalProperties': False,
+            },
+            {
+                'type': 'object',
+                'properties': {'b': {'type': 'integer'}},
+                'required': ['b'],
+                'additionalProperties': False,
+            },
+        ],
+    }
+
+    flattened = flatten_allof(copy.deepcopy(schema))
+
+    assert 'allOf' not in flattened
+    assert flattened['type'] == 'object'
+    assert flattened['properties']['a']['type'] == 'string'
+    assert flattened['properties']['b']['type'] == 'integer'
+    # union of required keys
+    assert set(flattened['required']) == {'a', 'b'}
+    # boolean AP should remain False when all are False
+    assert flattened.get('additionalProperties') is False
+
+
+def test_flatten_allof_nested_objects_and_pass_through_keywords() -> None:
+    from pydantic_ai._json_schema import flatten_allof
+
+    schema: dict[str, Any] = {
+        'type': 'object',
+        'title': 'Root',
+        'allOf': [
+            {
+                'type': 'object',
+                'properties': {
+                    'user': {
+                        'type': 'object',
+                        'properties': {'id': {'type': 'string'}},
+                        'required': ['id'],
+                    }
+                },
+                'required': ['user'],
+            },
+            {
+                'type': 'object',
+                'properties': {'age': {'type': 'integer'}},
+                'required': ['age'],
+            },
+        ],
+        'description': 'test',
+    }
+
+    flattened = flatten_allof(copy.deepcopy(schema))
+    assert flattened.get('title') == 'Root'
+    assert flattened.get('description') == 'test'
+    assert 'allOf' not in flattened
+    assert set(flattened['required']) == {'user', 'age'}
+    assert flattened['properties']['user']['type'] == 'object'
+    assert set(flattened['properties']['user']['required']) == {'id'}
+
+
+def test_flatten_allof_does_not_touch_unrelated_unions() -> None:
+    from pydantic_ai._json_schema import flatten_allof
+
+    schema: dict[str, Any] = {
+        'type': 'object',
+        'properties': {
+            'x': {
+                'anyOf': [
+                    {'type': 'string'},
+                    {'type': 'null'},
+                ]
+            }
+        },
+        'required': ['x'],
+    }
+
+    flattened = flatten_allof(copy.deepcopy(schema))
+    assert flattened['properties']['x'].get('anyOf') is not None
+
+
+def test_flatten_allof_non_object_members_are_left_as_is() -> None:
+    from pydantic_ai._json_schema import flatten_allof
+
+    schema: dict[str, Any] = {
+        'type': 'object',
+        'allOf': [
+            {'type': 'string'},
+            {'type': 'integer'},
+        ],
+    }
+
+    # Expect: we cannot sensibly merge non-object members; keep allOf
+    flattened = flatten_allof(copy.deepcopy(schema))
+    assert 'allOf' in flattened