ADK RAG grounding — retrieval, reranking, citations, hallucination defense

The boundary — memory recalls the user, RAG grounds on the corpus

ADK applications end up with two retrieval-shaped subsystems that look identical from a distance: both embed text, both run top-k similarity, both splice results into the prompt. The distinction that matters is the write path. A memory bank is written by conversations — claims extracted from finished sessions, scoped to one user, authored implicitly by whatever that user happened to say. A RAG corpus is written by an ingestion pipeline — documents your organisation authored deliberately, scoped by role or team, versioned somewhere outside the agent.

The failure modes diverge accordingly. A bad memory means the agent misremembers a preference: awkward, self-correcting, bounded to one user. A bad grounding means the agent states something false about your product and attaches a citation to it, which is the failure that ends up in a support escalation or a compliance review. Keep the two in separate stores behind separate tools. The companion article on the memory bank covers extraction, consolidation and decay — none of which apply here, because documents are not consolidated, they are re-ingested. The rule of thumb: if the claim’s source is something the user told you, it is memory; if the source is a document you publish, it is RAG, and it needs a citation.

Retrieval as a tool, or retrieval built into the turn

There are two places retrieval can attach, and the choice shapes latency, cost, and how the agent fails. Built into the turn means a pre-model step — a before_model_callback or an equivalent pre-processor — runs retrieval on every user message and injects the passages before the request goes out. It is deterministic: the model can never forget to ground, because it never had the choice. It also pays a retrieval hop on turns that did not need one (‘thanks, that worked’), and it can only ever retrieve once, on the raw user text.

Retrieval as a tool hands the decision to the model. It calls when it judges it needs the corpus, can rewrite a vague question into a better query first, and can retrieve twice when the first result set was thin. The cost is an extra model round trip per call, and a real failure mode: the model skips the tool and answers from parameters. The general-purpose tool loop — declarations, dispatch, ToolContext — is covered in the tools article; what is specific to grounding is the return shape. Every chunk must carry a stable id and a source reference, because that is what citations are later built from.

def search_docs(query: str, tool_context: ToolContext) -> dict:
    """Search the product documentation corpus.

    Args:
        query: A focused question. Prefer specific terms over the
            user's raw phrasing.
    """
    hits = retriever.search(query, top_k=8)
    # Record what was retrieved so a later check can score the answer
    # against exactly these passages, not against the whole corpus.
    tool_context.state["temp:retrieved"] = [h.chunk_id for h in hits]
    return {
        "chunks": [
            {"id": h.chunk_id, "text": h.text,
             "source": h.title, "locator": h.section, "score": h.score}
            for h in hits
        ]
    }

Hybrid designs are often right: retrieve on the first turn to seed context, then expose the tool for follow-ups.

Chunking is the decision everything downstream inherits

Chunking is the least glamorous stage and the one that bounds the ceiling of the whole system. A passage that answers the question perfectly is worthless if it was split across a boundary, because neither half retrieves well and neither half is self-contained enough to ground on. This is the single most common root cause behind ‘the model hallucinated’ reports that turn out, on inspection, to be retrieval failures.

Two adjustments do most of the work. Overlap — carry the last sentence or two of the previous chunk into the next — cheaply insures against boundary splits at the cost of some index bloat. And contextual headers: prepend the document title and heading path to each chunk’s indexed text, so a chunk saying ‘this is limited to 30 days’ still embeds near ‘refund’. Size itself is a genuine trade: small chunks give sharp relevance and thin context, large chunks give rich context and diluted embeddings.

Indexing: vector, keyword, and why you want both

Vector search embeds the query and returns the chunks nearest in embedding space. It is excellent at paraphrase — ‘can I send it back if it arrived broken’ finds a passage titled ‘damaged goods returns’ with no shared vocabulary. It is unreliable at exactly the tokens your users care about most: part numbers, error codes, API names, version strings. Embeddings compress, and identifiers are the first thing compression discards — ask a purely semantic index about ERR_4021 and you get chunks about errors in general.

