Build an agentic RAG agent with Pydantic AI

Introduction

When you build a RAG pipeline that retrieves once and answers from whatever it got, ambiguous queries and topic-spanning questions hand the model low-relevance chunks — and you ship answers that confidently miss the point. Agentic RAG turns retrieval into a controlled loop the LLM itself drives: score the context, reformulate when it's weak, stop when it's good enough or the retry budget runs out. By the end of this lesson you'll be able to wire up a Pydantic AI agent that calls a hybrid pgvector + BM25 retriever, gates output on a relevance threshold, and iterates until the threshold passes or max_iterations is hit.

Key Terminology

• pgvector — a PostgreSQL extension that adds a vector column type and approximate-nearest-neighbour indexes, so the agent can run cosine-similarity searches over embeddings without leaving the database.
• BM25 — a lexical scoring function that ranks documents by token-frequency match against the query; supplies the keyword half of hybrid retrieval that complements pgvector's semantic half.
• Reciprocal Rank Fusion (RRF) — a rank-aggregation method that merges two ranked lists into one by summing 1 / (k + rank) per document; how this agent blends semantic and lexical signals into the single relevance score it gates on.
• Agentic RAG — a retrieval pattern where the LLM decides per call whether the retrieved context is good enough and reformulates the query if not, instead of running a single fixed retrieve-then-generate step.

Concepts

Query Reformulation Strategy

The quality of agentic RAG depends heavily on how effectively the agent reformulates failed queries. The system prompt instructs the agent to "target missing information," but the actual reformulation strategy emerges from the LLM's reasoning. In practice, effective reformulations follow predictable patterns. When the initial query is too broad (e.g., "How does authentication work?"), the agent narrows to specific mechanisms ("JWT token validation in the middleware layer"). When retrieval returns tangentially related chunks, the agent adds discriminating terms or switches from abstract concepts to concrete implementation details. When chunks from different domains contaminate results, the agent adds exclusion context ("OAuth2 flow for the Python backend, not the JavaScript client").

You can influence reformulation quality by enriching the tool's return value. Including the actual text snippets in top_contexts lets the agent see what was retrieved and reason about why it scored poorly. Adding metadata like document titles or section headers gives the agent additional signals for reformulating. Some teams include a gap_analysis field where a secondary LLM call identifies what information the chunks are missing relative to the query—though this doubles latency per iteration and is typically reserved for complex multi-hop questions.

The max_iterations parameter of 3 balances thoroughness against latency. Each iteration adds roughly 1–2 seconds of retrieval time plus one LLM reasoning step. For user-facing applications, keeping the total under 3 iterations ensures sub-10-second response times. For batch processing or research applications, increasing to 5 iterations can improve answer quality on complex questions at the cost of higher token consumption and latency. Monitor the iterations field in the response to tune this threshold—if most queries resolve in 1 iteration, your base retrieval is strong and agentic overhead is minimal. If most queries require 3 iterations, investigate whether the chunking strategy from another goal or the hybrid weights from another goal need adjustment before adding more agent iterations.

Code Walkthrough

Agentic RAG Decision Loop

The following diagram illustrates how the Pydantic AI agent orchestrates iterative retrieval. Each cycle through the loop represents one tool call where the agent invokes retrieve_and_evaluate, inspects the relevance score, and decides whether to reformulate or finalize.

A Pydantic AI Agent orchestrates this hybrid RAG pipeline by routing each user query through retrieve_and_evaluate(), which invokes the HybridRetriever to combine pgvector semantic search with BM25 lexical matching. Reciprocal Rank Fusion merges both result sets, and a relevance scorer gates output at a 0.7 threshold. When context falls short, the agent reformulates the query up to max_iterations times before falling back to a best-effort answer with an explicit caveat and source attribution.

Loading diagram...

This loop is not hardcoded in application logic—the agent's system prompt instructs it to evaluate scores and reformulate autonomously. The LLM sees the relevance score returned by the tool and reasons about whether the context is adequate, what information is missing, and how to adjust the query. This is fundamentally different from a programmatic retry loop because the agent brings semantic understanding to the reformulation step.

Defining the Agent and Its Retrieval Tool

