Build a Multi-Agent Research Assistant

Lesson 2

Give it a tool

Ground your agent with search, stop it from inventing facts

Give it a tool, search + RAG

In Lesson 1 your agent confidently invented a benchmark that doesn't exist. The root cause: it had no way to look anything up. It could only pattern-match against its training data.

The fix is a tool, a Python function the agent can call to retrieve real information before it answers. This lesson adds the simplest possible tool: a search function. You will see hallucination disappear. Then you will discover a new, subtler problem.

What you're building

ResearchAgent, v0.2

Before answering, the agent searches for relevant documents and passes them as context to the LLM. The LLM answers only from that context, no invented facts.

See it first: the new architecture

Agent v0.2: search before you answer

The agent now searches first, feeds the retrieved docs to the LLM as context, then answers. Same loop, one new tool.

One box changed from Lesson 1. The agent now calls a search tool, and the documents it gets back become context for the LLM. The answer is built from retrieved text, not from memory.

What a tool actually is

What is a tool?

A tool is just a function

SearchWeb fetchCalculatorDatabase

Give the model functions it can call, and it stops guessing.

There is nothing mystical here. A tool is a function with a clean interface. Search today; the same pattern gives you calculators, databases, or web fetch later.

What the search tool hides: RAG and vector databases

Inside the search tool: RAG over a vector DB

Keyword search: "AlphaCode" misses "the model from DeepMind". Same meaning, no match.

Semantic search: embeddings sit near each other by meaning, so it finds it anyway.

vector stores: pgvector · Pinecone · FAISS · Weaviate

In production, SearchTool.search() is a RAG pipeline: documents are chunked and embedded into a vector database once; each query is embedded and matched by nearest-neighbour similarity. The MockSearchEngine is a stand-in for exactly this.

Our MockSearchEngine matches keywords. A real one runs retrieval-augmented generation (RAG): documents are chunked and embedded into a vector database once, then each query is embedded and matched by meaning, not exact words. The SearchTool interface stays identical, so none of the rest of the agent changes.

This course keeps retrieval behind that interface on purpose; the focus here is the agent. For the full production RAG system, chunking, hybrid retrieval, re-ranking, and evaluation, see the ML System Design course: RAG retrieval and re-ranking.

Why this matters: generated vs grounded

Generated vs grounded

Same question. The left answer is invented; the right one is grounded in a document the agent actually retrieved.

This is the payoff. The Lesson 1 hallucination disappears, not because the model got smarter, but because it is now reading a real document before it answers.

And the new limit you will hit

The new wall

A multi-part question needs several reasoned searches. Deciding what to look up next, in a loop, is Lesson 3: ReAct.

Grounding fixed one question. But ask the agent to compare two things and a single search only covers half of it. Fixing that needs reasoning across multiple searches, Lesson 3.

What a tool is

A tool is just a Python function with a clean interface:

Python

class SearchTool:
    def search(self, query: str) -> list[str]:
        raise NotImplementedError

The agent calls self.search_tool.search(query) and gets a list of text snippets back. Whether the real implementation hits a vector database, a web API, or a local index does not matter, the agent only sees the interface.

Same pattern as LLMClient: define the shape, swap the implementation.

The new agent loop

Without tools:

Text

query → LLM → answer (hallucinated)

With a search tool:

Text

query → search(query) → retrieved docs → LLM(docs + query) → grounded answer

This pattern has a name: Retrieval Augmented Generation (RAG). You are building it from scratch in about 10 lines.

Your tasks

Write the SearchTool abstract class (one method: search).
Update ResearchAgent.init to accept a search_tool.
Update ResearchAgent.run(): search first → build context string → call the LLM with context + question.
Run the test at the bottom, the AlphaCode answer should now be grounded in retrieved facts.

Python

# ── Shared interfaces ─────────────────────────────────────────────────────────
class LLMClient:
  def chat(self, messages: list) -> str:
      raise NotImplementedError

# TODO: write SearchTool, one method: search(self, query: str) -> list[str]
class SearchTool:
  pass  # ← replace this

