More test stuff

Author: g-eoj

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

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

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
 

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

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',