Before building the agent itself you need a dependency container that carries shared resources—the database session for hybrid retrieval, the embedding model for semantic search, and configuration parameters like the relevance threshold and maximum iteration count. Pydantic AI's RunContext injects this dependency object into every tool call, eliminating the need for global state or closure-captured variables. The same module also defines the retrieve_and_evaluate tool the agent will call: it runs the HybridRetriever (pgvector cosine similarity + ts_vector BM25 full-text), fuses the result sets with reciprocal rank fusion, scores the aggregate relevance, and returns a structured RetrievalResult the LLM can reason about. Accumulating contexts across iterations on the deps object lets the final answer cite chunks from every pass, not just the last one.

Code snippetpython
1import numpy as np
2from dataclasses import dataclass, field
3from sqlalchemy.ext.asyncio import AsyncSession
4from pydantic import BaseModel
5from pydantic_ai import Agent, RunContext
6
7from app.retrieval.hybrid import HybridRetriever
8
9@dataclass
10class AgenticRAGDeps:
11    """Dependency container injected into every agent tool call."""
12    db_session: AsyncSession
13    retriever: HybridRetriever
14    relevance_threshold: float = 0.7
15    max_iterations: int = 3
16    current_iteration: int = field(default=0, init=False)
17    accumulated_contexts: list[dict] = field(default_factory=list, init=False)
18
19class RetrievalResult(BaseModel):
20    """Structured result returned to the agent from retrieval."""
21    query_used: str
22    relevance_score: float
23    num_chunks: int
24    top_contexts: list[dict]
25    retries_remaining: int
26
27SYSTEM_PROMPT = """You are a research assistant that answers questions using
28a hybrid retrieval system. For each question:
29
301. Call retrieve_and_evaluate with the user's query.
312. Examine the relevance_score in the result.
323. If relevance_score >= 0.7, synthesize a final answer from the contexts.
334. If relevance_score < 0.7 and retries remain, reformulate the query to
34   target missing information, then call retrieve_and_evaluate again.
355. After max retries, generate a best-effort answer noting low confidence.
36
37Always cite source document IDs in your answer."""
38
39rag_agent = Agent(
40    "openai:gpt-4o",
41    system_prompt=SYSTEM_PROMPT,
42    deps_type=AgenticRAGDeps,
43    retries=2,
44)
45
46@rag_agent.tool
47async def retrieve_and_evaluate(
48    ctx: RunContext[AgenticRAGDeps],
49    query: str,
50) -> RetrievalResult:
51    """Retrieve documents via hybrid search and score relevance."""
52    deps = ctx.deps
53    deps.current_iteration += 1
54
55    # Execute hybrid retrieval: pgvector cosine + BM25 ts_vector
56    chunks = await deps.retriever.search(
57        query=query,
58        session=deps.db_session,
59        top_k=10,
60        semantic_weight=0.6,
61        bm25_weight=0.4,
62    )
63
64    # Score context relevance as mean RRF score of top chunks,
65    # normalized to 0-1 against the empirical RRF distribution
66    similarities = [c["rrf_score"] for c in chunks if "rrf_score" in c]
67    relevance_score = float(np.mean(similarities)) if similarities else 0.0
68    relevance_score = min(relevance_score / 0.05, 1.0)
69
70    top_contexts = [
71        {
72            "doc_id": c["doc_id"],
73            "text": c["text"][:500],
74            "score": round(c["rrf_score"], 4),
75        }
76        for c in chunks[:5]
77    ]
78
79    # Accumulate across iterations for final synthesis
80    deps.accumulated_contexts.extend(top_contexts)
81
82    retries_remaining = deps.max_iterations - deps.current_iteration
83    return RetrievalResult(
84        query_used=query,
85        relevance_score=round(relevance_score, 3),
86        num_chunks=len(chunks),
87        top_contexts=top_contexts,
88        retries_remaining=max(retries_remaining, 0),
89    )

