From 0ee01f18e9c63c96de29c2cf08ec983627ea9471 Mon Sep 17 00:00:00 2001 From: Joe G Date: Mon, 17 Nov 2025 10:22:04 -0800 Subject: [PATCH 1/9] Add output_json_schema property to Agents --- pydantic_ai_slim/pydantic_ai/_output.py | 27 ++++++++++++++++++- .../pydantic_ai/agent/__init__.py | 6 +++++ .../pydantic_ai/agent/abstract.py | 7 +++++ pydantic_ai_slim/pydantic_ai/agent/wrapper.py | 5 ++++ 4 files changed, 44 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index ebb737a1cf..919ba770b3 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/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__ diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 4cd353b44a..92dbd1e9e8 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/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`.""" diff --git a/pydantic_ai_slim/pydantic_ai/agent/abstract.py b/pydantic_ai_slim/pydantic_ai/agent/abstract.py index c7c1cb2b5c..c8de939707 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/abstract.py +++ b/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]: diff --git a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py index 38e832fa2b..ecb31dc802 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py +++ b/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 From 80dd6e38715bc0560fab76fb8d91e1f412ff2df4 Mon Sep 17 00:00:00 2001 From: Joe G Date: Mon, 17 Nov 2025 10:42:44 -0800 Subject: [PATCH 2/9] Add test outline --- tests/test_agent.py | 62 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_agent.py b/tests/test_agent.py index d8323f3c98..e93d02bd33 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -6125,3 +6125,65 @@ def llm(messages: list[ModelMessage], _info: AgentInfo) -> ModelResponse: ] ) assert run.all_messages_json().startswith(b'[{"parts":[{"content":"Hello",') + + +async def test_text_output_json_schema(): + def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart('')]) + + agent = Agent(FunctionModel(llm)) + output_schema = {'type': 'string'} + assert agent.output_json_schema == output_schema + + +async def test_tool_output_json_schema(): + def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart('')]) + + agent = Agent(FunctionModel(llm), output_type=[ToolOutput(bool)]) + assert agent.output_json_schema + + agent = Agent(FunctionModel(llm), output_type=[ToolOutput(bool), ToolOutput(bool)]) + assert agent.output_json_schema + + +async def test_native_output_json_schema(): + def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart('')]) + + agent = Agent( + FunctionModel(llm), + output_type=NativeOutput([bool], name='native_output_name', description='native_output_description') + ) + assert agent.output_json_schema + + +async def test_prompted_output_json_schema(): + def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart('')]) + + agent = Agent( + FunctionModel(llm), + output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description') + ) + assert agent.output_json_schema + + +async def test_custom_output_json_schema(): + def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart('')]) + + HumanDict = StructuredDict( + { + 'type': 'object', + 'properties': { + 'name': {'type': 'string'}, + 'age': {'type': 'integer'} + }, + 'required': ['name', 'age'] + }, + name='Human', + description='A human with a name and age', + ) + agent = Agent(FunctionModel(llm), output_type=HumanDict) + assert agent.output_json_schema From 7127de623484f1d3de5928983bba5cc6f5edd41a Mon Sep 17 00:00:00 2001 From: Joe G Date: Mon, 17 Nov 2025 10:57:11 -0800 Subject: [PATCH 3/9] Fix typecheck errors --- pydantic_ai_slim/pydantic_ai/_output.py | 10 +++++----- tests/test_agent.py | 15 ++++++--------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 919ba770b3..807c82233c 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -465,7 +465,7 @@ def __init__( self.processor = processor def dump(self) -> JsonSchema: - return self.object_def.json_schema + return self.object_def.json_schema # pyright: ignore[reportOptionalMemberAccess] class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]): @@ -541,9 +541,9 @@ 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 + toolset_processors = [self.toolset.processors[k] for k in self.toolset.processors] # pyright: ignore[reportOptionalMemberAccess] + processors_union = UnionOutputProcessor(toolset_processors) + return processors_union.object_def.json_schema class BaseOutputProcessor(ABC, Generic[OutputDataT]): @@ -736,7 +736,7 @@ class UnionOutputProcessor(BaseObjectOutputProcessor[OutputDataT]): def __init__( self, - outputs: Sequence[OutputTypeOrFunction[OutputDataT]], + outputs: Sequence[OutputTypeOrFunction[OutputDataT]|ObjectOutputProcessor[OutputDataT]], *, name: str | None = None, description: str | None = None, diff --git a/tests/test_agent.py b/tests/test_agent.py index e93d02bd33..aa830bf16a 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -6152,8 +6152,8 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart('')]) agent = Agent( - FunctionModel(llm), - output_type=NativeOutput([bool], name='native_output_name', description='native_output_description') + FunctionModel(llm), + output_type=NativeOutput([bool], name='native_output_name', description='native_output_description'), ) assert agent.output_json_schema @@ -6163,8 +6163,8 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart('')]) agent = Agent( - FunctionModel(llm), - output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description') + FunctionModel(llm), + output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description'), ) assert agent.output_json_schema @@ -6176,11 +6176,8 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: HumanDict = StructuredDict( { 'type': 'object', - 'properties': { - 'name': {'type': 'string'}, - 'age': {'type': 'integer'} - }, - 'required': ['name', 'age'] + 'properties': {'name': {'type': 'string'}, 'age': {'type': 'integer'}}, + 'required': ['name', 'age'], }, name='Human', description='A human with a name and age', From d1b2399c094080dd924ce9aa40aca00671b17047 Mon Sep 17 00:00:00 2001 From: Joe G Date: Mon, 17 Nov 2025 12:48:07 -0800 Subject: [PATCH 4/9] Capture title and desc for tool output --- pydantic_ai_slim/pydantic_ai/_output.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 59b751cfb6..655ed48137 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -1,5 +1,6 @@ from __future__ import annotations as _annotations +import copy import inspect import json import re @@ -411,9 +412,13 @@ def mode(self) -> OutputMode: 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 + processors = [] + for tool_def in self.toolset._tool_defs: + processor = copy.copy(self.toolset.processors[tool_def.name]) + processor.object_def.name = tool_def.name + processor.object_def.description = tool_def.description + processors.append(processor) + return UnionOutputProcessor(processors).object_def.json_schema return self.processor.object_def.json_schema @@ -465,7 +470,7 @@ def __init__( self.processor = processor def dump(self) -> JsonSchema: - return self.object_def.json_schema # pyright: ignore[reportOptionalMemberAccess] + return self.processor.object_def.json_schema class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]): @@ -541,9 +546,13 @@ def mode(self) -> OutputMode: return 'tool' def dump(self) -> JsonSchema: - toolset_processors = [self.toolset.processors[k] for k in self.toolset.processors] # pyright: ignore[reportOptionalMemberAccess] - processors_union = UnionOutputProcessor(toolset_processors) - return processors_union.object_def.json_schema + processors = [] + for tool_def in self.toolset._tool_defs: + processor = copy.copy(self.toolset.processors[tool_def.name]) + processor.object_def.name = tool_def.name + processor.object_def.description = tool_def.description + processors.append(processor) + return UnionOutputProcessor(processors).object_def.json_schema class BaseOutputProcessor(ABC, Generic[OutputDataT]): From 4454583983b529219cc158e3df15de90d972b673 Mon Sep 17 00:00:00 2001 From: Joe G Date: Mon, 17 Nov 2025 12:48:47 -0800 Subject: [PATCH 5/9] Add tests --- tests/test_agent.py | 123 ++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 114 insertions(+), 9 deletions(-) diff --git a/tests/test_agent.py b/tests/test_agent.py index aa830bf16a..b11ecf85cd 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -6132,19 +6132,93 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart('')]) agent = Agent(FunctionModel(llm)) - output_schema = {'type': 'string'} - assert agent.output_json_schema == output_schema + assert agent.output_json_schema == snapshot({'type': 'string'}) async def test_tool_output_json_schema(): def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart('')]) - agent = Agent(FunctionModel(llm), output_type=[ToolOutput(bool)]) - assert agent.output_json_schema + agent = Agent( + FunctionModel(llm), + output_type=[ToolOutput(bool, name='alice', description='Dreaming...')], + ) + assert agent.output_json_schema == snapshot( + { + 'type': 'object', + 'properties': { + 'result': { + 'anyOf': [ + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'alice'}, + 'data': { + 'properties': {'response': {'type': 'boolean'}}, + 'required': ['response'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'alice', + 'description': 'Dreaming...', + } + ] + } + }, + 'required': ['result'], + 'additionalProperties': False, + } + ) - agent = Agent(FunctionModel(llm), output_type=[ToolOutput(bool), ToolOutput(bool)]) - assert agent.output_json_schema + agent = Agent( + FunctionModel(llm), + output_type=[ToolOutput(bool, name='alice'), ToolOutput(bool, name='bob')], + ) + assert agent.output_json_schema == snapshot( + { + 'type': 'object', + 'properties': { + 'result': { + 'anyOf': [ + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'alice'}, + 'data': { + 'properties': {'response': {'type': 'boolean'}}, + 'required': ['response'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'alice', + 'description': 'bool: The final response which ends this conversation', + }, + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'bob'}, + 'data': { + 'properties': {'response': {'type': 'boolean'}}, + 'required': ['response'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'bob', + 'description': 'bool: The final response which ends this conversation', + }, + ] + } + }, + 'required': ['result'], + 'additionalProperties': False, + } + ) async def test_native_output_json_schema(): @@ -6155,7 +6229,9 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: FunctionModel(llm), output_type=NativeOutput([bool], name='native_output_name', description='native_output_description'), ) - assert agent.output_json_schema + assert agent.output_json_schema == snapshot( + {'properties': {'response': {'type': 'boolean'}}, 'required': ['response'], 'type': 'object'} + ) async def test_prompted_output_json_schema(): @@ -6166,7 +6242,9 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: FunctionModel(llm), output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description'), ) - assert agent.output_json_schema + assert agent.output_json_schema == snapshot( + {'properties': {'response': {'type': 'boolean'}}, 'required': ['response'], 'type': 'object'} + ) async def test_custom_output_json_schema(): @@ -6183,4 +6261,31 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: description='A human with a name and age', ) agent = Agent(FunctionModel(llm), output_type=HumanDict) - assert agent.output_json_schema + assert agent.output_json_schema == snapshot( + { + 'type': 'object', + 'properties': { + 'result': { + 'anyOf': [ + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'final_result'}, + 'data': { + 'properties': {'name': {'type': 'string'}, 'age': {'type': 'integer'}}, + 'required': ['name', 'age'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'final_result', + 'description': 'A human with a name and age', + } + ] + } + }, + 'required': ['result'], + 'additionalProperties': False, + } + ) From fdc88207bdc173f05579accbc08ffd724df57930 Mon Sep 17 00:00:00 2001 From: Joe G Date: Mon, 17 Nov 2025 13:16:49 -0800 Subject: [PATCH 6/9] Fix typecheck errors --- pydantic_ai_slim/pydantic_ai/_output.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 655ed48137..848fae752a 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -412,8 +412,8 @@ def mode(self) -> OutputMode: def dump(self) -> JsonSchema: if self.toolset: - processors = [] - for tool_def in self.toolset._tool_defs: + processors: list[ObjectOutputProcessor[OutputDataT]] = [] + for tool_def in self.toolset._tool_defs: # pyright: ignore [reportPrivateUsage] processor = copy.copy(self.toolset.processors[tool_def.name]) processor.object_def.name = tool_def.name processor.object_def.description = tool_def.description @@ -546,13 +546,16 @@ def mode(self) -> OutputMode: return 'tool' def dump(self) -> JsonSchema: - processors = [] - for tool_def in self.toolset._tool_defs: - processor = copy.copy(self.toolset.processors[tool_def.name]) - processor.object_def.name = tool_def.name - processor.object_def.description = tool_def.description - processors.append(processor) - return UnionOutputProcessor(processors).object_def.json_schema + if self.toolset: + processors: list[ObjectOutputProcessor[OutputDataT]] = [] + for tool_def in self.toolset._tool_defs: # pyright: ignore [reportPrivateUsage] + processor = copy.copy(self.toolset.processors[tool_def.name]) + processor.object_def.name = tool_def.name + processor.object_def.description = tool_def.description + processors.append(processor) + return UnionOutputProcessor(processors).object_def.json_schema + else: + raise RuntimeError('ToolOutputSchema has no toolset.') class BaseOutputProcessor(ABC, Generic[OutputDataT]): From 79253f1ea539c4090e2ba7488b7b01fc0d8b2d10 Mon Sep 17 00:00:00 2001 From: Joe G Date: Wed, 19 Nov 2025 08:08:10 -0800 Subject: [PATCH 7/9] Change to method --- pydantic_ai_slim/pydantic_ai/_output.py | 4 ++ .../pydantic_ai/agent/__init__.py | 10 ++-- .../pydantic_ai/agent/abstract.py | 11 ++--- pydantic_ai_slim/pydantic_ai/agent/wrapper.py | 7 ++- tests/test_agent.py | 49 ++++++++++++++++--- 5 files changed, 60 insertions(+), 21 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index 848fae752a..b8d311ecc0 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -228,6 +228,7 @@ def mode(self) -> OutputMode: def allows_text(self) -> bool: return self.text_processor is not None + @abstractmethod def dump(self) -> JsonSchema: raise NotImplementedError() @@ -453,6 +454,9 @@ def __init__(self, *, allows_deferred_tools: bool): def mode(self) -> OutputMode: return 'image' + def dump(self) -> JsonSchema: + raise NotImplementedError() + @dataclass(init=False) class StructuredTextOutputSchema(OutputSchema[OutputDataT], ABC): diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 92dbd1e9e8..599f2c2fb4 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -392,11 +392,6 @@ 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`.""" @@ -953,6 +948,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[OutputDataT] | None = None) -> JsonSchema: + """The output JSON schema.""" + output_schema = self._prepare_output_schema(output_type) + return output_schema.dump() + @overload def output_validator( self, func: Callable[[RunContext[AgentDepsT], OutputDataT], OutputDataT], / diff --git a/pydantic_ai_slim/pydantic_ai/agent/abstract.py b/pydantic_ai_slim/pydantic_ai/agent/abstract.py index c8de939707..d1970df6e0 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/abstract.py +++ b/pydantic_ai_slim/pydantic_ai/agent/abstract.py @@ -102,12 +102,6 @@ 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]: @@ -129,6 +123,11 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]: """ raise NotImplementedError + @abstractmethod + def output_json_schema(self, output_type: OutputSpec[OutputDataT] | None = None) -> JsonSchema: + """The output JSON schema.""" + raise NotImplementedError + @overload async def run( self, diff --git a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py index ecb31dc802..8291a32a8f 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py +++ b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py @@ -46,10 +46,6 @@ 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 @@ -72,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] | None = None) -> JsonSchema: + return self.wrapped.output_json_schema(output_type=output_type) + @overload def iter( self, diff --git a/tests/test_agent.py b/tests/test_agent.py index b11ecf85cd..fdc21e576b 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -6132,7 +6132,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart('')]) agent = Agent(FunctionModel(llm)) - assert agent.output_json_schema == snapshot({'type': 'string'}) + assert agent.output_json_schema() == snapshot({'type': 'string'}) async def test_tool_output_json_schema(): @@ -6143,7 +6143,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: FunctionModel(llm), output_type=[ToolOutput(bool, name='alice', description='Dreaming...')], ) - assert agent.output_json_schema == snapshot( + assert agent.output_json_schema() == snapshot( { 'type': 'object', 'properties': { @@ -6176,7 +6176,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: FunctionModel(llm), output_type=[ToolOutput(bool, name='alice'), ToolOutput(bool, name='bob')], ) - assert agent.output_json_schema == snapshot( + assert agent.output_json_schema() == snapshot( { 'type': 'object', 'properties': { @@ -6229,7 +6229,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: FunctionModel(llm), output_type=NativeOutput([bool], name='native_output_name', description='native_output_description'), ) - assert agent.output_json_schema == snapshot( + assert agent.output_json_schema() == snapshot( {'properties': {'response': {'type': 'boolean'}}, 'required': ['response'], 'type': 'object'} ) @@ -6242,7 +6242,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: FunctionModel(llm), output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description'), ) - assert agent.output_json_schema == snapshot( + assert agent.output_json_schema() == snapshot( {'properties': {'response': {'type': 'boolean'}}, 'required': ['response'], 'type': 'object'} ) @@ -6261,7 +6261,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: description='A human with a name and age', ) agent = Agent(FunctionModel(llm), output_type=HumanDict) - assert agent.output_json_schema == snapshot( + assert agent.output_json_schema() == snapshot( { 'type': 'object', 'properties': { @@ -6289,3 +6289,40 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: 'additionalProperties': False, } ) + + +async def test_override_output_json_schema(): + def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse(parts=[TextPart('')]) + + agent = Agent(FunctionModel(llm)) + assert agent.output_json_schema() == snapshot({'type': 'string'}) + output_type = ([ToolOutput(bool, name='alice', description='Dreaming...')],) + assert agent.output_json_schema(output_type=output_type) == snapshot( + { + 'type': 'object', + 'properties': { + 'result': { + 'anyOf': [ + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'alice'}, + 'data': { + 'properties': {'response': {'type': 'boolean'}}, + 'required': ['response'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'alice', + 'description': 'Dreaming...', + } + ] + } + }, + 'required': ['result'], + 'additionalProperties': False, + } + ) From 09184e2d526ea57a6638c188d613d8b4b0524200 Mon Sep 17 00:00:00 2001 From: Joe G Date: Wed, 19 Nov 2025 11:23:32 -0800 Subject: [PATCH 8/9] More test stuff --- pydantic_ai_slim/pydantic_ai/_output.py | 21 ++++--- .../pydantic_ai/agent/__init__.py | 2 +- .../pydantic_ai/agent/abstract.py | 2 +- pydantic_ai_slim/pydantic_ai/agent/wrapper.py | 2 +- tests/test_agent.py | 63 +++++++++++-------- 5 files changed, 51 insertions(+), 39 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index b8d311ecc0..50b6fe38eb 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -420,6 +420,7 @@ def dump(self) -> JsonSchema: processor.object_def.description = tool_def.description processors.append(processor) return UnionOutputProcessor(processors).object_def.json_schema + return self.processor.object_def.json_schema @@ -550,16 +551,16 @@ def mode(self) -> OutputMode: return 'tool' def dump(self) -> JsonSchema: - if self.toolset: - processors: list[ObjectOutputProcessor[OutputDataT]] = [] - for tool_def in self.toolset._tool_defs: # pyright: ignore [reportPrivateUsage] - processor = copy.copy(self.toolset.processors[tool_def.name]) - processor.object_def.name = tool_def.name - processor.object_def.description = tool_def.description - processors.append(processor) - return UnionOutputProcessor(processors).object_def.json_schema - else: - raise RuntimeError('ToolOutputSchema has no toolset.') + if self.toolset is None: + # need to check expected behavior + raise NotImplementedError() + processors: list[ObjectOutputProcessor[OutputDataT]] = [] + for tool_def in self.toolset._tool_defs: # pyright: ignore [reportPrivateUsage] + processor = copy.copy(self.toolset.processors[tool_def.name]) + processor.object_def.name = tool_def.name + processor.object_def.description = tool_def.description + processors.append(processor) + return UnionOutputProcessor(processors).object_def.json_schema class BaseOutputProcessor(ABC, Generic[OutputDataT]): diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 599f2c2fb4..54c98e87ab 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -948,7 +948,7 @@ 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] | None = None) -> JsonSchema: + 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.dump() diff --git a/pydantic_ai_slim/pydantic_ai/agent/abstract.py b/pydantic_ai_slim/pydantic_ai/agent/abstract.py index d1970df6e0..8c7c81ee41 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/abstract.py +++ b/pydantic_ai_slim/pydantic_ai/agent/abstract.py @@ -124,7 +124,7 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]: raise NotImplementedError @abstractmethod - def output_json_schema(self, output_type: OutputSpec[OutputDataT] | None = None) -> JsonSchema: + def output_json_schema(self, output_type: OutputSpec[RunOutputDataT] | None = None) -> JsonSchema: """The output JSON schema.""" raise NotImplementedError diff --git a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py index 8291a32a8f..6431be7297 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py +++ b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py @@ -68,7 +68,7 @@ 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] | None = None) -> JsonSchema: + def output_json_schema(self, output_type: OutputSpec[RunOutputDataT] | None = None) -> JsonSchema: return self.wrapped.output_json_schema(output_type=output_type) @overload diff --git a/tests/test_agent.py b/tests/test_agent.py index fbbe2072b3..18d7b84971 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -5155,6 +5155,35 @@ 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': { + 'result': { + 'anyOf': [ + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'final_result'}, + 'data': { + 'properties': {'a': {'type': 'integer'}, 'b': {'type': 'string'}}, + 'required': ['a', 'b'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'final_result', + 'description': 'The final response which ends this conversation', + } + ] + } + }, + 'required': ['result'], + 'additionalProperties': False, + } + ) + assert wrapper_agent.output_json_schema(output_type=str) == snapshot({'type': 'string'}) bar_toolset = FunctionToolset() @@ -6151,19 +6180,13 @@ def test_message_history_cannot_start_with_model_response(): async def test_text_output_json_schema(): - def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: - return ModelResponse(parts=[TextPart('')]) - - agent = Agent(FunctionModel(llm)) + agent = Agent('test') assert agent.output_json_schema() == snapshot({'type': 'string'}) async def test_tool_output_json_schema(): - def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: - return ModelResponse(parts=[TextPart('')]) - agent = Agent( - FunctionModel(llm), + 'test', output_type=[ToolOutput(bool, name='alice', description='Dreaming...')], ) assert agent.output_json_schema() == snapshot( @@ -6196,7 +6219,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: ) agent = Agent( - FunctionModel(llm), + 'test', output_type=[ToolOutput(bool, name='alice'), ToolOutput(bool, name='bob')], ) assert agent.output_json_schema() == snapshot( @@ -6245,11 +6268,8 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: async def test_native_output_json_schema(): - def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: - return ModelResponse(parts=[TextPart('')]) - agent = Agent( - FunctionModel(llm), + 'test', output_type=NativeOutput([bool], name='native_output_name', description='native_output_description'), ) assert agent.output_json_schema() == snapshot( @@ -6258,11 +6278,8 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: async def test_prompted_output_json_schema(): - def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: - return ModelResponse(parts=[TextPart('')]) - agent = Agent( - FunctionModel(llm), + 'test', output_type=PromptedOutput([bool], name='prompted_output_name', description='prompted_output_description'), ) assert agent.output_json_schema() == snapshot( @@ -6271,9 +6288,6 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: async def test_custom_output_json_schema(): - def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: - return ModelResponse(parts=[TextPart('')]) - HumanDict = StructuredDict( { 'type': 'object', @@ -6283,7 +6297,7 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: name='Human', description='A human with a name and age', ) - agent = Agent(FunctionModel(llm), output_type=HumanDict) + agent = Agent('test', output_type=HumanDict) assert agent.output_json_schema() == snapshot( { 'type': 'object', @@ -6315,12 +6329,9 @@ def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: async def test_override_output_json_schema(): - def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: - return ModelResponse(parts=[TextPart('')]) - - agent = Agent(FunctionModel(llm)) + agent = Agent('test') assert agent.output_json_schema() == snapshot({'type': 'string'}) - output_type = ([ToolOutput(bool, name='alice', description='Dreaming...')],) + output_type = [ToolOutput(bool, name='alice', description='Dreaming...')] assert agent.output_json_schema(output_type=output_type) == snapshot( { 'type': 'object', From 0d587620dc9ae168eff4156a34d4195830001fad Mon Sep 17 00:00:00 2001 From: Joe G Date: Wed, 19 Nov 2025 11:46:35 -0800 Subject: [PATCH 9/9] Improve test coverage --- tests/test_agent.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/test_agent.py b/tests/test_agent.py index 18d7b84971..fda72265cb 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -6184,6 +6184,38 @@ async def test_text_output_json_schema(): assert agent.output_json_schema() == snapshot({'type': 'string'}) +async def test_auto_output_json_schema(): + agent = Agent('test', output_type=bool) + assert agent.output_json_schema() == snapshot( + { + 'type': 'object', + 'properties': { + 'result': { + 'anyOf': [ + { + 'type': 'object', + 'properties': { + 'kind': {'type': 'string', 'const': 'final_result'}, + 'data': { + 'properties': {'response': {'type': 'boolean'}}, + 'required': ['response'], + 'type': 'object', + }, + }, + 'required': ['kind', 'data'], + 'additionalProperties': False, + 'title': 'final_result', + 'description': 'The final response which ends this conversation', + } + ] + } + }, + 'required': ['result'], + 'additionalProperties': False, + } + ) + + async def test_tool_output_json_schema(): agent = Agent( 'test',