Add output_json_schema property to Agents

Author: g-eoj

pydantic_ai_slim/pydantic_ai/_output.py

@@ -15,6 +15,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 +227,9 @@ def mode(self) -> OutputMode:
     def allows_text(self) -> bool:
         return self.text_processor is not None
 
+    def dump(self) -> JsonSchema:
+        raise NotImplementedError()
+
     @classmethod
     def build(  # noqa: C901
         cls,
@@ -405,6 +409,13 @@ def __init__(
     def mode(self) -> OutputMode:
         return 'auto'
 
+    def dump(self) -> JsonSchema:
+        if self.toolset:
+            toolset_processors = [self.toolset.processors[k] for k in self.toolset.processors]
+            processors_union = UnionOutputProcessor(toolset_processors).object_def.json_schema
+            return processors_union
+        return self.processor.object_def.json_schema
+
 
 @dataclass(init=False)
 class TextOutputSchema(OutputSchema[OutputDataT]):
@@ -425,6 +436,9 @@ def __init__(
     def mode(self) -> OutputMode:
         return 'text'
 
+    def dump(self) -> JsonSchema:
+        return {'type': 'string'}
+
 
 class ImageOutputSchema(OutputSchema[OutputDataT]):
     def __init__(self, *, allows_deferred_tools: bool):
@@ -450,6 +464,9 @@ def __init__(
         )
         self.processor = processor
 
+    def dump(self) -> JsonSchema:
+        return self.object_def.json_schema
+
 
 class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]):
     @property
@@ -523,6 +540,11 @@ def __init__(
     def mode(self) -> OutputMode:
         return 'tool'
 
+    def dump(self) -> JsonSchema:
+        toolset_processors = [self.toolset.processors[k] for k in self.toolset.processors]
+        processors_union = UnionOutputProcessor(toolset_processors).object_def.json_schema
+        return processors_union
+
 
 class BaseOutputProcessor(ABC, Generic[OutputDataT]):
     @abstractmethod
@@ -725,7 +747,10 @@ def __init__(
         json_schemas: list[ObjectJsonSchema] = []
         self._processors = {}
         for output in outputs:
-            processor = ObjectOutputProcessor(output=output, strict=strict)
+            if isinstance(output, ObjectOutputProcessor):
+                processor = output
+            else:
+                processor = ObjectOutputProcessor(output=output, strict=strict)
             object_def = processor.object_def
 
             object_key = object_def.name or output.__name__

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
@@ -391,6 +392,11 @@ def deps_type(self) -> type:
         """The type of dependencies used by the agent."""
         return self._deps_type
 
+    @property
+    def output_json_schema(self) -> JsonSchema:
+        """The output JSON schema."""
+        return self._output_schema.dump()
+
     @property
     def output_type(self) -> OutputSpec[OutputDataT]:
         """The type of data output by agent runs, used to validate the data returned by the model, defaults to `str`."""

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
@@ -101,6 +102,12 @@ def deps_type(self) -> type:
         """The type of dependencies used by the agent."""
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def output_json_schema(self) -> JsonSchema:
+        """The output JSON schema."""
+        raise NotImplementedError
+
     @property
     @abstractmethod
     def output_type(self) -> OutputSpec[OutputDataT]:

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
@@ -45,6 +46,10 @@ def name(self) -> str | None:
     def name(self, value: str | None) -> None:
         self.wrapped.name = value
 
+    @property
+    def output_json_schema(self) -> JsonSchema:
+        return self.wrapped.output_json_schema
+
     @property
     def deps_type(self) -> type:
         return self.wrapped.deps_type