Also include return (type) json schema for tools (similar how tool parameters are handled)

[LysanderKie]

Description

Feature Request: Include Tool Return Type Schema in Tool Definitions (passed into LLM call as part of tools json)

Background
We use pydantic-ai as an agent framework in a complex business domain. It's critical for us that in addition to the tool parameters schema and description, the schema and description of tool return types are also included in the JSON passed to the LLM as part of tool definitions.

Current Problem: The LLM doesn't know which properties/fields are available in a tool's return value without actually executing the tool call. This leads to:

• Unnecessary tool calls just to discover what data is available
• The LLM being unable to plan multi-step operations effectively
• Confusion about which tool provides which data points

Use Cases

1. Rich DTOs as Tool Returns
   All of our tools return DTOs (Pydantic models) instead of plain strings or primitives. These models contain detailed information about domain entities like invoices, projects, and contracts.

2. Sequential (Chained) Tool Calls
   This is foundational for the sequential tool call improvements discussed in Code Mode: Tool Injection for mcp-run-python: Enable Python code to call back to agent tools with request/response flow #2037. If the LLM knows what data a tool returns, it can:

   • Plan a sequence of tool calls where one tool's output feeds into another
   • Determine upfront if a requested data point is available
   • Chain operations more intelligently (e.g., "get invoice details, then get the contract details for that invoice")

3. Field Descriptions in Complex Domains
   In complex business domains, field names alone are often insufficient for the LLM to understand what data represents. For example:

   • audited_amount vs approved_amount vs paid_amount on an invoice
   • date_received vs date_posted vs date_due
   • Domain-specific terminology that needs context

   By including field descriptions from Pydantic models in the return schema, the LLM can:

   • Understand the semantic meaning of each field
   • Provide more accurate answers about what data is available
   • Better explain differences between similar-sounding fields to users

4. Better User Experience
   With return schemas available, the LLM can:

   • Politely decline requests for data that isn't available in any tool return
   • Suggest alternative data points that are available
   • Be more precise about what information it can provide

Proposed Solution
Similar to how parameters work today, each tool definition would include a returns field containing the JSON schema of the return type.

How this could be implemented

1. Auto-generate from type annotations: Since tools must return JSON-serializable data, we can extract the return type from function signatures and generate the schema
2. Docstring parsing: Parse return type descriptions from docstrings (similar to current parameter handling)
3. Opt-in flag: To avoid breaking changes and manage context length, introduce an opt-in flag (e.g., include_return_schema=True) in ToolDefinition or during tool preparation

I would be curious if you already thought about this and have a preferred way of how this should be implemented in mind.

Example: Current vs. Proposed

Current tool definition:

{
    "type": "function",
    "function": {
        "name": "get_contract_details",
        "description": "Get details about a specific contract by ID.",
        "parameters": {
            "additionalProperties": false,
            "properties": {
                "contract_id": {
                    "type": "integer",
                    "description": "ID of the contract"
                }
            },
            "required": ["contract_id"],
            "type": "object",
        },
        "strict": true
    }
}

Proposed with return schema:

{
    "type": "function",
    "function": {
        "name": "get_contract_details",
        "description": "Get details about a specific contract by ID.",
        "parameters": {
            "additionalProperties": false,
            "properties": {
                "contract_id": {
                    "type": "integer"
                }
            },
            "required": ["contract_id"],
            "type": "object"
        },
        "strict": true,
        "returns": {
            // Either with 'properties' or using #/$defs/... which is used for parameters
            "type": "object",
            "description": "Details about a contract within a project",
            "properties": {
                "id": {
                    "type": "integer",
                    "description": "Unique contract identifier"
                },
                "project_id": {
                    "type": "integer",
                    "description": "ID of the parent project"
                },
                "contractor_name": {
                    "type": "string",
                    "description": "Name of the contractor company"
                },
                "amount": {
                    "type": "number",
                    "description": "Total contract amount"
                },
                "signed_date": {
                    "type": "string",
                    "format": "date",
                    "description": "Date the contract was signed"
                }
            },
            "required": ["id", "project_id", "contractor_name", "amount", "signed_date"]
        }
    }
}

Minimal Example

Here's a minimal example demonstrating the need, especially for complex domain models with field descriptions:

from datetime import date
from decimal import Decimal
import pydantic
from pydantic_ai import Agent, RunContext

class InvoiceDetails(pydantic.BaseModel):
    """Details about an invoice for a contract"""
    id: int = pydantic.Field(description="Unique invoice identifier")
    contract_id: int = pydantic.Field(description="ID of the associated contract")
    invoice_number: str = pydantic.Field(description="External invoice number")
    
    # Multiple amount fields - descriptions are crucial for LLM to understand differences
    initial_amount: Decimal = pydantic.Field(description="Initial invoice amount as submitted by contractor")
    audited_amount: Decimal = pydantic.Field(description="Amount verified after checking the invoice")
    approved_amount: Decimal = pydantic.Field(description="Amount released for payment after approval")
    paid_amount: Decimal = pydantic.Field(description="Amount actually paid out")
    
    # Multiple date fields - descriptions clarify the workflow
    date_received: date = pydantic.Field(description="Date invoice was received")
    date_posted: date = pydantic.Field(description="Date invoice was posted in accounting")
    status: str = pydantic.Field(description="Invoice status (new, approved, paid)")

