Add StructuredDict for structured outputs with custom JSON schema (#2157)

Co-authored-by: Douwe Maan <[email protected]>

Author: fswair

docs/api/output.md

@@ -9,3 +9,4 @@
             - NativeOutput
             - PromptedOutput
             - TextOutput
+            - StructuredDict

docs/output.md

@@ -31,7 +31,7 @@ _(This example is complete, it can be run "as is")_
 
 ## Output data {#structured-output}
 
-The [`Agent`][pydantic_ai.Agent] class constructor takes an `output_type` argument that takes one or more types or [output functions](#output-functions). It supports simple scalar types, list and dict types, dataclasses and Pydantic models, as well as type unions -- generally everything supported as type hints in a Pydantic model. You can also pass a list of multiple choices.
+The [`Agent`][pydantic_ai.Agent] class constructor takes an `output_type` argument that takes one or more types or [output functions](#output-functions). It supports simple scalar types, list and dict types (including `TypedDict`s and [`StructuredDict`s](#structured-dict)), dataclasses and Pydantic models, as well as type unions -- generally everything supported as type hints in a Pydantic model. You can also pass a list of multiple choices.
 
 By default, Pydantic AI leverages the model's tool calling capability to make it return structured data. When multiple output types are specified (in a union or list), each member is registered with the model as a separate output tool in order to reduce the complexity of the schema and maximise the chances a model will respond correctly. This has been shown to work well across a wide range of models. If you'd like to change the names of the output tools, use a model's native structured output feature, or pass the output schema to the model in its [instructions](agents.md#instructions), you can use an [output mode](#output-modes) marker class.
 
@@ -117,7 +117,6 @@ print(result.output)
 
 _(This example is complete, it can be run "as is")_
 
-
 ### Output functions
 
 Instead of plain text or structured data, you may want the output of your agent run to be the result of a function called with arguments provided by the model, for example to further process or validate the data provided through the arguments (with the option to tell the model to try again), or to hand off to another agent.
@@ -387,6 +386,37 @@ print(repr(result.output))
 
 _(This example is complete, it can be run "as is")_
 
+### Custom JSON schema {#structured-dict}
+
+If it's not feasible to define your desired structured output object using a Pydantic `BaseModel`, dataclass, or `TypedDict`, for example when you get a JSON schema from an external source or generate it dynamically, you can use the [`StructuredDict()`][pydantic_ai.output.StructuredDict] helper function to generate a `dict[str, Any]` subclass with a JSON schema attached that Pydantic AI will pass to the model.
+
+Note that Pydantic AI will not perform any validation of the received JSON object and it's up to the model to correctly interpret the schema and any constraints expressed in it, like required fields or integer value ranges.
+
+The output type will be a `dict[str, Any]` and it's up to your code to defensively read from it in case the model made a mistake. You can use an [output validator](#output-validator-functions) to reflect validation errors back to the model and get it to try again.
+
+Along with the JSON schema, you can optionally pass `name` and `description` arguments to provide additional context to the model:
+
+```python
+from pydantic_ai import Agent, StructuredDict
+
+HumanDict = StructuredDict(
+    {
+        "type": "object",
+        "properties": {
+            "name": {"type": "string"},
+            "age": {"type": "integer"}
+        },
+        "required": ["name", "age"]
+    },
+    name="Human",
+    description="A human with a name and age",
+)
+
+agent = Agent('openai:gpt-4o', output_type=HumanDict)
+result = agent.run_sync("Create a person")
+#> {'name': 'John Doe', 'age': 30}
+```
+
 ### Output validators {#output-validator-functions}
 
 Some validation is inconvenient or impossible to do in Pydantic validators, in particular when the validation requires IO and is asynchronous. PydanticAI provides a way to add validation functions via the [`agent.output_validator`][pydantic_ai.Agent.output_validator] decorator.

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -12,7 +12,7 @@
 )
 from .format_prompt import format_as_xml
 from .messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl, VideoUrl
-from .output import NativeOutput, PromptedOutput, TextOutput, ToolOutput
+from .output import NativeOutput, PromptedOutput, StructuredDict, TextOutput, ToolOutput
 from .tools import RunContext, Tool
 
 __all__ = (
@@ -46,6 +46,7 @@
     'NativeOutput',
     'PromptedOutput',
     'TextOutput',
+    'StructuredDict',
     # format_prompt
     'format_as_xml',
 )

pydantic_ai_slim/pydantic_ai/_output.py

@@ -264,22 +264,23 @@ def _build_tools(
 
                 output = output.output
 
+            description = description or default_description
+            if strict is None:
+                strict = default_strict
+
+            processor = ObjectOutputProcessor(output=output, description=description, strict=strict)
+
             if name is None:
                 name = default_name
                 if multiple:
-                    name += f'_{output.__name__}'
+                    name += f'_{processor.object_def.name}'
 
             i = 1
             original_name = name
             while name in tools:
                 i += 1
                 name = f'{original_name}_{i}'
 
-            description = description or default_description
-            if strict is None:
-                strict = default_strict
-
-            processor = ObjectOutputProcessor(output=output, description=description, strict=strict)
             tools[name] = OutputTool(name=name, processor=processor, multiple=multiple)
 
         return tools
@@ -616,6 +617,9 @@ def __init__(
                 # including `response_data_typed_dict` as a title here doesn't add anything and could confuse the LLM
                 json_schema.pop('title')
 
+        if name is None and (json_schema_title := json_schema.get('title', None)):
+            name = json_schema_title
+
         if json_schema_description := json_schema.pop('description', None):
             if description is None:
                 description = json_schema_description

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -60,7 +60,12 @@ def is_model_like(type_: Any) -> bool:
     return (
         isinstance(type_, type)
         and not isinstance(type_, GenericAlias)
-        and (issubclass(type_, BaseModel) or is_dataclass(type_) or is_typeddict(type_))  # pyright: ignore[reportUnknownArgumentType]
+        and (
+            issubclass(type_, BaseModel)
+            or is_dataclass(type_)  # pyright: ignore[reportUnknownArgumentType]
+            or is_typeddict(type_)  # pyright: ignore[reportUnknownArgumentType]
+            or getattr(type_, '__is_model_like__', False)  # pyright: ignore[reportUnknownArgumentType]
+        )
     )
 
 

pydantic_ai_slim/pydantic_ai/output.py

@@ -2,10 +2,14 @@
 
 from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass
-from typing import Callable, Generic, Literal, Union
+from typing import Any, Callable, Generic, Literal, Union
 
+from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema
 from typing_extensions import TypeAliasType, TypeVar
 
+from . import _utils
 from .tools import RunContext
 
 __all__ = (
@@ -14,6 +18,7 @@
     'NativeOutput',
     'PromptedOutput',
     'TextOutput',
+    'StructuredDict',
     # types
     'OutputDataT',
     'OutputMode',
@@ -266,6 +271,65 @@ def split_into_words(text: str) -> list[str]:
     """The function that will be called to process the model's plain text output. The function must take a single string argument."""
 
 
+def StructuredDict(
+    json_schema: JsonSchemaValue, name: str | None = None, description: str | None = None
+) -> type[JsonSchemaValue]:
+    """Returns a `dict[str, Any]` subclass with a JSON schema attached that will be used for structured output.
+
+    Args:
+        json_schema: A JSON schema of type `object` defining the structure of the dictionary content.
+        name: Optional name of the structured output. If not provided, the `title` field of the JSON schema will be used if it's present.
+        description: Optional description of the structured output. If not provided, the `description` field of the JSON schema will be used if it's present.
+
+    Example:
+    ```python {title="structured_dict.py"}
+    from pydantic_ai import Agent, StructuredDict
+
+
+    schema = {
+        "type": "object",
+        "properties": {
+            "name": {"type": "string"},
+            "age": {"type": "integer"}
+        },
+        "required": ["name", "age"]
+    }
+
+    agent = Agent('openai:gpt-4o', output_type=StructuredDict(schema))
+    result = agent.run_sync("Create a person")
+    print(result.output)
+    #> {'name': 'John Doe', 'age': 30}
+    ```
+    """
+    json_schema = _utils.check_object_json_schema(json_schema)
+
+    if name:
+        json_schema['title'] = name
+
+    if description:
+        json_schema['description'] = description
+
+    class _StructuredDict(JsonSchemaValue):
+        __is_model_like__ = True
+
+        @classmethod
+        def __get_pydantic_core_schema__(
+            cls, source_type: Any, handler: GetCoreSchemaHandler
+        ) -> core_schema.CoreSchema:
+            return core_schema.dict_schema(
+                keys_schema=core_schema.str_schema(),
+                values_schema=core_schema.any_schema(),
+            )
+
+        @classmethod
+        def __get_pydantic_json_schema__(
+            cls, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+        ) -> JsonSchemaValue:
+            return json_schema
+
+    return _StructuredDict
+
+
 OutputSpec = TypeAliasType(
     'OutputSpec',
     Union[

tests/test_agent.py

@@ -41,7 +41,7 @@
 )
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.output import ToolOutput
+from pydantic_ai.output import StructuredDict, ToolOutput
 from pydantic_ai.profiles import ModelProfile
 from pydantic_ai.result import Usage
 from pydantic_ai.tools import ToolDefinition
@@ -1266,6 +1266,77 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
     )
 
 
+def test_output_type_structured_dict():
+    PersonDict = StructuredDict(
+        {
+            'type': 'object',
+            'properties': {
+                'name': {'type': 'string'},
+                'age': {'type': 'integer'},
+            },
+            'required': ['name', 'age'],
+        },
+        name='Person',
+        description='A person',
+    )
+    AnimalDict = StructuredDict(
+        {
+            'type': 'object',
+            'properties': {
+                'name': {'type': 'string'},
+                'species': {'type': 'string'},
+            },
+            'required': ['name', 'species'],
+        },
+        name='Animal',
+        description='An animal',
+    )
+
+    output_tools = None
+
+    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        assert info.output_tools is not None
+
+        nonlocal output_tools
+        output_tools = info.output_tools
+
+        args_json = '{"name": "John Doe", "age": 30}'
+        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])
+
+    agent = Agent(
+        FunctionModel(call_tool),
+        output_type=[PersonDict, AnimalDict],
+    )
+
+    result = agent.run_sync('Generate a person')
+
+    assert result.output == snapshot({'name': 'John Doe', 'age': 30})
+    assert output_tools == snapshot(
+        [
+            ToolDefinition(
+                name='final_result_Person',
+                parameters_json_schema={
+                    'properties': {'name': {'type': 'string'}, 'age': {'type': 'integer'}},
+                    'required': ['name', 'age'],
+                    'title': 'Person',
+                    'type': 'object',
+                },
+                description='A person',
+            ),
+            ToolDefinition(
+                name='final_result_Animal',
+                parameters_json_schema={
+                    'properties': {'name': {'type': 'string'}, 'species': {'type': 'string'}},
+                    'required': ['name', 'species'],
+                    'title': 'Animal',
+                    'type': 'object',
+                },
+                description='An animal',
+            ),
+        ]
+    )
+
+
 def test_default_structured_output_mode():
     def hello(_: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
         return ModelResponse(parts=[TextPart(content='hello')])  # pragma: no cover

tests/test_examples.py

@@ -444,6 +444,10 @@ async def list_tools() -> list[None]:
     'What is a Ford Explorer?': '{"result": {"kind": "Vehicle", "data": {"name": "Ford Explorer", "wheels": 4}}}',
     'What is a MacBook?': '{"result": {"kind": "Device", "data": {"name": "MacBook", "kind": "laptop"}}}',
     'Write a creative story about space exploration': 'In the year 2157, Captain Maya Chen piloted her spacecraft through the vast expanse of the Andromeda Galaxy. As she discovered a planet with crystalline mountains that sang in harmony with the cosmic winds, she realized that space exploration was not just about finding new worlds, but about finding new ways to understand the universe and our place within it.',
+    'Create a person': ToolCallPart(
+        tool_name='final_result',
+        args={'name': 'John Doe', 'age': 30},
+    ),
 }
 
 tool_responses: dict[tuple[str, str], str] = {

tests/typed_agent.py

@@ -4,13 +4,13 @@
 from collections.abc import Awaitable
 from dataclasses import dataclass
 from decimal import Decimal
-from typing import Callable, TypeAlias, Union
+from typing import Any, Callable, TypeAlias, Union
 
 from typing_extensions import assert_type
 
 from pydantic_ai import Agent, ModelRetry, RunContext, Tool
 from pydantic_ai.agent import AgentRunResult
-from pydantic_ai.output import TextOutput, ToolOutput
+from pydantic_ai.output import StructuredDict, TextOutput, ToolOutput
 from pydantic_ai.tools import ToolDefinition
 
 # Define here so we can check `if MYPY` below. This will not be executed, MYPY will always set it to True
@@ -170,6 +170,16 @@ def run_sync3() -> None:
 union_agent2: Agent[None, MyUnion] = Agent(output_type=MyUnion)  # type: ignore[call-overload]
 assert_type(union_agent2, Agent[None, MyUnion])
 
+structured_dict = StructuredDict(
+    {
+        'type': 'object',
+        'properties': {'name': {'type': 'string'}, 'age': {'type': 'integer'}},
+        'required': ['name', 'age'],
+    }
+)
+structured_dict_agent = Agent(output_type=structured_dict)
+assert_type(structured_dict_agent, Agent[None, dict[str, Any]])
+
 
 def foobar_ctx(ctx: RunContext[int], x: str, y: int) -> Decimal:
     return Decimal(x) + y