Datalumina

Agent and grounding

The research agent

Turn retrieval into a Pydantic AI tool loop that searches, reads, and returns a structured, citable answer.

This stage puts an LLM in charge of the search layer you just built. When it is done, an analyst question produces a GroundedAnswer. Plain text with [n] citation markers, a citations list with verbatim excerpts, and an honest flag when the corpus cannot support an answer. The agent decides for itself when to search, what to read in full, and when it has enough evidence.

daveebbelaar/document-copilot2:04:00

Tool loop, not context stuffing

Most RAG tutorials are one-shot. Embed the query, fetch top-k chunks, paste them into a system prompt, and tell the model "here's everything, go answer." Dave calls this out directly in the video, saying "there's no way to have a fallback or to course correct, because in case it did not find enough evidence, the default loop is not implemented."

The alternative is an agentic loop. The model gets tools instead of pre-stuffed context. In Dave's words, "with Pydantic AI out of the box, when you have an agent and you give it tools, the agent can use those tools, check what kind of answer it gets back, and then use tools again." If the first search misses, it reformulates. If an excerpt cuts off mid-table, it reads the surrounding chunks. Hybrid search and agentic search combined, which is why answers take a minute but come back grounded.

Every passage a tool returns also lands in the turn registry on the way through. That detail looks minor here and becomes the entire trust story two sections down.

Define the agent

The agent lives in backend/app/assistant/agent.py. It is fully typed. Dependencies in, structured output out.

# backend/app/assistant/agent.py
def get_document_agent() -> Agent[DocumentAgentDeps, GroundedAnswer]:
    global _document_agent
    if _document_agent is None:
        model = OpenAIChatModel(
            settings.openai_chat_model,
            provider=OpenAIProvider(api_key=settings.openai_api_key),
        )
        _document_agent = Agent(
            model,
            deps_type=DocumentAgentDeps,
            output_type=GroundedAnswer,
            instructions=INSTRUCTIONS,
            tools=[search_filings, read_chunks, read_chunk, read_surrounding_chunks],
        )
    return _document_agent

INSTRUCTIONS is loaded from instructions.md next to the module (more on that below). The run itself is capped so a confused model cannot loop forever. run_document_agent passes UsageLimits(request_limit=settings.openai_agent_request_limit), which defaults to 20 requests in app/config.py, alongside openai_agent_temperature: float = 0.0 for repeatable behavior.

The four tools

backend/app/assistant/tools.py wraps the retrieval layer in four bounded tools:

ToolPurpose
search_filingsHybrid retrieval with optional ticker, form, fiscal_years filters
read_chunksBatch-fetch full text for multiple chunk UUIDs in one call
read_chunkFetch a single chunk by UUID
read_surrounding_chunksRead chunks before and after an anchor chunk in the same filing

The search tool shows the pattern all four follow:

# backend/app/assistant/tools.py
async def search_filings(
    ctx: RunContext[DocumentAgentDeps],
    query: str,
    ticker: str | None = None,
    form: str | None = None,
    fiscal_years: str | None = None,
) -> str:
    """Search SEC filings with hybrid retrieval. Optional filters: ticker, form, fiscal_years (comma-separated)."""
    ...
    passages = await _run_tool(
        ctx.deps, "search_filings", detail, _search_sync,
        ctx.deps, query, ticker=ticker, form=form, fiscal_years=fiscal_years,
    )
    ctx.deps.registry.register_many(passages)
    return format_passages_for_agent(passages)

Three implementation details carry the design:

  • Blocking work goes to a thread. The retriever and SQLAlchemy sessions are synchronous, so _run_tool runs them via asyncio.to_thread. That keeps the event loop free, and Pydantic AI can execute multiple tool calls from one model response concurrently.
  • Every tool emits status events. _run_tool calls emit_tool_start before the work and reports duration and result counts after, which later powers the "searching filings..." updates in the chat UI.
  • Output is bounded. format_passages_for_agent caps each excerpt at MAX_PASSAGE_EXCERPT_CHARS (800) and the whole tool response at MAX_AGENT_OUTPUT_CHARS (12,000), so a broad search cannot blow up the context window. The agent reads excerpts first and only fetches full text when it needs it.

The turn registry

The most important object in this phase is the smallest. TurnRegistry in backend/app/assistant/deps.py records every passage any tool returns during a turn:

