Build a Pydantic AI regeneration agent with typed tools

Introduction

When a user thumbs-down a generated response, the simplest recovery is to re-run the same prompt and hope for better output—but the same model, temperature, and system prompt will produce statistically similar text, and the silent retry teaches users that feedback does nothing. Teams that ship LLM products quickly find that meaningful regeneration requires reasoning about why the response failed and adjusting generation parameters accordingly. By the end of this lesson you'll be able to build a Pydantic AI agent that reads a structured feedback signal, selects typed corrective tools (temperature delta, model switch, system-prompt append), and dispatches the regeneration through LiteLLM—producing auditable adjustments instead of indistinguishable retries.

Key Terminology

• FeedbackSignal: an enum that classifies a negative rating's reason (e.g. TOO_GENERIC, WRONG_TONE, TOO_SHORT) and drives which corrective tool the agent selects.
• RegenerationDeps: the Pydantic AI dependency container that injects per-request state—database session, conversation history, feedback signal, mutable current_params—into every tool call.
• Typed tool: a function registered on the agent with @agent.tool whose parameters are validated by Pydantic; the regeneration agent exposes three (adjust_temperature, switch_model, modify_system_prompt) as the only legal corrective actions.

Concepts

Why an Agent Instead of a Function

A conventional approach would encode regeneration logic in a chain of if/elif blocks: if the rating is negative and the comment mentions "too short," raise max_tokens; if the comment mentions "wrong tone," swap the system prompt. This works for three or four rules but collapses under combinatorial pressure. What happens when a message is simultaneously "too short" and "wrong tone" and the user wants a different model? You need a reasoning layer that can compose multiple adjustments in a single pass, which is exactly what a Pydantic AI agent provides. The agent receives a natural-language planning prompt enriched with structured feedback data, and it calls one or more typed tools—adjust_temperature, switch_model, modify_system_prompt—in whatever combination the situation requires. Each tool call is validated by Pydantic models, logged for analytics, and reversible for A/B comparison workflows described later in this chapter.

Code Walkthrough

Building on the agent-versus-function reasoning above, the next step is to wire that reasoning layer into code: a dependency container that carries feedback context into every tool call, the three typed tools the agent composes, and the end-to-end flow that turns a thumbs-down into a parameter-adjusted regeneration.

The Dependency Injection Context

Before the agent can make decisions, it needs access to the feedback record, the original message, and the set of available models. Pydantic AI's deps_type mechanism injects this context into every tool call without global state or closures. The following code defines the RegenerationDeps dataclass that encapsulates all request-scoped data the agent needs. The feedback_signal field carries the structured feedback collected by the PostgreSQL-backed storage models described in another goal—a FeedbackSignal enum value—alongside the free-text feedback_comment. The available_models list is populated from the application's configuration and determines which providers the switch_model tool can select. The current_params dictionary holds the mutable generation parameters that tools will modify in place before the final LiteLLM call.

1from __future__ import annotations
2
3import enum
4from dataclasses import dataclass, field
5from sqlalchemy.ext.asyncio import AsyncSession
6
7class FeedbackSignal(enum.Enum):
8    TOO_GENERIC = "too_generic"
9    FACTUALLY_WRONG = "factually_wrong"
10    WRONG_TONE = "wrong_tone"
11    TOO_LONG = "too_long"
12    TOO_SHORT = "too_short"
13    OTHER = "other"
14
15@dataclass
16class RegenerationDeps:
17    db_session: AsyncSession
18    message_id: str
19    conversation_history: list[dict[str, str]]
20    feedback_signal: FeedbackSignal
21    feedback_comment: str
22    available_models: list[str] = field(
23        default_factory=lambda: [
24            "gpt-4o",
25            "gemini/gemini-2.5-flash",
26            "anthropic/claude-sonnet-4-20250514",
27        ]
28    )
29    current_params: dict = field(
30        default_factory=lambda: {
31            "temperature": 0.7,
32            "model": "gpt-4o",
33            "system_prompt": "",
34            "max_tokens": 2048,
35        }
36    )
37    applied_adjustments: list[str] = field(default_factory=list)

Registering Typed Tools on the Agent

With the dependency context defined, the next step is building the agent and registering the three corrective tools. Each tool receives a RunContext[RegenerationDeps] as its first parameter, giving it typed access to the mutable current_params dictionary. The adjust_temperature tool clamps the new value between 0.0 and 2.0 to prevent invalid API calls. The switch_model tool validates that the requested model exists in available_models and raises a ValueError if it does not—Pydantic AI surfaces this as a tool error the agent can recover from. The modify_system_prompt tool appends behavioral instructions rather than replacing the entire prompt, preserving the original context while steering the regeneration. All three tools return a confirmation string that the agent uses to plan subsequent actions.

