FunctionModel: write a stand-in model

from pydantic_ai import Agent, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel


def reply(messages, info: AgentInfo) -> ModelResponse:
    ticket = messages[-1].parts[-1].content
    return ModelResponse(parts=[TextPart(f"You wrote {len(ticket.split())} words.")])


agent = Agent(FunctionModel(reply))
print(agent.run_sync("I was charged twice for one order").output)

def peek(messages, info: AgentInfo) -> ModelResponse:
    print("instructions:", info.instructions)
    print("tools:", [tool.name for tool in info.function_tools])
    print("text allowed:", info.allow_text_output)
    return ModelResponse(parts=[TextPart("ok")])


agent = Agent(FunctionModel(peek), instructions="Be short.")
agent.run_sync("hello")

import re

from pydantic_ai import ModelResponse, TextPart, ToolCallPart
from pydantic_ai.models.function import AgentInfo, FunctionModel


def sort_ticket(text):
    text = text.lower()
    if "charged" in text or "refund" in text:
        return "billing", 4
    if "parcel" in text or "arrived" in text:
        return "shipping", 3
    return "other", 1


def shop_reply(messages, info: AgentInfo) -> ModelResponse:
    prompts = [p.content for m in messages for p in m.parts if p.part_kind == "user-prompt"]
    ticket = prompts[-1]
    last = messages[-1].parts[-1]
    order = re.search(r"A-\d{4}", ticket)

    # 1. The ticket names an order and the agent has a tool: ask for it.
    if order and info.function_tools and last.part_kind == "user-prompt":
        tool = info.function_tools[0].name
        return ModelResponse(parts=[ToolCallPart(tool, {"order_id": order.group()})])

    # 2. A tool answered: write the reply from what it said.
    if last.part_kind == "tool-return" and info.allow_text_output:
        return ModelResponse(parts=[TextPart(f"Order {order.group()}: {last.content}.")])

    # 3. The agent wants a typed answer: fill in its output tool.
    category, priority = sort_ticket(ticket)
    if info.output_tools:
        args = {"category": category, "priority": priority}
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args)])

    # 4. Otherwise, plain text.
    return ModelResponse(parts=[TextPart(f"Sorted as {category}.")])


shop_model = FunctionModel(shop_reply, model_name="shop")

from pydantic_ai import Agent

from shop_model import shop_model

agent = Agent(shop_model)
for ticket in ["I was charged twice for one order", "My parcel never arrived", "How do I change my email?"]:
    print(agent.run_sync(ticket).output)