def get_invoice_details(ctx: RunContext, invoice_id: int) -> InvoiceDetails:
    """Get details about a specific invoice by ID."""
    return InvoiceDetails(
        id=invoice_id,
        contract_id=1,
        invoice_number="INV-2024-001",
        initial_amount=Decimal("50000.00"),
        audited_amount=Decimal("48000.00"),
        approved_amount=Decimal("48000.00"),
        paid_amount=Decimal("48000.00"),
        date_received=date(2024, 3, 15),
        date_posted=date(2024, 3, 20),
        status="approved",
    )

agent = Agent("openai:gpt-4o", tools=[get_invoice_details])

# Without return schema: 
# - LLM doesn't know what fields InvoiceDetails contains
# - Can't distinguish between initial_amount, audited_amount, approved_amount, paid_amount
# - User asks "what's the invoice amount?" - which one do they mean?

# With return schema + field descriptions:
# - LLM knows exactly what fields are available
# - Understands the difference between amounts (initial vs audited vs approved vs paid)
# - Can ask clarifying questions or provide the most relevant amount based on context

Related Issues

• Code Mode: Tool Injection for mcp-run-python: Enable Python code to call back to agent tools with request/response flow #2037 - Sequential tool calls (this feature would be foundational for that)

Open Questions

1. Opt-in vs. opt-out: Should this be enabled by default or require explicit configuration?
2. Context length concerns: Should there be automatic truncation/summarization for very large return schemas?
3. Docstring parsing: How should we handle return type descriptions from docstrings?

I'm happy to contribute to implementing this feature if there's interest! Let me know if you have preferences on the approach.

---

Note: This feature would significantly improve pydantic-ai's capabilities for complex, multi-step agentic workflows in production environments, especially in domains where field descriptions are essential for LLM comprehension.

References

No response

[DouweM]

@LysanderKie Thanks a ton for filing the issue! I've also been thinking about this in context of #2037, which I expect to dive into later this month.

1. Opt-in vs. opt-out: Should this be enabled by default or require explicit configuration?

The opt-in flag you suggested for inclusion in the tool definition that's sent to models makes sense to me:

Opt-in flag: To avoid breaking changes and manage context length, introduce an opt-in flag (e.g., include_return_schema=True) in ToolDefinition or during tool preparation

I don't know if model APIs already have a dedicated API field for this, or if we'd need to inject it into the description. If it's the latter, it should always be opt-in to not be a breaking/context-enlarging change. If there's a dedicated API, I could be convinced to always send it.

Either way, we'll have the schema available as a field on the ToolDefinition, and can use it in the generated function signatures in the RunPythonToolset I suggested in #2037.

3. Context length concerns: Should there be automatic truncation/summarization for very large return schemas?

We should try to handle this using at least $defs/$refs.

LLMs may be able to "resolve" them across the "tool return schema (possibly in description)" <> "tool arguments schema" boundary, or at least reference $defs from other tool descriptions which I assume they all get in context.

For this to work right, include_return_schema may need to be turned on for an entire toolset, as otherwise the defs may be missing.

5. Docstring parsing: How should we handle return type descriptions from docstrings?

We should merge them into the schema as field descriptions, like we do for arguments in the docstring.

---

Parsing schemas from the return type should be reasonably straightforward, as we already parse arguments. We will need to take some special consideration with:

• MCP tools, that may have an outputSchema in their ToolDefinition.metadata. We should verify it works with both FastAPIs (from the mcp and fastmcp packages)
• ToolReturn which has a return_value: Any that holds the value returned to the model as the tool result, and content that can hold things like FileUrl and BinaryContent, which will be returned to the model as follow-up user parts because tools can not return them directly. What should the output schema be? When chaining tool calls, the value passed around should probably be the content, or both. So we'd need to create a list/union schema of str | UserContent.
• Tools can also have a return type of FileUrl, BinaryContent, or one of their subclasses, which may also require special consideration.
• Tools that return a BaseModel or dataclass.
• Handoff tools that return (await subagent.run(...)).result and have a complex signature including things like DeferredToolRequests (see OutputSpec in the codebase).
  • We may have to handle those deferred tools (frontend tools or approval-required tools) in a special way when used in chaining context, perhaps by causing the super-agent to end on DeferredToolRequests itself.
  • We may also need a way to generate a JSON schema for any agent/agent._output_schema: OutputSchema): Add Agent.to_mcp() method #3076 (comment)

Pydantic TypeAdapter would likely help with a lot of this, and there may be more special things we support as tool returns that we'd need to handle :)

I'm happy to contribute to implementing this feature if there's interest! Let me know if you have preferences on the approach.

That'd be great! I'm sure we'll come up with more edge cases in code review :)