1from pydantic_ai import Agent, RunContext
2
3regeneration_agent = Agent(
4    "gpt-4o",
5    deps_type=RegenerationDeps,
6    system_prompt=(
7        "You are a regeneration controller. Analyze the feedback signal "
8        "and comment, then call one or more tools to adjust generation "
9        "parameters before the message is regenerated. Always explain "
10        "your reasoning before calling tools."
11    ),
12)
13
14@regeneration_agent.tool
15async def adjust_temperature(
16    ctx: RunContext[RegenerationDeps], delta: float
17) -> str:
18    current = ctx.deps.current_params["temperature"]
19    new_temp = max(0.0, min(2.0, current + delta))
20    ctx.deps.current_params["temperature"] = new_temp
21    adjustment = f"Temperature: {current:.2f} -> {new_temp:.2f}"
22    ctx.deps.applied_adjustments.append(adjustment)
23    return adjustment
24
25@regeneration_agent.tool
26async def switch_model(
27    ctx: RunContext[RegenerationDeps], model_name: str
28) -> str:
29    if model_name not in ctx.deps.available_models:
30        raise ValueError(
31            f"{model_name} not in {ctx.deps.available_models}"
32        )
33    old = ctx.deps.current_params["model"]
34    ctx.deps.current_params["model"] = model_name
35    adjustment = f"Model: {old} -> {model_name}"
36    ctx.deps.applied_adjustments.append(adjustment)
37    return adjustment
38
39@regeneration_agent.tool
40async def modify_system_prompt(
41    ctx: RunContext[RegenerationDeps], instruction: str
42) -> str:
43    existing = ctx.deps.current_params["system_prompt"]
44    separator = "\n\n" if existing else ""
45    ctx.deps.current_params["system_prompt"] = (
46        f"{existing}{separator}{instruction}"
47    )
48    adjustment = f"System prompt appended: {instruction[:80]}"
49    ctx.deps.applied_adjustments.append(adjustment)
50    return adjustment

Agent Execution and Regeneration Flow

The following diagram shows how a feedback event flows through the agent into a regenerated response. The agent sits between the feedback ingestion layer from another goal and the LiteLLM dispatch layer used by the concurrent A/B comparison in another goal. This positioning allows the same regeneration logic to feed into both single-model regeneration and multi-model comparison workflows.

Loading diagram...

The sequence makes three architectural decisions explicit. First, the agent does not call LiteLLM itself—it only adjusts parameters. This separation of concerns means the same current_params dictionary can be passed to the single-model litellm.acompletion call or fanned out to the concurrent multi-model comparison endpoint from another goal. Second, the adjustment metadata is persisted alongside the new message, enabling the materialized views from another goal to compute regeneration success rates partitioned by adjustment type. Third, the conversation history loaded from PostgreSQL follows the branching model from another goal: if the user edited a message before requesting regeneration, the history reflects the pruned branch, not the stale original.

Wiring the Agent into the Endpoint

To plug the agent into a FastAPI route, load the original Message and its Feedback row from PostgreSQL, classify the comment to a FeedbackSignal (keyword matching is enough to demonstrate the typed-tool architecture; an embedding-based classifier slots in later), reconstruct the active branch's conversation history, and build a RegenerationDeps seeded with the original message's system_prompt and model_used. Calling await regeneration_agent.run(...) with a prompt that names the signal, the user comment, and the current model lets the agent mutate deps.current_params through its typed tools. Then await litellm.acompletion(**deps.current_params, messages=[...]) produces the regenerated text; persist the new Message with regeneration_of=message_id and adjustments_applied=deps.applied_adjustments, and return both the new content and the adjustment list so the UI can show "Regenerated with lower temperature and modified instructions" instead of an opaque retry.

You'll know it works when a thumbs-down on a too-generic answer triggers a regeneration whose adjustments payload includes at least one tool invocation (e.g. Temperature: 0.70 -> 0.40) and the new message body differs meaningfully from the original.

Do's and Don'ts

Do's

1. ✓Do declare RegenerationDeps as the deps_type on the agent and accept RunContext[RegenerationDeps] as the first parameter of every tool — this gives each tool typed, request-scoped access to current_params and applied_adjustments without closures or global state, so parameter mutations stay isolated to a single regeneration request and remain fully auditable.
2. ✓Do clamp the result inside adjust_temperature with max(0.0, min(2.0, current + delta)) before writing back to current_params["temperature"] — the agent supplies an unconstrained float delta, and a value outside [0.0, 2.0] causes an invalid LiteLLM API call that surfaces as an unrecoverable dispatch error rather than a typed tool error the agent can reason about and correct.
3. ✓Do append behavioral instructions in modify_system_prompt rather than replacing current_params["system_prompt"] — overwriting the original system prompt discards the context the first response was built on; appending a targeted instruction (e.g., a tone or factual-accuracy directive) steers the regeneration without invalidating the original framing.

Don'ts

1. ✗Don't skip the available_models guard in switch_model and write an arbitrary string directly into current_params["model"] — without validating against the known provider list, an unrecognized or malformed model name (e.g., missing the gemini/ prefix) silently reaches the LiteLLM dispatch call and raises an unrecoverable provider error instead of a typed ValueError the agent can catch and retry with a valid alternative.
2. ✗Don't regenerate a thumbs-down response by re-dispatching the original current_params without any tool invocation — the same temperature, model, and system prompt produce statistically similar output, and a silent retry with no parameter change teaches users that submitting feedback does nothing, undermining the entire feedback-driven regeneration loop.
3. ✗Don't mutate current_params directly in application code outside the registered tool functions — changes made outside adjust_temperature, switch_model, or modify_system_prompt bypass the applied_adjustments.append(...) call in each tool, breaking the audit trail that records which parameters were adjusted and by how much for downstream regeneration success tracking.