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,47 @@ 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 JSON schema."""
+ if output_type is None:
+ output_type = self.output_type
+
+ # call this first to force output_type to an iterable
+ output_type = list(_flatten_output_spec(output_type))
+
+ # flatten special outputs
+ for i, _ in enumerate(output_type):
+ if isinstance(_, _output.NativeOutput):
+ output_type[i] = _flatten_output_spec(_.outputs)
+ if isinstance(_, _output.PromptedOutput):
+ output_type[i] = _flatten_output_spec(_.outputs)
+ if isinstance(_, _output.ToolOutput):
+ output_type[i] = _flatten_output_spec(_.output)
+
+ # final flattening
+ output_type = _flatten_output_spec(output_type)
+
+ json_schemas: list[JsonSchema] = []
+ for _ in output_type:
+ if inspect.isfunction(_) or inspect.ismethod(_):
+ json_schema = TypeAdapter(inspect.signature(_).return_annotation).json_schema(mode='serialization')
+ elif isinstance(_, _output.TextOutput):
+ json_schema = TypeAdapter(str).json_schema(mode='serialization')
+ else:
+ json_schema = TypeAdapter(_).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
  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,

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()