• AgenticRAGDeps: The relevance_threshold of 0.7 and max_iterations of 3 are configurable at construction time, while current_iteration and accumulated_contexts are mutable runtime state tracked across tool calls within a single agent run.
• RetrievalResult: Pydantic AI serializes this BaseModel to JSON before passing it back to the LLM, so the agent sees structured fields like relevance_score and retries_remaining that it can reason about explicitly. The top_contexts field carries the actual text chunks the agent will use for synthesis.
• Agent instantiation: The deps_type parameter tells the framework to inject AgenticRAGDeps into every tool's RunContext. The retries=2 parameter controls Pydantic AI's internal retry on validation errors—separate from the agent's semantic retry loop.
• retrieve_and_evaluate: The @rag_agent.tool decorator registers this async function as a callable tool whose JSON schema Pydantic AI auto-generates from the signature and docstring. The aggregate RRF score is normalized into 0-1 before being returned so the agent's threshold check is stable across query distributions.

Running the Agent and Handling the Response

With the agent and tool defined, running the agentic RAG loop requires constructing the dependencies, calling agent.run, and extracting both the final answer and the full message history for observability. The message history is critical for production systems because it captures every tool call, every relevance score, and every query reformulation the agent performed—data you need for debugging retrieval quality issues and tuning the relevance threshold. The following snippet demonstrates the complete invocation pattern, including how to extract iteration metadata from the agent's run result for logging to your OpenTelemetry tracing pipeline built in another goal.

1from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
2
3from app.retrieval.hybrid import HybridRetriever
4
5async def agentic_rag_query(user_question: str) -> dict:
6    """Execute agentic RAG with iterative retrieval."""
7    engine = create_async_engine(
8        "postgresql+asyncpg://user:pass@localhost:5432/ragdb"
9    )
10    async_session = async_sessionmaker(engine, class_=AsyncSession)
11
12    async with async_session() as session:
13        deps = AgenticRAGDeps(
14            db_session=session,
15            retriever=HybridRetriever(embedding_model="text-embedding-3-small"),
16            relevance_threshold=0.7,
17            max_iterations=3,
18        )
19
20        result = await rag_agent.run(user_question, deps=deps)
21
22        # Extract iteration metadata for observability
23        tool_calls = [
24            msg for msg in result.all_messages()
25            if hasattr(msg, "tool_name")
26        ]
27
28        return {
29            "answer": result.data,
30            "iterations": deps.current_iteration,
31            "total_contexts": len(deps.accumulated_contexts),
32            "tool_calls": len(tool_calls),
33            "sources": list({
34                ctx["doc_id"] for ctx in deps.accumulated_contexts
35            }),
36        }

Do's and Don'ts

Do's

1. ✓Do inject shared retrieval state via AgenticRAGDeps and Pydantic AI's RunContext — passing the AsyncSession, HybridRetriever, relevance_threshold, and max_iterations through the dependency container keeps every retrieve_and_evaluate tool call self-contained and eliminates the global-state and closure-captured-variable patterns that break async concurrency.
2. ✓Do accumulate retrieved chunks in accumulated_contexts across every iteration, not just the final one — because the agent may find the highest-relevance evidence in an earlier pass before reformulating, discarding prior-iteration chunks means the final answer loses citable sources from those passes and narrows its grounding unnecessarily.
3. ✓Do let the system prompt drive query reformulation by exposing relevance_score and retries_remaining in the RetrievalResult — when the LLM reads the score and reasons about what information is missing, it brings semantic understanding to rewriting that a hardcoded programmatic retry loop cannot; the reformulation targets the specific gap rather than mechanically resubmitting the original query.

Don'ts

1. ✗Don't gate output on the raw RRF score without normalizing it to the 0–1 range first — raw Reciprocal Rank Fusion values cluster well below 1.0 (the walkthrough divides by 0.05 then clamps), so comparing an unnormalized score against the 0.7 threshold will cause the condition to fail on virtually every query and exhaust max_iterations regardless of retrieval quality.
2. ✗Don't replace the agent-driven reformulation loop with an external programmatic retry — wrapping retrieve_and_evaluate in your own for loop and regenerating queries without LLM reasoning removes the semantic reformulation step; you get mechanical repetition rather than targeted query rewriting, which is the core failure mode agentic RAG is designed to fix.
3. ✗Don't return a low-confidence answer silently when max_iterations is exhausted — omitting the explicit caveat (as specified in the system prompt's step 5) lets the model present an inadequately grounded response with the same confidence as a threshold-passing answer, which reintroduces the confident-but-wrong failure mode the entire iterative loop exists to prevent.