# backend/app/assistant/deps.py
@dataclass
class TurnRegistry:
    passages_by_chunk_id: dict[UUID, RetrievedPassage] = field(default_factory=dict)

    def register(self, passage: RetrievedPassage) -> None:
        self.passages_by_chunk_id[passage.chunk_id] = passage
        for neighbor in passage.neighbors:
            self.passages_by_chunk_id[neighbor.chunk_id] = neighbor

    def register_many(self, passages: list[RetrievedPassage]) -> None:
        for passage in passages:
            self.register(passage)

Notice that every tool ends with ctx.deps.registry.register(...) before returning text to the model. That makes the registry a citation allowlist. Only chunk IDs that were actually retrieved this turn may appear in the final answer. The next stage, grounding validation, enforces exactly that. Without the registry, a model could invent a plausible-looking UUID and nobody could tell.

Structured output

The agent cannot return free text. Its output type in backend/app/assistant/outputs.py forces the answer and its evidence into one object:

# backend/app/assistant/outputs.py
class Citation(BaseModel):
    citation_index: int = Field(description="1-based index referenced as [n] in the answer text")
    chunk_id: UUID = Field(description="UUID of the cited document chunk")
    excerpt: str = Field(description="Verbatim substring from the chunk text supporting the claim")

class GroundedAnswer(BaseModel):
    answer: str = Field(description="Plain-English answer with [n] citation markers")
    citations: list[Citation] = Field(default_factory=list)
    insufficient_evidence: bool = Field(default=False)

insufficient_evidence is the escape hatch that makes honesty structurally possible. Instead of forcing an answer, the model can say the corpus does not contain one.

Instructions as the product contract

backend/app/assistant/instructions.md is where the client brief becomes enforceable behavior. The core rules:

- Answer **only** from passages returned by your tools. Never invent facts, numbers, or filing language.
- **Cite every factual claim** with `[n]` markers in the answer text that match `citation_index` in your citations list.
- Each citation must include a **verbatim excerpt** copied from the retrieved chunk text.
- If the corpus does not contain enough evidence, set `insufficient_evidence` to true, explain what is missing, and return an **empty** citations list.
- **No stock picks**, trading recommendations, or investment advice.
- Do not infer causation or conclusions beyond what the filings explicitly state (e.g. do not claim generative AI improved margins unless a filing directly says so).

That last rule targets a specific trap question from the client brief, "did generative AI improve margins?" Filings mention AI and filings mention margins, and a helpful model will happily connect them. The contract says it may not. The file also steers tool efficiency. Answer from search excerpts when possible, batch reads through read_chunks in one call, and use read_surrounding_chunks only when neighbors are not enough.

Watch the size of AI-generated phases

Dave's honest note on this phase is that it "one-shotted" far more than he asked for, and "that phase number six, that was definitely way too big." His recovery move is worth copying. He had the agent generate README files for the new modules, then read those docs to reverse-engineer what was actually built. When a coding agent overshoots, make it explain itself before you trust the code.

Test it interactively

backend/scripts/smoke_assistant.py runs one query end to end with timestamped progress logging. Pick a query by editing QUERY_KEY at the top (apple-mix, nvda-datacenter, q10-refusal, or underspecified), then run from backend/:

uv run python -u scripts/smoke_assistant.py

The script wires up real deps and prints every tool call as it happens:

# backend/scripts/smoke_assistant.py
registry = TurnRegistry()
deps = DocumentAgentDeps(
    retriever=DocumentRetriever(),
    registry=registry,
    thread_id=uuid.uuid4(),
    user_id=uuid.uuid4(),
)
answer = prune_unreferenced_citations(run_document_agent(query, deps))
validation = asyncio.run(GroundingValidator().validate(answer, registry))

Two practical notes from the video. The script calls nest_asyncio.apply() because Pydantic AI runs an event loop under the hood, and running it inside a Jupyter or interactive session fails without it. And the progress logging exists because Dave's first run looked frozen. As he put it, "we need to check, is it actually stuck or is it actually taking that long?" With timestamps you can see it is thinking, not hanging.

Checkpoint

  • uv run python -u scripts/smoke_assistant.py streams progress lines (agent start, each tool call with duration, agent done with token usage)
  • The apple-mix query returns a GroundedAnswer with [n] markers and matching citations
  • Every cited chunk_id exists in the printed registry metadata (ticker, form, page)
  • The run stays under the request limit and finishes in roughly a minute

Next: Grounding validation

Verify every citation against the turn registry before the user ever sees the answer.

On this page