Output functions and several output types

An agent can accept more than one kind of answer, such as a ticket or a hand-over to a person. An output function runs your code on the answer before the run ends.

Example

from typing import Literal

from pydantic import BaseModel
from pydantic_ai import Agent, ModelResponse, ToolCallPart
from pydantic_ai.models.function import FunctionModel


class Ticket(BaseModel):
    category: Literal["billing", "shipping", "other"]
    priority: int


class NeedsHuman(BaseModel):
    reason: str


def triage(messages, info):
    ticket = messages[0].parts[-1].content
    if "lawyer" in ticket:
        return ModelResponse(parts=[ToolCallPart("final_result_NeedsHuman", {"reason": "legal threat"})])
    return ModelResponse(parts=[ToolCallPart("final_result_Ticket", {"category": "billing", "priority": 4})])

Example

agent = Agent(FunctionModel(triage), output_type=[Ticket, NeedsHuman])

for text in ["I was charged twice", "My lawyer will hear about this"]:
    output = agent.run_sync(text).output
    if isinstance(output, NeedsHuman):
        print("to a person:", output.reason)
    else:
        print("queued:", output)

A list of types gives the model one output tool per type, named final_result_Ticket and final_result_NeedsHuman, and it picks one. Your code checks which with isinstance. Add str to the list and a plain text answer is accepted too.

Output functions

Example

from pydantic_ai import Agent, ModelResponse, ToolCallPart
from pydantic_ai.models.function import FunctionModel

queue = []


def open_ticket(category: str, priority: int) -> str:
    """Put a ticket in the support queue."""
    queue.append({"category": category, "priority": priority})
    return f"T-{len(queue)}"

A model function that prints the output tool it is offered and calls it:

Example

def triage(messages, info):
    tool = info.output_tools[0]
    print("output tool:", tool.name, "-", tool.description)
    return ModelResponse(parts=[ToolCallPart(tool.name, {"category": "billing", "priority": 4})])

Example

agent = Agent(FunctionModel(triage), output_type=open_ticket)
result = agent.run_sync("I was charged twice")
print(result.output)
print(queue)

With a function as output_type, its parameters become the output tool's arguments and its docstring the tool's description. The model calls it to finish; Pydantic validates the arguments, the function runs, and what it returns is result.output. The model never sees that return value, unlike a tool's, because the run is over.

An output function can raise ModelRetry, like a validator, and can take RunContext as its first parameter. It is the place for work that should happen once the answer is known, such as saving it.

Try it yourself

Add a third type, Spam with no fields, to the list, and a rule in triage for it.
Put open_ticket and NeedsHuman in one list.
Raise ModelRetry in open_ticket when priority is above 5.

PreviousOutput validators: rules Pydantic cannot check

This is what real progress feels like.