Add Agent.output_json_schema() method (#3454)

Co-authored-by: Douwe Maan <douwe@pydantic.dev>

Author: g-eoj

pydantic_ai_slim/pydantic_ai/_output.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
 
 from pydantic import Json, TypeAdapter, ValidationError
+from pydantic._internal._typing_extra import get_function_type_hints
 from pydantic_core import SchemaValidator, to_json
 from typing_extensions import Self, TypedDict, TypeVar
 
@@ -1012,3 +1013,34 @@ def _flatten_output_spec(output_spec: OutputSpec[T]) -> Sequence[_OutputSpecItem
         else:
             outputs_flat.append(cast(_OutputSpecItem[T], output))
     return outputs_flat
+
+
+def types_from_output_spec(output_spec: OutputSpec[T]) -> Sequence[T | type[str]]:
+    outputs: Sequence[OutputSpec[T]]
+    if isinstance(output_spec, Sequence):
+        outputs = output_spec
+    else:
+        outputs = (output_spec,)
+
+    outputs_flat: list[T | type[str]] = []
+    for output in outputs:
+        if isinstance(output, NativeOutput):
+            outputs_flat.extend(types_from_output_spec(output.outputs))
+        elif isinstance(output, PromptedOutput):
+            outputs_flat.extend(types_from_output_spec(output.outputs))
+        elif isinstance(output, TextOutput):
+            outputs_flat.append(str)
+        elif isinstance(output, ToolOutput):
+            outputs_flat.extend(types_from_output_spec(output.output))
+        elif union_types := _utils.get_union_args(output):
+            outputs_flat.extend(union_types)
+        elif inspect.isfunction(output) or inspect.ismethod(output):
+            type_hints = get_function_type_hints(output)
+            if return_annotation := type_hints.get('return', None):
+                outputs_flat.extend(types_from_output_spec(return_annotation))
+            else:
+                outputs_flat.append(str)
+        else:
+            outputs_flat.append(cast(T, output))
+
+    return outputs_flat

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Generic, TypeAlias, cast, overload
 
 import anyio
+from pydantic import TypeAdapter
 from typing_extensions import Self, TypeIs, TypeVar
 
 from pydantic_graph import End
@@ -23,6 +24,8 @@
     result,
     usage as _usage,
 )
+from .._json_schema import JsonSchema
+from .._output import types_from_output_spec
 from .._tool_manager import ToolManager
 from ..builtin_tools import AbstractBuiltinTool
 from ..output import OutputDataT, OutputSpec
@@ -123,6 +126,28 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
         """
         raise NotImplementedError
 
+    def output_json_schema(self, output_type: OutputSpec[OutputDataT | RunOutputDataT] | None = None) -> JsonSchema:
+        """The output return JSON schema."""
+        if output_type is None:
+            output_type = self.output_type
+
+        return_types = types_from_output_spec(output_spec=output_type)
+
+        json_schemas: list[JsonSchema] = []
+        for return_type in return_types:
+            json_schema = TypeAdapter(return_type).json_schema(mode='serialization')
+            if json_schema not in json_schemas:
+                json_schemas.append(json_schema)
+
+        if len(json_schemas) == 1:
+            return json_schemas[0]
+        else:
+            json_schemas, all_defs = _utils.merge_json_schema_defs(json_schemas)
+            json_schema: JsonSchema = {'anyOf': json_schemas}
+            if all_defs:
+                json_schema['$defs'] = all_defs
+            return json_schema
+
     @overload
     async def run(
         self,

pydantic_ai_slim/pydantic_ai/agent/wrapper.py

@@ -10,6 +10,7 @@
     models,
     usage as _usage,
 )
+from .._json_schema import JsonSchema
 from ..builtin_tools import AbstractBuiltinTool
 from ..output import OutputDataT, OutputSpec
 from ..run import AgentRun
@@ -68,6 +69,9 @@ async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]:
     async def __aexit__(self, *args: Any) -> bool | None:
         return await self.wrapped.__aexit__(*args)
 
+    def output_json_schema(self, output_type: OutputSpec[OutputDataT | RunOutputDataT] | None = None) -> JsonSchema:
+        return self.wrapped.output_json_schema(output_type=output_type)
+
     @overload
     def iter(
         self,

tests/test_agent.py

@@ -5256,6 +5256,15 @@ def foo() -> str:
     assert wrapper_agent.name == 'wrapped'
     assert wrapper_agent.output_type == agent.output_type
     assert wrapper_agent.event_stream_handler == agent.event_stream_handler
+    assert wrapper_agent.output_json_schema() == snapshot(
+        {
+            'type': 'object',
+            'properties': {'a': {'title': 'A', 'type': 'integer'}, 'b': {'title': 'B', 'type': 'string'}},
+            'title': 'Foo',
+            'required': ['a', 'b'],
+        }
+    )
+    assert wrapper_agent.output_json_schema(output_type=str) == snapshot({'type': 'string'})
 
     bar_toolset = FunctionToolset()