Change to method

Author: g-eoj

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

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], /

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,

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,

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,
+        }
+    )