Add Agent.output_json_schema() method

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
+from pydantic import TypeAdapter
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Self, TypeVar, deprecated
 
@@ -34,7 +35,8 @@
     UserPromptNode,
     capture_run_messages,
 )
-from .._output import OutputToolset
+from .._json_schema import JsonSchema
+from .._output import OutputToolset, _flatten_output_spec  # pyright: ignore[reportPrivateUsage]
 from .._tool_manager import ToolManager
 from ..builtin_tools import AbstractBuiltinTool
 from ..models.instrumented import InstrumentationSettings, InstrumentedModel, instrument_model
@@ -955,6 +957,55 @@ def decorator(
             self._system_prompt_functions.append(_system_prompt.SystemPromptRunner[AgentDepsT](func, dynamic=dynamic))
             return func
 
+    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
+
+        # forces output_type to an iterable
+        output_type = _flatten_output_spec(output_type)
+
+        flat_output_type: list[OutputSpec[OutputDataT | RunOutputDataT] | type[str]] = []
+        for output_spec in output_type:
+            if isinstance(output_spec, _output.NativeOutput):
+                flat_output_type += _flatten_output_spec(output_spec.outputs)
+            elif isinstance(output_spec, _output.PromptedOutput):
+                flat_output_type += _flatten_output_spec(output_spec.outputs)
+            elif isinstance(output_spec, _output.TextOutput):
+                flat_output_type.append(str)
+            elif isinstance(output_spec, _output.ToolOutput):
+                flat_output_type += _flatten_output_spec(output_spec.output)
+            elif inspect.isfunction(output_spec) or inspect.ismethod(output_spec):
+                return_annotation = inspect.signature(output_spec).return_annotation
+                flat_output_type += _flatten_output_spec(return_annotation)
+            else:
+                flat_output_type.append(output_spec)
+
+        json_schemas: list[JsonSchema] = []
+        for output_spec in flat_output_type:
+            json_schema = TypeAdapter(output_spec).json_schema(mode='serialization')
+
+            if output_spec == _messages.BinaryImage:
+                properties = json_schema.get('properties', {})
+                kept_properties = {}
+                for key, value in properties.items():
+                    if key in ['data', 'media_type']:
+                        kept_properties[key] = value
+                json_schema['properties'] = kept_properties
+                json_schema.pop('required', None)
+
+            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
     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[OutputDataT | 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[OutputDataT | RunOutputDataT] | None = None) -> JsonSchema:
+        return self.wrapped.output_json_schema(output_type=output_type)
+
     @overload
     def iter(
         self,

pydantic_ai_slim/pydantic_ai/messages.py

@@ -611,14 +611,17 @@ def format(self) -> str:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+@dataclass(init=False, repr=False)
 class BinaryImage(BinaryContent):
     """Binary content that's guaranteed to be an image."""
 
+    media_type: ImageMediaType | str
+
     def __init__(
         self,
         data: bytes,
         *,
-        media_type: str,
+        media_type: ImageMediaType | str,
         identifier: str | None = None,
         vendor_metadata: dict[str, Any] | None = None,
         # Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.

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': {'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()