Keyword search (BM25 and friends) has the mirror-image profile: exact on identifiers, blind to paraphrase. Hybrid retrieval runs both and fuses the results — reciprocal rank fusion is the usual default because it needs no score calibration between two systems whose scores are not comparable. For most document corpora hybrid is not an optimisation, it is the baseline; pure-vector setups tend to look fine in demos and fall over on the first real support ticket containing a SKU. Alongside this, index metadata you can filter on before ranking: product, version, locale, effective date, audience. A filter that removes two-thirds of the corpus improves relevance more than any reranker, and costs nothing.

Reranking: relevant versus merely similar

Retrieval optimises for recall at speed. It compares a query vector against millions of chunk vectors, which forces a cheap comparison — the query and the chunk are embedded independently, so nothing in the scoring ever looks at them together. That is why the top-k contains chunks that are topically adjacent but do not answer the question.

A reranker fixes this by being expensive on purpose. A cross-encoder takes the query and one candidate chunk jointly and scores actual relevance. You cannot run it over a corpus; you can easily run it over 30 candidates. The standard shape is therefore retrieve wide, rerank narrow, ground on few: fetch 30–50 candidates from hybrid search, rerank, keep the top 4–6. In most systems this is the largest single quality gain after chunking, and it is the stage teams skip because retrieval ‘already works’.

Reranking also gives you a genuine relevance floor. Raw similarity scores are not comparable across queries, so thresholding on them is guesswork; reranker scores are calibrated enough to say ‘nothing clears 0.3, so retrieval found nothing’ — the signal that lets an agent decline rather than ground on the least-bad chunk it happened to get.

Managed retrieval or your own vector store

Google’s managed options — Vertex AI Search, and the RAG Engine in Vertex — take a corpus of documents and handle ingestion, parsing, chunking, embedding, indexing, retrieval and grounding metadata as a service. The two things you cannot easily replicate are the document parsing (tables, PDFs, scanned layouts) and the grounding metadata: per-claim support scores linking spans of the generated answer back to the passages that support them.

Either way, the agent side is the same: retrieval is a function taking a query and returning chunks with ids and sources, wrapped in a tool or a callback. Keep that seam clean and swapping a pgvector prototype for a managed corpus is a one-file change; let retrieval-client specifics leak into prompt construction and the migration becomes a rewrite.

Spending the context budget

Retrieved context is not free real estate. It competes with the system instruction, the tool declarations, and the conversation history for the same window, and long contexts degrade before they overflow — instructions placed above a wall of retrieved text get followed less reliably, and material in the middle of a long block gets attended to least. Ten chunks is usually worse than four, even when they all fit.

def inject_context(callback_context, llm_request):
    chunks = callback_context.state.get("temp:chunks") or []
    budget, used, kept = 3000, 0, []
    for c in chunks:                      # already reranked, best first
        cost = estimate_tokens(c["text"])
        if used + cost > budget:
            break
        used += cost
        kept.append(c)
    if not kept:
        block = "NO SOURCES RETRIEVED. Say you do not have this information."
    else:
        block = "\n\n".join(
            f'[S{i}] source={c["source"]} {c["locator"]}\n{c["text"]}'
            for i, c in enumerate(kept, 1))
    llm_request.append_instructions([f"RETRIEVED SOURCES\n{block}"])

Three things that block does deliberately. It truncates by whole chunks, never mid-passage, because half a policy clause grounds worse than none. It labels each passage with a handle the model can cite. And it handles the empty case explicitly — empty retrieval must produce an instruction to decline, not an absent block the model is free to improvise around.

The grounding instruction, and permission to say ‘I do not know’

Retrieved passages in the context do not by themselves make an answer grounded. The model still has parametric knowledge about refunds, HTTP status codes and your industry, and absent instruction it will blend the two — producing an answer that is mostly from your document with a confident sentence of invention wedged into it. That blended answer is harder to catch than a pure hallucination, because most of it checks out.

