RAG Agents

Real-World Scenario

A developer asks a support agent: “My endpoint is slow, about 12 seconds per response. Will I still get shipment events?” The answer is spread across two documentation pages. One says a delivery taking over 10 seconds counts as failed; the other says failed deliveries are retried 5 times. Neither page uses the words “slow” or “12 seconds”. One search with the user’s words doesn’t find either page. An agent that reads its first results, then searches in the documentation’s own vocabulary, finds both.

The search here is real: BM25, the keyword ranking behind most search engines, implemented in a dozen lines. The model’s decisions, meaning which query to try next and how to phrase the answer, are scripted as a stand-in and marked as such. The loop, the evidence handling and the checks are the code a real RAG agent needs, and none of it needs an API key.

The Documents and the Search Tool

Nine short pages of API documentation. Save them as kb.py.

DOCS = {
    "auth-1":    "Every request needs a bearer token. Tokens expire after 60 minutes; request a new one from /oauth/token.",
    "rate-1":    "The API allows 120 requests per minute per key. Exceeding it returns HTTP 429 with a Retry-After header.",
    "rate-2":    "Burst traffic is smoothed by a token bucket: up to 20 requests may arrive in the same second.",
    "hooks-1":   "Webhooks deliver shipment events to your HTTPS endpoint as JSON, signed with an HMAC-SHA256 header.",
    "hooks-2":   "Failed webhook deliveries are retried 5 times with exponential backoff starting at 30 seconds.",
    "hooks-3":   "A webhook delivery counts as failed when your endpoint returns a non-2xx status or takes over 10 seconds.",
    "errors-1":  "HTTP 409 means the shipment was already created with the same idempotency key.",
    "pages-1":   "List endpoints return 50 items per page; follow the next_cursor field until it is null.",
    "sandbox-1": "Sandbox keys start with sk_test_ and never create real shipments or charges.",
}

BM25 scores a document higher when it contains query words that are rare across the collection, repeated, and in a short document. Save this as search.py.

import math, re
from collections import Counter
from kb import DOCS

def terms(text):
    return re.findall(r"[a-z0-9]+", text.lower())

N = len(DOCS)
AVG = sum(len(terms(t)) for t in DOCS.values()) / N
DF = Counter(w for t in DOCS.values() for w in set(terms(t)))

def bm25(query, k=2, k1=1.5, b=0.75):
    """Keyword ranking used by most search engines: rare, repeated query words score high."""
    scores = {}
    for doc_id, text in DOCS.items():
        tf, length = Counter(terms(text)), len(terms(text))
        scores[doc_id] = sum(
            math.log(1 + (N - DF[w] + 0.5) / (DF[w] + 0.5)) * tf[w] * (k1 + 1)
            / (tf[w] + k1 * (1 - b + b * length / AVG)) for w in terms(query) if w in tf)
    return sorted(scores.items(), key=lambda p: -p[1])[:k]

if __name__ == "__main__":
    q = "If my webhook endpoint is down, how many times is a delivery retried?"
    for doc_id, score in bm25(q, k=3):
        print(f"{score:5.2f}  {doc_id:9} {DOCS[doc_id][:60]}")

$ python search.py
 5.52  hooks-2   Failed webhook deliveries are retried 5 times with exponenti
 5.42  hooks-3   A webhook delivery counts as failed when your endpoint retur
 3.40  rate-2    Burst traffic is smoothed by a token bucket: up to 20 reques

A question phrased in the documentation’s own terms finds the right pages: the retry policy ranks first, and the failure definition second. The third result is noise, matching only common words. Real questions are rarely phrased this helpfully.

The Failure: One Search With the User’s Words

This is basic RAG: search once with the question as asked, and answer from the top two results.

from search import bm25

question = "My endpoint is slow, about 12 seconds per response. Will I still get shipment events?"
for doc_id, score in bm25(question):              # retrieve once, top 2, then answer
    print(f"{score:4.2f}  {doc_id}")

$ python single_shot.py
4.70  hooks-1
2.79  pages-1

The top two results are the webhook overview and the pagination page. Neither says what happens to a slow delivery; the pagination page matched only because both texts contain words like “per”. A model answering from this context can say events are sent as JSON, then guess the rest. The two pages that hold the answer, about the 10-second limit and the retries, never reach it.

The Agent Loop: Read, Rephrase, Search Again

The agent treats search as a tool it can call repeatedly. Code enforces a minimum score, so weak matches aren’t treated as evidence, and a search budget.

