Handle TextOutput

Author: g-eoj

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -984,6 +984,8 @@ def output_json_schema(self, output_type: OutputSpec[OutputDataT | RunOutputData
                 function_schema = _function_schema.function_schema(_, GenerateToolJsonSchema)
                 json_schema = function_schema.json_schema
                 json_schema['description'] = function_schema.description
+            elif isinstance(_, _output.TextOutput):
+                json_schema = TypeAdapter(str).json_schema(mode='serialization')
             else:
                 json_schema = TypeAdapter(_).json_schema(mode='serialization')
 

tests/test_agent_output_schemas.py

@@ -9,6 +9,7 @@
     NativeOutput,
     PromptedOutput,
     StructuredDict,
+    TextOutput,
     ToolOutput,
 )
 
@@ -28,6 +29,28 @@ async def test_text_output_json_schema():
     agent = Agent('test')
     assert agent.output_json_schema() == snapshot({'type': 'string'})
 
+    def func(x: str) -> str:
+        return x
+
+    agent = Agent('test', output_type=TextOutput(func))
+    assert agent.output_json_schema() == snapshot({'type': 'string'})
+
+
+async def test_function_output_json_schema():
+    def func(x: int) -> int:
+        return x
+
+    agent = Agent('test', output_type=[func])
+    assert agent.output_json_schema() == snapshot(
+        {
+            'additionalProperties': False,
+            'properties': {'x': {'type': 'integer'}},
+            'required': ['x'],
+            'type': 'object',
+            'description': None,
+        }
+    )
+
 
 async def test_auto_output_json_schema():
     # one output