The instruction has to do four things, and it belongs in the agent’s system instruction, not appended after the sources. Restrict the answer to the provided sources. Require an inline marker for each factual claim. Explicitly authorise refusal when the sources do not cover the question. And forbid inference beyond what is written — the over-extension failure, where a policy covering damaged items in transit gets stretched to cover damage after delivery, is a reasoning error the model considers helpful.

The refusal clause is the one teams soften under pressure, because ‘I don’t have that information’ reads as a product failure. It is not. An agent that declines on 8% of questions and is right on the rest is usable; one that answers everything and is wrong on 8% is not.

Citations that survive generation

A citation is only worth having if it is verifiable, which means it must resolve to a specific passage a human can open. Ask the model to cite document titles and it will produce plausible titles, including for documents that were never retrieved — a fabricated citation is strictly worse than none, because it manufactures the appearance of diligence.

The pattern that holds up is to hand the model opaque handles and resolve them yourself. Label the injected passages [S1]…[S6], instruct the model to cite only those handles, then post-process: parse the markers out of the response, map each back to the chunk you actually injected, and render the real title, section and URL. A marker that does not resolve — [S9] when you injected six — is a hard signal that the model is confabulating, and it is trivially detectable. Keep the mapping in session state alongside the response so the citation trail is reconstructible later during an audit.

Two refinements pay off. Cite at the claim level, not once at the end, so a reader can check the sentence they doubt. And flag factual sentences carrying no marker — the ungrounded assertions are where errors concentrate.

Measuring groundedness

Every stage needs its own metric, because an aggregate ‘answer quality’ number tells you nothing about where to work. Retrieval is measured on a labelled set of questions with known answer passages: recall@k asks whether the right chunk was retrieved at all, and it is the ceiling on everything after. If recall@20 is 0.6, no prompt engineering gets you past 60% — fix chunking or go hybrid.

Groundedness is measured on the answer: decompose it into atomic claims and check each against the retrieved passages. A judge model does this well because it is a narrow entailment question, not an open one. Run it offline over an eval set, and online as a sampled after_model_callback so production drift shows up.

async def check_grounding(callback_context, llm_response):
    chunks = callback_context.state.get("temp:chunks") or []
    if not chunks or not sampled(rate=0.05):
        return None                        # None = leave the response alone
    verdict = await judge.score(
        answer=text_of(llm_response), sources=chunks)
    if verdict["unsupported"]:
        telemetry.emit("grounding.unsupported",
                       claims=verdict["unsupported"],
                       session=callback_context.state.get("session_id"))
    return None

Track the pair together: groundedness and refusal rate. Groundedness alone is trivially gamed by an agent that declines everything, and a sudden drop in refusals usually means retrieval started returning noise the model is now willing to use.

Freshness and permission-aware retrieval

Two governance properties determine whether a grounded agent is safe to point at a real corpus, and neither is a prompt concern. Freshness: the index is a copy, and a copy drifts. A policy updated on Monday and re-indexed on Friday means four days of confidently cited, correctly formatted, outdated answers — worse than no answer, because the citation vouches for it. Treat index lag as an SLO with an alert, drive re-indexing off document-change events rather than a nightly sweep, and carry an updated_at on every chunk so the agent can surface effective dates on time-sensitive claims.

Permissions: retrieval is a read path into your document store, and an agent that retrieves without regard to the caller’s entitlements will happily paraphrase a restricted document into an answer. The leak has no audit trail, because nobody opened the file. Filter at the query — pass the caller’s identity and groups into the retrieval call so the store never returns unauthorised chunks — rather than filtering results afterwards, which fails open the first time someone forgets the check. And accept the consequence: two users asking the same question must be able to get different answers.

Telling the failure modes apart

Grounded systems fail in a handful of characteristic ways, and each has a different fix at a different stage. The diagnostic habit worth building is to log the retrieved chunk ids with every response, because most arguments about ‘the model hallucinated’ are settled by looking at what was actually in the context.

Note that only one row is fixed by changing the prompt. Grounding is mostly a data-pipeline discipline wearing a prompt-engineering costume, and the highest-leverage stage is almost always the earliest one you have not measured.