# ── MockLLM, now context-aware (given, do not modify) ────────────────────────
class MockLLM(LLMClient):
  def chat(self, messages: list) -> str:
      user = messages[-1]["content"]

# Grounded mode: context was injected
      if "Context:" in user and "Question:" in user:
          context = user.split("Context:")[1].split("Question:")[0].strip().lower()
          question = user.split("Question:")[1].strip().lower()
          has_alphacode = "alphacode" in context or "codecontests" in context
          asks_codex    = "codex" in question

if has_alphacode and asks_codex:
              return (
                  "Based on retrieved context: AlphaCode (DeepMind, 2022) benchmarked on "
                  "CodeContests using a generate-and-filter approach. However, my retrieved "
                  "context does not include Codex data, I cannot compare them accurately."
              )
          if has_alphacode:
              return (
                  "Based on retrieved context: AlphaCode evaluated on CodeContests (not "
                  "HumanEval), used generate-and-filter with test-case reranking, and "
                  "matched the median competitive programmer's performance."
              )
          return f"Based on retrieved context: {user.split('Context:')[1].split('Question:')[0].strip()[:180]}..."

# No context injected, hallucination mode (same as Lesson 1)
      return "Based on my training: [confident but unverifiable, no retrieval was done]"

# ── MockSearchEngine (given, do not modify) ───────────────────────────────────
class MockSearchEngine(SearchTool):
  DB = {
      "alphacode": [
          "AlphaCode (DeepMind, 2022) was evaluated on CodeContests, not HumanEval.",
          "AlphaCode used a generate-and-filter strategy: sample many programs, then "
          "filter by test cases and rerank by a separate model.",
          "AlphaCode achieved performance comparable to the median competitive programmer.",
      ],
      "codex": [
          "OpenAI Codex was trained on public GitHub code and evaluated on HumanEval.",
          "HumanEval consists of 164 hand-written Python programming problems.",
          "Codex (code-davinci-002) achieved ~72% pass@1 on HumanEval.",
      ],
  }

def search(self, query: str) -> list[str]:
      query_lower = query.lower()
      results = []
      for keyword, docs in self.DB.items():
          if keyword in query_lower:
              results.extend(docs)
      return results if results else [f"No documents found for: '{query}'"]

# ── Your turn: update ResearchAgent ───────────────────────────────────────────
class ResearchAgent:
  def __init__(self, llm: LLMClient, search_tool: SearchTool):
      self.llm         = llm
      self.search_tool = search_tool

def run(self, query: str) -> str:
      # Step 1: retrieve relevant documents
      docs = self.search_tool.search(query)

# Step 2: format them as a context string
      context = "\n".join(f"- {d}" for d in docs)

# Step 3: call the LLM with context + question
      # Use this message format:
      #   system: "Answer based ONLY on the provided context. If the context
      #            doesn't cover the question, say so."
      #   user:   f"Context:\n{context}\n\nQuestion: {query}"
      pass  # ← replace this

# ── Test: AlphaCode query (the Lesson 1 hallucination case) ───────────────────
agent = ResearchAgent(llm=MockLLM(), search_tool=MockSearchEngine())
result = agent.run("What did the AlphaCode paper report about its benchmark?")
print(result)

First run loads the Python runtime (~10 MB) — takes ~5–10 seconds. Subsequent runs are instant.

The answer now references CodeContests, not the invented HumanEval figure from Lesson 1. The agent is grounded.

Show solutiontry it yourself first

What's happening: this is still one agent. It does two things, retrieve then answer, and it answers only from what it retrieved. Nothing runs in parallel yet, and there are no sub-agents (those start in Lesson 4's Planner and Lesson 5's parallel executors).

Two pieces to fill in, the SearchTool interface and run():

Python

# the tool interface: any search backend implements this one method
class SearchTool:
    def search(self, query: str) -> list[str]:
        raise NotImplementedError
 
