Pydantic AI Agents

Useful evidence: a labelled test set, a confusion table, and examples of messages sent for human review. Show which mistakes are tolerable and which would send a learner to the wrong team.

Eligibility is a hard condition: the learner has completed the prerequisites, the exercise is published, and their account can access it. Preference helps choose between eligible options: they have twenty minutes, enjoy debugging, or want revision. Filter hard conditions in application code before asking a model to compare the remaining options.

Step 3: Expose one narrow tool

from dataclasses import dataclass, field
from typing import Literal

from pydantic import BaseModel
from pydantic_ai import Agent, RunContext, models
from pydantic_ai.messages import ModelResponse, ToolCallPart, ToolReturnPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.usage import UsageLimits

models.ALLOW_MODEL_REQUESTS = False
Status = Literal["available", "maintenance", "review"]
LABS = {
    "LAB-101": {"owner": "learner-1", "status": "available"},
    "LAB-202": {"owner": "learner-2", "status": "maintenance"},
}

class Reply(BaseModel):
    lab_id: str
    status: Status

@dataclass
class Context:
    user_id: str
    evidence: dict[str, str] = field(default_factory=dict)

def permitted_status(user_id: str, lab_id: str) -> str:
    record = LABS.get(lab_id)
    if record is None or record["owner"] != user_id:
        return "review"
    return record["status"]

def build_agent(model):
    agent = Agent(
        model,
        deps_type=Context,
        output_type=Reply,
        instructions="Look up the requested lab and report its status.",
    )

    @agent.tool
    def lookup_lab(ctx: RunContext[Context], lab_id: str) -> dict:
        status = permitted_status(ctx.deps.user_id, lab_id)
        ctx.deps.evidence[lab_id] = status
        return {"lab_id": lab_id, "status": status}

    return agent

def fixture(lab_id: str, forge_status: bool = False):
    def respond(messages, info: AgentInfo):
        returns = [
            part for message in messages for part in message.parts
            if isinstance(part, ToolReturnPart)
        ]
        if not returns:
            return ModelResponse([
                ToolCallPart("lookup_lab", {"lab_id": lab_id})
            ])
        candidate = dict(returns[-1].content)
        if forge_status:
            candidate["status"] = "available"
        return ModelResponse([
            ToolCallPart(info.output_tools[0].name, candidate)
        ])
    return FunctionModel(respond)

def checked_status(context: Context, requested_id: str, reply: Reply) -> str:
    observed = context.evidence.get(requested_id)
    if reply.lab_id != requested_id or reply.status != observed:
        return "review"
    return observed

def scenario(lab_id: str, expected: str, forge_status: bool = False):
    context = Context(user_id="learner-1")
    result = build_agent(fixture(lab_id, forge_status)).run_sync(
        f"Please check {lab_id}.",
        deps=context,
        usage_limits=UsageLimits(request_limit=3),
    )
    status = checked_status(context, lab_id, result.output)
    assert status == expected
    print(lab_id, status)

if __name__ == "__main__":
    scenario("LAB-101", "available")
    scenario("LAB-202", "review")
    scenario("LAB-999", "review")
    scenario("LAB-202", "review", forge_status=True)

Store history by authenticated user and conversation ID. Check ownership when loading it, apply a retention policy, and keep a new conversation empty unless the user intentionally continues a previous one. Never keep all users’ messages in one mutable global list.