Support desk project: the finished agent

from dataclasses import dataclass

from pydantic_ai import Agent, ApprovalRequired, DeferredToolRequests, ModelRetry, RunContext


@dataclass
class Desk:
    orders: dict[str, dict]
    refund_limit: float = 50.0


agent = Agent(
    "openai:gpt-4.1",
    defer_model_check=True,
    deps_type=Desk,
    output_type=[str, DeferredToolRequests],
    instructions="You answer support tickets for an online shop. Be short and kind.",
)


@agent.tool
def lookup_order(ctx: RunContext[Desk], order_id: str) -> str:
    """Look up an order.

    Args:
        order_id: The order id, like A-1001.
    """
    order = ctx.deps.orders.get(order_id)
    if order is None:
        raise ModelRetry(f"There is no order {order_id}.")
    return f"Order {order_id} is {order['status']}"


@agent.tool
def refund_order(ctx: RunContext[Desk], order_id: str, amount: float) -> str:
    """Refund part or all of an order.

    Args:
        order_id: The order id, like A-1001.
        amount: The amount to refund, in euros.
    """
    order = ctx.deps.orders.get(order_id)
    if order is None:
        raise ModelRetry(f"There is no order {order_id}.")
    if amount > order["total"]:
        raise ModelRetry(f"Order {order_id} only cost {order['total']:.2f} euros.")
    if amount > ctx.deps.refund_limit and not ctx.tool_call_approved:
        raise ApprovalRequired(metadata={"reason": f"{amount:.2f} euros is over the limit"})
    order["refunded"] = amount
    return f"Refunded {amount:.2f} euros on order {order_id}"

import re

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


def desk_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].lower()
    last = messages[-1].parts[-1]
    order = re.search(r"a-\d{4}", ticket)
    tools = {t.name for t in info.function_tools}

    if last.part_kind == "user-prompt" and order:
        order_id = order.group().upper()
        if "refund" in ticket and "refund_order" in tools:
            amount = float(re.search(r"(\d+(?:\.\d+)?) euros", ticket).group(1))
            return ModelResponse(parts=[ToolCallPart("refund_order", {"order_id": order_id, "amount": amount})])
        return ModelResponse(parts=[ToolCallPart("lookup_order", {"order_id": order_id})])

    if last.part_kind == "retry-prompt":
        return ModelResponse(parts=[TextPart(f"Sorry, I could not do that. {last.content}")])

    if last.part_kind == "tool-return":
        return ModelResponse(parts=[TextPart(last.content.rstrip(".") + ".")])

    return ModelResponse(parts=[TextPart("Could you send me your order number, like A-1001?")])


desk_model = FunctionModel(desk_reply, model_name="desk")

from pydantic_ai import DeferredToolRequests, DeferredToolResults, models

from desk import Desk, agent
from desk_model import desk_model

models.ALLOW_MODEL_REQUESTS = False


def make_desk():
    return Desk(orders={"A-1001": {"status": "shipped", "total": 40.0},
                        "A-1002": {"status": "waiting for stock", "total": 120.0}})


def test_small_refund_needs_no_approval():
    desk = make_desk()
    with agent.override(model=desk_model):
        result = agent.run_sync("Refund 20 euros on A-1001", deps=desk)
    assert result.output == "Refunded 20.00 euros on order A-1001."
    assert desk.orders["A-1001"]["refunded"] == 20.0


def test_large_refund_waits_for_a_person():
    desk = make_desk()
    with agent.override(model=desk_model):
        result = agent.run_sync("Refund 90 euros on A-1002", deps=desk)
    assert isinstance(result.output, DeferredToolRequests)
    assert "refunded" not in desk.orders["A-1002"]

    call = result.output.approvals[0]
    with agent.override(model=desk_model):
        final = agent.run_sync(message_history=result.all_messages(), deps=desk,
                               deferred_tool_results=DeferredToolResults(approvals={call.tool_call_id: True}))
    assert desk.orders["A-1002"]["refunded"] == 90.0
    assert "Refunded 90.00" in final.output


def test_refund_above_order_total_is_refused():
    desk = make_desk()
    with agent.override(model=desk_model):
        result = agent.run_sync("Refund 45 euros on A-1001", deps=desk)
    assert "refunded" not in desk.orders["A-1001"]
    assert result.output == "Sorry, I could not do that. Order A-1001 only cost 40.00 euros."


def test_unknown_order():
    with agent.override(model=desk_model):
        result = agent.run_sync("Where is A-7777?", deps=make_desk())
    assert result.output == "Sorry, I could not do that. There is no order A-7777."

from pydantic_evals import Case, Dataset
from pydantic_evals.evaluators import Contains, Evaluator, EvaluatorContext

from desk import Desk, agent
from desk_model import desk_model


class ShortReply(Evaluator):
    def evaluate(self, ctx: EvaluatorContext) -> bool:
        return len(ctx.output) <= 80


async def answer(ticket: str) -> str:
    desk = Desk(orders={"A-1001": {"status": "shipped", "total": 40.0}})
    result = await agent.run(ticket, deps=desk, model=desk_model)
    return str(result.output)


dataset = Dataset(
    name="desk",
    cases=[
        Case(name="status", inputs="Where is A-1001?", evaluators=[Contains("shipped")]),
        Case(name="small refund", inputs="Refund 10 euros on A-1001", evaluators=[Contains("Refunded 10.00")]),
        Case(name="no order id", inputs="Where is my parcel?", evaluators=[Contains("order number")]),
        Case(name="lower case id", inputs="where is a-1001", evaluators=[Contains("shipped")]),
        Case(name="unknown order", inputs="Where is A-7777?", evaluators=[Contains("no order A-7777")]),
        Case(name="money back", inputs="Can I get 25 euros back on A-1001?", evaluators=[Contains("Refunded 25.00")]),
    ],
    evaluators=[ShortReply()],
)

report = dataset.evaluate_sync(answer, progress=False)
for case in report.cases:
    failed = [name for name, result in case.assertions.items() if not result.value]
    print(f"{case.name:14} {case.output}")
    if failed:
        print("    failed:", failed)
print(f"passed: {report.averages().assertions:.1%}")