from kb import DOCS
from search import bm25

def run(model_steps, max_searches=3, min_score=2.0):
    """The agent loop: the model chooses searches until it can answer; code keeps the evidence."""
    evidence, searches = {}, 0
    for action, arg in model_steps:
        if action == "answer":
            return arg, evidence
        searches += 1
        if searches > max_searches:
            return "[stopped: search budget used up]", evidence
        hits = [d for d, score in bm25(arg) if score >= min_score]   # weak matches aren't evidence
        evidence.update({d: DOCS[d] for d in hits})
        print(f"  search {arg[:44]!r:47} -> {hits or 'nothing relevant'}")

QUESTION = "My endpoint is slow, about 12 seconds per response. Will I still get shipment events?"
# What a model decides after reading each result (scripted here):
STEPS = [("search", QUESTION),                                       # events exist, but slowness?
         ("search", "webhook delivery failed timeout seconds"),      # rephrased after reading hooks-1
         ("answer", "Not reliably: responses over 10 seconds count as failed [hooks-3]. Failed "
                    "deliveries are retried 5 times with backoff from 30 seconds [hooks-2].")]

if __name__ == "__main__":
    answer, evidence = run(STEPS)
    print("evidence:", sorted(evidence))
    print("answer:", answer)

$ python agent.py
  search 'My endpoint is slow, about 12 seconds per re'  -> ['hooks-1', 'pages-1']
  search 'webhook delivery failed timeout seconds'       -> ['hooks-3', 'hooks-2']
evidence: ['hooks-1', 'hooks-2', 'hooks-3', 'pages-1']
answer: Not reliably: responses over 10 seconds count as failed [hooks-3]. Failed deliveries are retried 5 times with backoff from 30 seconds [hooks-2].

The first search, with the user’s words, returns the same two pages as before. The model reads them, learns the product calls these “webhook deliveries”, and searches again in the documentation’s vocabulary. That search finds both missing pages. The answer then cites exactly the two documents that support it. pages-1 is in the evidence but is never cited, which is fine: evidence is what was available, and citations are what was used.

The scripted STEPS are the model’s contribution: after each search, a real model sees the results and chooses the next query or the answer. Everything else is fixed code. Evidence is tracked by document id, low scores are excluded, and the loop can’t run past its budget.

Verifying Citations

Citations are only useful if they’re true. Check them in code before showing the answer.

import re
from agent import run, STEPS

def check_citations(answer, evidence):
    """Each cited id must be retrieved evidence, and its sentence's numbers must appear in it."""
    problems = []
    for sentence in re.split(r"(?<=\.)\s+", answer):
        for doc_id in re.findall(r"\[([\w-]+)\]", sentence):
            if doc_id not in evidence:
                problems.append(f"[{doc_id}] was never retrieved")
                continue
            claims = re.sub(r"\[[\w-]+\]", "", sentence)          # don't read "hooks-3" as a 3
            for n in re.findall(r"\b\d+\b", claims):
                if not re.search(rf"\b{n}\b", evidence[doc_id]):
                    problems.append(f"{n} is not in [{doc_id}]")
    return problems or ["ok"]

_, evidence = run(STEPS)
print(check_citations(STEPS[-1][1], evidence))
print(check_citations("Deliveries are retried 3 times [hooks-2]. Tokens last 60 minutes [auth-1].",
                      evidence))

$ python citations.py
  search 'My endpoint is slow, about 12 seconds per re'  -> ['hooks-1', 'pages-1']
  search 'webhook delivery failed timeout seconds'       -> ['hooks-3', 'hooks-2']
['ok']
['3 is not in [hooks-2]', '[auth-1] was never retrieved']

The agent’s answer passes: both cited pages were retrieved, and every number in each sentence appears in the page it cites. The second answer is the kind of error models make. It gets the retry count wrong (3 instead of 5), and it cites auth-1, which was never retrieved, so the claim can’t have come from the evidence. Both are caught.

The re.sub line fixes a real bug in the first version of this check. The number regex found the “3” in [hooks-3] and reported it missing from the page. Strip citation markers before looking for claims in the text.

Knowing When Not to Answer

When nothing relevant is retrieved, the right answer is to say so.

from agent import run

