Simple Agent with RAG

Welcome! Today, we’ll see how to build a simple AI agent that answers questions using external knowledge. Even advanced AI models sometimes need to look up information — just as you might check your notes or Google before answering a tough question.

Retrieval-Augmented Generation (RAG) helps AI agents find and use relevant information from a knowledge base for more accurate answers. By the end, you’ll know how to create a basic agent that uses RAG to answer user queries more effectively.

First, the agent needs access to knowledge. Here, our knowledge base is a JSON file with study tips and resources.

Here’s how to load it in Python:

import os
import json

def load_data(file_name):
    """Load knowledge base from JSON file."""
    current_dir = os.path.dirname(__file__)
    dataset_file = os.path.join(current_dir, "data", file_name)
    with open(dataset_file, 'r') as file:
        return json.load(file)

data = load_data("data.json")
print(data)  # [{'id': 1, 'content': "Read Chapter 7 of 'Clean Code' — focus on writing small, single-purpose functions."}, {'id': 2, 'content': "Experiment with React's useContext by creating a theme toggler component."}, {'id': 3, 'content': 'Review different types of SQL joins — especially LEFT and FULL OUTER joins.'}, {'id': 4, 'content': 'Watch lecture on consensus algorithms — try to summarize Paxos in your own words.'}, {'id': 5, 'content': 'Figure out the difference between React Query and Redux for async data handling.'}, {'id': 6, 'content': "Draft blog post: '3 Things I Learned About Writing Better Functions.'"}, {'id': 7, 'content': 'Ask Aram about good beginner-friendly open-source projects to contribute to.'}]

This code locates and opens data.json, loading its content as a list of dictionaries. Each dictionary is a document, for example:

{
    "id": 2,
    "content": "Experiment with React's useContext by creating a theme toggler component."
}

This simple structure works well for small knowledge sets. Larger, real-world knowledge bases follow the same principle: load and access information for the agent.

With the knowledge base loaded, the next step is retrieval — finding the most relevant document for a user’s question. We use a basic method: word overlap. We count how many words from the user’s question appear in each document. The document with the most overlap is chosen.

Here’s the code:

def rag_retrieval(query, knowledge_base):
    query_words = set(query.lower().split())
    best_doc = None
    best_overlap = 0
    for doc in knowledge_base:
        doc_words = set(doc["content"].lower().split())
        overlap = len(query_words & doc_words)
        if overlap > best_overlap:
            best_overlap = overlap
            best_doc = doc
    return best_doc if best_overlap > 0 else None

# Example usage:
query = "How should I start tackling React useContext"
best_doc = rag_retrieval(query, data)
print(best_doc)  # {'id': 2, 'content': "Experiment with React's useContext by creating a theme toggler component."}

The query and each document are split into lowercase words, and the code counts overlapping words. The document with the highest overlap is returned. This method is simple and fast. More advanced systems use semantic search, but word overlap is a great way to start understanding retrieval.

After finding the most relevant document, we give this context to the agent. The agent uses both the user’s question and the retrieved document to answer. We build a prompt that combines both:

def build_prompt(user_prompt, rag_document=None):
    if rag_document:
        return f"Based on the following document: {rag_document['content']}, {user_prompt}"
    return user_prompt

# Example usage:
user_prompt = "How should I start tackling React useContext"
prompt_with_context = build_prompt(user_prompt, best_doc)
print(prompt_with_context)  # Based on the following document: Experiment with React's useContext by creating a theme toggler component., How should I start tackling React useContext

If a document is found, the prompt includes its content. If not, the agent just gets the original question. Providing context helps the agent focus on useful information, leading to better answers. This is the core of RAG: retrieval plus generation.

Now, let’s see how the agent uses the prompt to generate a response. The agent is defined with instructions and run using a helper:

from agents import Agent, Runner

AGENT = Agent(
    name="Learning Assistant",
    instructions=(
        "You are a personal learning assistant. "
        "Whenever asked a question about learning plans, use the context provided to answer questions."
    )
)

def ask_agent(prompt):
    try:
        result = Runner.run_sync(AGENT, prompt)
        return result.final_output
    except Exception as e:
        print(f"Agent error: {e}")
        raise

# Example usage:
response = ask_agent(prompt_with_context)
print(response)  # (Output will depend on the agent's implementation, e.g., "To start with React's useContext, try building a simple theme toggler as described in the document.")

The agent gets the prompt (with or without context), generates a response using its instructions and the context, and prints the response. This workflow — retrieval, prompt building, and response generation — is the foundation of a simple RAG agent.

You’ve learned how to build a simple agent using Retrieval-Augmented Generation (RAG) for better answers. We covered loading a knowledge base from JSON, retrieving the most relevant document using word overlap, building a prompt with the user’s question and context, and running the agent to generate a response. This approach makes AI agents more helpful, especially for specific topics or documents.

Now it’s your turn! Next, you’ll practice building and running your own RAG-powered agent. You’ll load data, retrieve relevant information, and generate answers using an agent. This hands-on work will reinforce your understanding and prepare you for more advanced RAG techniques.