# retrieve first, then answer ONLY from the retrieved context
def run(self, query: str) -> str:
    docs = self.search_tool.search(query)
    context = "\n".join(f"- {d}" for d in docs)
    messages = [
        {"role": "system", "content": "Answer based ONLY on the provided context. "
                                       "If the context doesn't cover the question, say so."},
        {"role": "user",   "content": f"Context:\n{context}\n\nQuestion: {query}"},
    ]
    return self.llm.chat(messages)

Expected output for agent.run("What did the AlphaCode paper report about its benchmark?"):

Text

Based on retrieved context: AlphaCode evaluated on CodeContests (not HumanEval),
used generate-and-filter with test-case reranking, and matched the median
competitive programmer's performance.

See what the agent did: a verbose trace

Add a verbose flag and print each step. This is the trace for one agent; in Lesson 5 the same idea, printed per executor, becomes a parallel multi-agent trace.

Python

def run(self, query: str, verbose: bool = False) -> str:
    if verbose: print(f"[agent] SEARCH    {query!r}")
    docs = self.search_tool.search(query)
    if verbose:
        print(f"[agent] RETRIEVED {len(docs)} doc(s)")
        for d in docs:
            print(f"          - {d[:60]}...")
    context = "\n".join(f"- {d}" for d in docs)
    messages = [
        {"role": "system", "content": "Answer based ONLY on the provided context."},
        {"role": "user",   "content": f"Context:\n{context}\n\nQuestion: {query}"},
    ]
    answer = self.llm.chat(messages)
    if verbose: print(f"[agent] ANSWER    grounded in {len(docs)} doc(s)")
    return answer

Trace output for the AlphaCode query (run(query, verbose=True)):

Text

[agent] SEARCH    'What did the AlphaCode paper report about its benchmark?'
[agent] RETRIEVED 3 doc(s)
          - AlphaCode (DeepMind, 2022) was evaluated on CodeConte...
          - AlphaCode used a generate-and-filter strategy: sample...
          - AlphaCode achieved performance comparable to the medi...
[agent] ANSWER    grounded in 3 doc(s)
 
Based on retrieved context: AlphaCode evaluated on CodeContests (not HumanEval),
used generate-and-filter with test-case reranking, and matched the median
competitive programmer's performance.

Reading the trace: one query goes in, one SEARCH fires, three documents come back, and the answer is built only from them. That is the whole job of a Lesson 2 agent. The moment you want several of these running at once, one per sub-question, you need a Planner to split the work (Lesson 4) and asyncio.gather to run the executors in parallel (Lesson 5). That is where the trace gains multiple interleaved [exec 1] / [exec 2] / [exec 3] lines, the multi-agent trace.

Now try a two-entity question

Your agent can search once. That works when one query fetches everything needed. Research rarely works that way.

The new wall

Run the cell below. The question asks about two papers. Notice what the agent says about the one it can't find in context.

Python

# Self-contained, includes the full working agent from above

class LLMClient:
  def chat(self, messages: list) -> str:
      raise NotImplementedError

class SearchTool:
  def search(self, query: str) -> list[str]:
      raise NotImplementedError

class MockLLM(LLMClient):
  def chat(self, messages: list) -> str:
      user = messages[-1]["content"]
      if "Context:" in user and "Question:" in user:
          context  = user.split("Context:")[1].split("Question:")[0].strip().lower()
          question = user.split("Question:")[1].strip().lower()
          has_alphacode = "alphacode" in context or "codecontests" in context
          has_codex     = "codex" in context
          asks_codex    = "codex" in question

if has_alphacode and not has_codex and asks_codex:
              return (
                  "Based on retrieved context: AlphaCode used CodeContests and a "
                  "generate-and-filter approach. However, my retrieved context does "
                  "not contain information about Codex, I cannot compare them."
              )
          if has_alphacode and has_codex:
              return (
                  "Based on retrieved context: AlphaCode (CodeContests, generate-and-filter) "
                  "vs Codex (HumanEval, ~72% pass@1). Both are code generation models "
                  "using different benchmarks and training strategies."
              )
      return "No context retrieved, cannot answer reliably."