for question in ["How long is a bearer token valid?", "Do you support cash on delivery payments?"]:
    print(question)
    answer, evidence = run([("search", question), ("answer", None)])
    if not evidence:                          # nothing cleared the score bar: don't guess
        print("  -> I couldn't find this in the API docs; please contact support.")
    else:
        print(f"  -> answer from {sorted(evidence)}")

$ python abstain.py
How long is a bearer token valid?
  search 'How long is a bearer token valid?'             -> ['auth-1', 'rate-2']
  -> answer from ['auth-1', 'rate-2']
Do you support cash on delivery payments?
  search 'Do you support cash on delivery payments?'     -> nothing relevant
  -> I couldn't find this in the API docs; please contact support.

The token question retrieves the authentication page and is answered from it. The cash-on-delivery question matches nothing above the minimum score, so the evidence is empty and the agent says it couldn’t find an answer instead of improvising a payment policy. The rate-2 page came along with auth-1 because it mentions tokens (bucket tokens, not auth tokens). A shared word isn’t shared meaning, which is why the answer step should cite only the pages it actually used.

Practice

1. The agent uses min_score=2.0. Print what a single search admits as evidence at thresholds 0, 2 and 3, for the slow-endpoint question and the cash-on-delivery question (use k=4).

   Solution

   from search import bm25
   from agent import QUESTION
   
   for q in [QUESTION, "Do you support cash on delivery payments?"]:
       scored = bm25(q, k=4)
       print(q[:40])
       for bar in (0.0, 2.0, 3.0):
           print(f"  min_score {bar}: {[d for d, s in scored if s > 0 and s >= bar]}")

   $ python sol_threshold.py
   My endpoint is slow, about 12 seconds pe
     min_score 0.0: ['hooks-1', 'pages-1', 'hooks-3', 'rate-1']
     min_score 2.0: ['hooks-1', 'pages-1', 'hooks-3']
     min_score 3.0: ['hooks-1']
   Do you support cash on delivery payments
     min_score 0.0: ['hooks-3']
     min_score 2.0: []
     min_score 3.0: []

   At 0, the cash question admits the webhook failure page as “evidence” for a payments question, and the slow-endpoint question picks up the rate-limit page. At 3, the slow-endpoint question keeps only the overview and loses hooks-3, the page it needs. 2.0 keeps what matters for both. Scores depend on the collection and the ranking method, so choose the threshold from real questions, including ones your documents can’t answer.

2. A model can get stuck repeating one search. Script six identical searches before an answer and confirm the loop stops at the budget.

   Solution

   from agent import run
   
   looping = [("search", "failed webhook delivery retries")] * 6 + [("answer", "done")]
   answer, evidence = run(looping)
   print(answer, "| evidence:", sorted(evidence))

   $ python sol_budget.py
     search 'failed webhook delivery retries'               -> ['hooks-3', 'hooks-2']
     search 'failed webhook delivery retries'               -> ['hooks-3', 'hooks-2']
     search 'failed webhook delivery retries'               -> ['hooks-3', 'hooks-2']
   [stopped: search budget used up] | evidence: ['hooks-2', 'hooks-3']

   The loop runs three searches and returns the stop message instead of the answer. The evidence gathered is kept, so a real agent can fall back to answering from it, with a note that the search ran out. Without the budget, each repeated search costs a model call and a search call.

3. The budget stops a loop late. Report after each search whether it found anything new, so the model is told to stop as soon as searching stops helping.

   Solution

   from kb import DOCS
   from search import bm25
   
   def run(model_steps, min_score=2.0):
       evidence = {}
       for action, arg in model_steps:
           if action == "answer":
               return arg, evidence
           new = [d for d, s in bm25(arg) if s >= min_score and d not in evidence]
           evidence.update({d: DOCS[d] for d in new})
           print(f"  search {arg!r:40} new: {new or 'none, tell the model to stop searching'}")
   
   run([("search", "webhook delivery failed timeout"), ("search", "failed webhook delivery retries"),
        ("search", "webhook retry backoff"), ("answer", "...")])

   $ python sol_nothing_new.py
     search 'webhook delivery failed timeout'        new: ['hooks-3', 'hooks-2']
     search 'failed webhook delivery retries'        new: none, tell the model to stop searching
     search 'webhook retry backoff'                  new: none, tell the model to stop searching

   The first search finds both relevant pages. The next two find nothing that isn’t already in the evidence, and the message for that result says so plainly, which is what a model needs to decide to answer now. “No new evidence” is a better stop signal than a fixed count, and the budget remains as the backstop.