Add Agent.output_json_schema() method

pydantic_ai_slim/pydantic_ai/_output.py

@@ -6,6 +6,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
 
 from pydantic import Json, TypeAdapter, ValidationError
@@ -15,6 +16,7 @@
 from pydantic_ai._instrumentation import InstrumentationNames
 
 from . import _function_schema, _utils, messages as _messages
+from ._json_schema import JsonSchema
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UserError
 from .output import (
@@ -226,6 +228,10 @@ def mode(self) -> OutputMode:
     def allows_text(self) -> bool:
         return self.text_processor is not None
 
+    @cached_property
+    def json_schema(self) -> JsonSchema:
+        raise NotImplementedError()
+
     @classmethod
     def build(  # noqa: C901
         cls,
@@ -353,7 +359,9 @@ def build(  # noqa: C901
 
         if len(other_outputs) > 0:
             return AutoOutputSchema(
-                processor=cls._build_processor(other_outputs, name=name, description=description, strict=strict),
+                processor=cls._build_processor(
+                    outputs=other_outputs, name=name, description=description, strict=strict
+                ),
                 toolset=toolset,
                 allows_deferred_tools=allows_deferred_tools,
                 allows_image=allows_image,
@@ -377,6 +385,56 @@ def _build_processor(
 
         return UnionOutputProcessor(outputs=outputs, strict=strict, name=name, description=description)
 
+    def build_json_schema(self) -> JsonSchema:  # noqa: C901
+        # allow any output with {'type': 'string'} if no constraints
+        if not any([self.allows_deferred_tools, self.allows_image, self.object_def, self.toolset]):
+            return TypeAdapter(str).json_schema()
+
+        json_schemas: list[ObjectJsonSchema] = []
+
+        processor = getattr(self, 'processor', None)
+        if isinstance(processor, ObjectOutputProcessor):
+            json_schema = processor.object_def.json_schema
+            if k := processor.outer_typed_dict_key:
+                json_schema = json_schema['properties'][k]
+            json_schemas.append(json_schema)
+
+        elif self.toolset:
+            if self.allows_text:
+                json_schema = TypeAdapter(str).json_schema()
+                json_schemas.append(json_schema)
+            for tool_processor in self.toolset.processors.values():
+                json_schema = tool_processor.object_def.json_schema
+                if k := tool_processor.outer_typed_dict_key:
+                    json_schema = json_schema['properties'][k]
+                if json_schema not in json_schemas:
+                    json_schemas.append(json_schema)
+
+        elif self.allows_text:
+            json_schema = TypeAdapter(str).json_schema()
+            json_schemas.append(json_schema)
+
+        if self.allows_deferred_tools:
+            json_schema = TypeAdapter(DeferredToolRequests).json_schema(mode='serialization')
+            if json_schema not in json_schemas:
+                json_schemas.append(json_schema)
+
+        if self.allows_image:
+            json_schema = TypeAdapter(_messages.BinaryImage).json_schema()
+            json_schema = {k: v for k, v in json_schema['properties'].items() if k in ['data', 'media_type']}
+            if json_schema not in json_schemas:
+                json_schemas.append(json_schema)
+
+        if len(json_schemas) == 1:
+            return json_schemas[0]
+
+        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
+
 
 @dataclass(init=False)
 class AutoOutputSchema(OutputSchema[OutputDataT]):
@@ -405,6 +463,10 @@ def __init__(
     def mode(self) -> OutputMode:
         return 'auto'
 
+    @cached_property
+    def json_schema(self) -> JsonSchema:
+        return self.build_json_schema()
+
 
 @dataclass(init=False)
 class TextOutputSchema(OutputSchema[OutputDataT]):
@@ -425,6 +487,10 @@ def __init__(
     def mode(self) -> OutputMode:
         return 'text'
 
+    @cached_property
+    def json_schema(self) -> JsonSchema:
+        return self.build_json_schema()
+
 
 class ImageOutputSchema(OutputSchema[OutputDataT]):
     def __init__(self, *, allows_deferred_tools: bool):
@@ -434,6 +500,10 @@ def __init__(self, *, allows_deferred_tools: bool):
     def mode(self) -> OutputMode:
         return 'image'
 
+    @cached_property
+    def json_schema(self) -> JsonSchema:
+        return self.build_json_schema()
+
 
 @dataclass(init=False)
 class StructuredTextOutputSchema(OutputSchema[OutputDataT], ABC):
@@ -450,6 +520,10 @@ def __init__(
         )
         self.processor = processor
 
+    @cached_property
+    def json_schema(self) -> JsonSchema:
+        return self.build_json_schema()
+
 
 class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]):
     @property
@@ -516,6 +590,10 @@ def __init__(
     def mode(self) -> OutputMode:
         return 'tool'
 
+    @cached_property
+    def json_schema(self) -> JsonSchema:
+        return self.build_json_schema()
+
 
 class BaseOutputProcessor(ABC, Generic[OutputDataT]):
     @abstractmethod

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -34,6 +34,7 @@
     UserPromptNode,
     capture_run_messages,
 )
+from .._json_schema import JsonSchema
 from .._output import OutputToolset
 from .._tool_manager import ToolManager
 from ..builtin_tools import AbstractBuiltinTool
@@ -955,6 +956,11 @@ def decorator(
             self._system_prompt_functions.append(_system_prompt.SystemPromptRunner[AgentDepsT](func, dynamic=dynamic))
             return func
 
+    def output_json_schema(self, output_type: OutputSpec[RunOutputDataT] | None = None) -> JsonSchema:
+        """The output JSON schema."""
+        output_schema = self._prepare_output_schema(output_type)
+        return output_schema.json_schema
+
     @overload
     def output_validator(
         self, func: Callable[[RunContext[AgentDepsT], OutputDataT], OutputDataT], /

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -23,6 +23,7 @@
     result,
     usage as _usage,
 )
+from .._json_schema import JsonSchema
 from .._tool_manager import ToolManager
 from ..builtin_tools import AbstractBuiltinTool
 from ..output import OutputDataT, OutputSpec
@@ -122,6 +123,11 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def output_json_schema(self, output_type: OutputSpec[RunOutputDataT] | None = None) -> JsonSchema:
+        """The output JSON schema."""
+        raise NotImplementedError
+
     @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
@@ -67,6 +68,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[RunOutputDataT] | None = None) -> JsonSchema:
+        return self.wrapped.output_json_schema(output_type=output_type)
+
     @overload
     def iter(
         self,

tests/test_agent.py

@@ -5248,6 +5248,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': {'type': 'integer'}, 'b': {'type': 'string'}},
+            'title': 'Foo',
+            'required': ['a', 'b'],
+        }
+    )
+    assert wrapper_agent.output_json_schema(output_type=str) == snapshot({'type': 'string'})
 
     bar_toolset = FunctionToolset()