class MockSearchEngine(SearchTool):
  DB = {
      "alphacode": [
          "AlphaCode evaluated on CodeContests, not HumanEval.",
          "AlphaCode used a generate-and-filter approach with test-case reranking.",
      ],
      "codex": [
          "Codex evaluated on HumanEval (164 Python problems).",
          "Codex (code-davinci-002) achieved ~72% pass@1.",
      ],
  }
  def search(self, query: str) -> list[str]:
      results = []
      for kw, docs in self.DB.items():
          if kw in query.lower():
              results.extend(docs)
      return results or [f"No documents found for: '{query}'"]

class ResearchAgent:
  def __init__(self, llm, search_tool):
      self.llm         = llm
      self.search_tool = search_tool
  def run(self, query: str) -> str:
      docs    = self.search_tool.search(query)
      context = "\n".join(f"- {d}" for d in docs)
      messages = [
          {"role": "system", "content": "Answer ONLY from the provided context. If context is insufficient, say so."},
          {"role": "user",   "content": f"Context:\n{context}\n\nQuestion: {query}"},
      ]
      return self.llm.chat(messages)

agent = ResearchAgent(llm=MockLLM(), search_tool=MockSearchEngine())

# ── The wall: two papers, one search ──────────────────────────────────────────
query = "Compare AlphaCode and Codex, what benchmark did each use and what was the core algorithmic difference?"
print("Query:", query)
print()

# What did the search actually retrieve?
docs = MockSearchEngine().search(query)
print("Retrieved docs:")
for d in docs: print(" ", d)
print()

print("Agent answer:")
print(agent.run(query))

First run loads the Python runtime (~10 MB) — takes ~5–10 seconds. Subsequent runs are instant.

What just happened

The search matched "alphacode" in the query and returned AlphaCode docs, but "codex" also appears in the query, so Codex docs were returned too.

Run it again and look at the retrieved docs section. Both papers were found this time, because both keywords happened to be in the single query string.

Now imagine a query where the two topics have no lexical overlap. Or a question that requires three sub-searches with different keywords. A single search(full_query) call is brittle:

No control over what to search: the query is whatever the user typed, verbatim
No iteration: the agent cannot look at search results, notice a gap, and search again
No decomposition: multi-part questions become one undifferentiated search blob

The pattern you just built has a name

Retrieval Augmented Generation (RAG), retrieve relevant documents, inject them as context, let the LLM answer from context only. It eliminates hallucination for covered topics. The limitation is that it retrieves once, blindly, before any reasoning happens. The agent cannot adapt its search strategy based on what it finds.

The fix requires the agent to think before searching, observe what it got, and search again if needed. That is a loop, and it has a name too: ReAct.

Practice it yourself

Build the RAG pipeline by hand — retrieve, format, ground. Real Python, hidden tests, no setup.

Format docs as contextEasy

One bullet per snippet — the block the LLM reads.

MockSearchEngine retrievalEasy

Keyword match the query, or return "not found".

The RAG runMedium

Search → ground → answer, only from retrieved context.

Where you are

Capability	Status
LLM call	✓ Lesson 1
Tool interface	✓ This lesson
Search + RAG	✓ This lesson
Adaptive, multi-step search	✗ Lesson 3
Planning and decomposition	✗ Lesson 4

Quiz

Check your understanding

1 / 6

What is a 'tool' in agent terms, and what changes when you give the agent one?

Practice it yourself

A real search tool ranks and fuses results before the LLM sees them. Build the retrieval-quality primitives underneath, in the browser with hidden tests, no setup:

Cross-encoder rerank top-kEasy

Re-score retrieved candidates and keep the top-k, the step that sharpens what your tool returns.

Recall@k (retrieval)Easy

Measure whether the right document made it into the top-k, how you'd know the tool actually works.

Reciprocal rank fusionMedium

Merge several ranked lists into one, exactly what you'd need to combine multi-search results.

Give it a tool

See it first: the new architecture

What a tool actually is

What the search tool hides: RAG and vector databases

Why this matters: generated vs grounded

And the new limit you will hit

You are here

What a tool is

The new agent loop

See what the agent did: a verbose trace

Now try a two-entity question

What just happened

Practice it yourself

Where you are

Quiz

Practice it yourself