Add an Agent¶

This guide walks through adding a new AI agent to the pipeline, following the established agent pattern.

1. Define the Output Schema¶

Add a Pydantic model to src/agents/schemas.py:

class MyAgentOutput(BaseModel):
    model_config = ConfigDict(strict=True)

    passed: bool
    score: float
    feedback: str

2. Create the Agent Module¶

Create src/agents/my_agent.py:

import json

from src.agents.base import call_agent
from src.agents.schemas import MyAgentOutput

async def run_my_agent(
    draft: str,
    *,
    article_id: uuid.UUID | None = None,
    session: AsyncSession | None = None,
) -> MyAgentOutput:
    messages = [
        {"role": "user", "content": f"<draft>{draft}</draft>"}
    ]

    response = await call_agent(
        "my_agent",
        "article",
        messages,
        article_id=article_id,
        session=session,
    )

    return MyAgentOutput.model_validate_json(
        extract_json(response.content)
    )

If your agent needs tools, add tools=_TOOL_MAP and tool_definitions=_TOOL_DEFINITIONS (see the factuality checker for an example).

3. Seed the Prompt¶

Create an Alembic data migration:

uv run alembic revision -m "seed my_agent prompt v1"

In the migration, insert the system prompt:

def upgrade() -> None:
    op.execute(
        sa.text(
            "INSERT INTO agent_prompts "
            "(agent_name, version, model, system_prompt, active) "
            "VALUES (:name, :version, :model, :prompt, true)"
        ).bindparams(
            name="my_agent",
            version=1,
            model="claude-haiku-4-5",
            prompt="You are a specialist agent that...",
        )
    )

Apply with uv run alembic upgrade head.

4. Add Tools (Optional)¶

If your agent needs external data access, create or reuse tool modules from src/tools/:

Export a *_SCHEMA constant (JSON dict for Anthropic's tool protocol)
Export an async handler function
In your agent module, wrap the handler to return JSON strings
Build _TOOL_MAP and _TOOL_DEFINITIONS

5. Write Tests¶

Add tests in tests/unit/test_agents/test_my_agent.py:

Mock call_agent at the import site: @patch("src.agents.my_agent.call_agent")
Test that the agent assembles the correct messages
Test that the output is parsed correctly
Test error handling (invalid JSON, missing fields)

6. Integrate into a Flow¶

Call your agent from the appropriate Prefect flow in src/pipeline/. Remember: no business logic in flows — the flow just calls the agent and persists results.

Checklist¶

[ ] Output schema in schemas.py
[ ] Agent module in src/agents/
[ ] Prompt seeded via Alembic migration
[ ] Unit tests with mocked call_agent
[ ] Integrated into pipeline flow
[ ] Type check passes: uv run mypy src/