The first feature I want to build with this (which you/we may end up implementing in the process to verify that the LLM can use chaining across all these types) is #3093, which we could implement using https://github.com/danthedeckie/simpleeval. That'd probably only work for cases where the argument and output schemas match cleanly, but would still be useful in some cases.

[phemmer]

I ran across this as well after noticing that my agents were having difficulty using tools, acting as if they weren't aware of how to interpret the tool output.

I think for MCP tools, including the output schema must be a default behavior. outputSchema is part of the MCP spec. Not passing that data to the agent is unexpected behavior IMHO.

For local (non-MCP) tools, I can see how it might be debatable, as tools might be derived from normal python functions which weren't intended to be exposed to an LLM, and thus may contain large and verbose output type schemas. However (again, IMHO), I think for consistency if MCP tools are going to include it, then non-MCP tools should as well.

However, as a potential alternative that might satisfy the reason behind both of the arguments I made above: maybe the tool output schema should only be provided in the result of the tool call. This way the agent can see and understand the schema only when it needs to interpret the output of the tool. This satisfies the MCP argument as the schema is passed to the agent and not discarded, and it ensures consistency between the MCP and non-MCP tools. In fact, I might argue this is the ideal behavior, as if the schema is only provided on the tool itself in the tool list, then when the LLM receives the result of the tool, it's possible that tool may no longer be in the tool list, and thus the schema is not available. Including it in the tool result handles that scenario (I actually have this exact behavior in one of my toolsets).
As another "however", I think it should still be possible to enable the tool output definition in the tool list. I have situations where the agent needs to know the structure of a tool's output before the tool is called (#3093 is very similar to one of those situations). But if the chosen solution is to have the schema in the output by default, but not the list, I would be fine with this behavior needing to be explicitly enabled.

As for how to control this behavior, the option that makes sense to me is a toolset filter. Basically a subclass of WrapperToolset which sets the desired schema format. Like if the default is to not include output schema in the tool list, then there could be a OutputSchemaEnabledInToolListWrapperToolset (not literally that name, but you get the point).
My reason for this is that I don't see much other reasonable alternatives. I don't think adding some flag to the Agent (or Agent.run) is viable, as this only allows you to turn it on or off for everything, where as it may be desired to turn it on or off for individual tools or toolsets. And putting it on the toolset or even the tool construction isn't viable, as toolsets may be provided by external libraries, so you may not get the ability to control their construction parameters. A WrapperToolset based option allows you to be as granular as you like, and can be used at any point. It can be used immediately after the toolset is constructed, when constructing the Agent, when calling Agent.run(), or anywhere in between. You can use it on an individual toolset. You can use FilteredToolset for individual tools. You can act on multiple toolsets by combining them with CombinedToolset first. Etc...

[DouweM]

I think for MCP tools, including the output schema must be a default behavior. outputSchema is part of the MCP spec. Not passing that data to the agent is unexpected behavior IMHO.

@phemmer If the model APIs had a dedicated output_schema field on tool definitions, I would agree with you that we should send it by default and could most likely justify any additional token usage as the model API's responsibility -- we're just giving it the full picture. But as we can only include it on the tool description, "smuggling" it in like that feels like it should be opt-in, especially as significant token usage increases are arguably a breaking change we can't make until v2.

Your point about whether to send this in the list, with tool results, or both, is a good one. But there's no dedicated way in the model APIs to send metadata along with a tool result, so it's not clear how we would send it. Including it in the tool list is most often going to be what we want, I think, so it can be taken into account when making a plan / choosing a tool, although having it available when interpreting the tool return would be valuable as well... I'm not sure yet how to best do that in the scenario you pointed out where the tool is no longer available in the list.

Either way, I'm imaging a ToolDefinition.return_schema field holding the schema itself (which MCPServer can start setting today in addition to metadata['outputSchema']), and then you can start with something like this, which is pretty straightforward use of PreparedToolset if you've looked into ToolDefinitions in the past:

toolset = PreparedToolset(lambda tool_def: replace(tool_def, description='\n\n'join([tool_def.description, 'Return schema:', str(tool_def.return_schema)]) if tool_def.return_schema else tool_def)

A dedicated toolset for that could be useful, but I worry about the unwieldy name. On MCPServer and FunctionToolset (once we can read these from type hints), I'd also expect a way to enable this for the entire toolset with a constructor flag, especially because it's a native MCP feature.

So I think the first PR here could be to add the info on ToolDefinition, plus an option on MCPServer to include them in tool descriptions.

Then we can work out the rest once we start inferring output schemas from tool functions.

[phemmer]

But there's no dedicated way in the model APIs to send metadata along with a tool result

Well, the counter point to that is that the API doesn't even handle structured tool results at all. The tool call content field is plain text. We just serialize the structured data from the tool and stuff it in there. But there's nothing that says the content can't be an object with a schema and an output, or even text with <schema>...</schema><output>...</output>. It's up to the LLM to interpret it.

[DouweM]

@phemmer True, I agree we could make that opt-in as well.