Testing a document-grounded AI one stage at a time
When a retrieval-augmented system gives a wrong answer, the cause can sit in the index, the retriever, the model's reading or the final answer. How I test each stage separately so the failure points to a fix.
Retrieval-augmented generation pairs a language model with a search over your own documents. The model answers from passages it is given rather than from memory alone. Lewis and colleagues showed the approach in 2020, and it is now the usual way to make a model answer from a specific body of material. [1]
Building one is quick. Knowing whether it is any good takes longer, mostly because a wrong answer at the end tells you very little about where it went wrong.
Four stages, four different failures
A grounded answer passes through four stages. The document has to be in the index in a usable form. The retriever has to rank the right passage high enough to be passed to the model. The model has to actually use that passage. The final answer has to stay within what the passages support.
Each stage fails differently and has a different fix. Changing the embedding model will not help if the passage was never indexed. A better prompt will not help if the right passage was ranked 40th.
Start with a question set you can mark
The most useful thing I build is a list of real questions, each paired with the passage or passages that answer it. Thirty to fifty is enough to start. Include questions whose answer sits in a table, in a footnote or across two documents, because those are where retrieval tends to struggle.
Also include a few questions the documents cannot answer. A good system should say it does not know. A system that answers them confidently has a stage 4 problem.
Measure each stage on its own
With the question set, each stage gets its own check.
This turns one vague score into four specific ones. A low recall at k points to chunking, the embedding model or hybrid keyword search. A stage 3 failure with the right passage in hand points to the prompt, the context order or the model. A stage 4 failure points to stricter citation instructions and refusing when evidence is missing.
- IndexFor every expected passage, confirm it exists in the index as a complete chunk. Search for its text directly. Chunk boundaries that split a table or a procedure in half show up here.
- RetrieveRun each question through the retriever only and record whether an expected passage appears in the top k. This is recall at k, the fraction of relevant items that were retrieved. [2] Record its rank as well.
- ReadGive the model the correct passage directly and ask the question. If the answer is still wrong, retrieval is not the problem.
- AnswerCheck each claim in the final answer against the passages that were supplied. Any claim without a supporting passage is a failure, even if it happens to be true.
Sources: [2]
Keep the checks running
Documents change, and so do models. I keep the question set beside the code and run the stage checks whenever the index is rebuilt, the chunking changes or the model is swapped. A change that improves one stage can quietly break another, and a fixed question set is the cheapest way to notice.
The same approach works for a large deployment or a local model on one machine. It is the habit that matters: test the parts separately, so a wrong answer points to a fix rather than a guess.
Sources and further reading
- Lewis et al., Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks, NeurIPS 2020 (arXiv:2005.11401)
- Thakur et al., BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models, NeurIPS 2021 Datasets and Benchmarks (arXiv:2104.08663). Uses nDCG@10 and Recall@100.