An AI Agent That Searches Your Documents

📘 AI Agents & RAG Course · Lesson 5 of 7

In Lesson 3 every question triggered exactly one search. An agent can do better: it can rephrase and search again, search two topics for one question, or skip searching when the question is a simple calculation.

Retrieval as a tool

tools = [{
    "name": "search_policies",
    "description": "Search company policy documents. Returns passages with source names and scores. "
                   "Search again with different words if the first results don't answer the question.",
    "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]},
}]

def run_tool(name, args):
    hits = [h for h in search(args["query"], 4) if h[0] >= MIN_SCORE]
    if not hits:
        return "No matching passages."
    return "\n\n".join(f"[{d['source']}] (score {s:.1f})\n{d['text']}" for s, d in hits)

Reuse the loop from Lesson 2 with this tool and a system prompt:

SYSTEM = ("You answer employee questions using the search_policies tool. "
          "Always search before answering policy questions. Cite sources in [brackets]. "
          "If searches return nothing relevant, say you don't know and suggest who to ask.")

What changes in practice

Ask “I’m travelling to Pune for 3 days — what can I claim in total?” and watch the log. A good run searches the travel policy, finds the daily allowance, maybe searches again for hotel limits, and then does the arithmetic in its answer with both sources cited.

Add a calculator

Models are good at arithmetic but not perfect. For money, give the agent a tiny calculator tool and tell it to use it for every sum:

{"name": "calc", "description": "Evaluate an arithmetic expression like 3*2500+1800. Use for all money calculations.",
 "input_schema": {"type": "object", "properties": {"expr": {"type": "string"}}, "required": ["expr"]}}

import ast, operator as op
OPS = {ast.Add: op.add, ast.Sub: op.sub, ast.Mult: op.mul, ast.Div: op.truediv}
def calc(expr):
    def ev(n):
        if isinstance(n, ast.Constant): return n.value
        if isinstance(n, ast.BinOp): return OPS[type(n.op)](ev(n.left), ev(n.right))
        raise ValueError("unsupported")
    return str(ev(ast.parse(expr, mode="eval").body))