Fix function output

Author: g-eoj

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -19,7 +19,6 @@
 
 from .. import (
     _agent_graph,
-    _function_schema,
     _output,
     _system_prompt,
     _utils,
@@ -981,9 +980,7 @@ def output_json_schema(self, output_type: OutputSpec[OutputDataT | RunOutputData
         json_schemas: list[JsonSchema] = []
         for _ in output_type:
             if inspect.isfunction(_) or inspect.ismethod(_):
-                function_schema = _function_schema.function_schema(_, GenerateToolJsonSchema)
-                json_schema = function_schema.json_schema
-                json_schema['description'] = function_schema.description
+                json_schema = TypeAdapter(inspect.signature(_).return_annotation).json_schema(mode='serialization')
             elif isinstance(_, _output.TextOutput):
                 json_schema = TypeAdapter(str).json_schema(mode='serialization')
             else:

tests/test_agent_output_schemas.py

@@ -41,15 +41,7 @@ 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,
-        }
-    )
+    assert agent.output_json_schema() == snapshot({'type': 'integer'})
 
 
 async def test_auto_output_json_schema():