Build a RAG Knowledge Base with Gemini
Build a CLI knowledge base that answers questions using RAG and Gemini.
Introduction
30 Second Summary
Ever asked an AI a question about your own notes or documents and gotten a completely made-up answer? That's because LLMs can only answer from their training data. They have no idea what's in your files.
In this project, you will build a command-line Q&A system that answers natural language questions about your documents. It retrieves the most relevant passages from a local vector database and generates cited answers using Gemini AI.
What You'll Build
A Python-based RAG pipeline that ingests markdown documents, chunks and embeds them into a ChromaDB vector store, then retrieves relevant passages to generate cited answers via the Gemini API.
By the end of this project, you'll have:
- A document ingestion script that loads, chunks, and embeds markdown files into a local ChromaDB vector store.
- A query script that takes a question, finds the most relevant chunks via vector similarity search, and generates an answer using Gemini with source citations.
- A set of sample markdown documents as your knowledge base.
- Secret Mission: Add relevance scoring so the system says "I don't know" when no chunk is relevant enough, and display similarity scores to the user.
Are there any prerequisites?
You should be comfortable with basic Python (functions, loops, file I/O) and the command line. No prior experience with AI, embeddings, or vector databases is needed.
Set Up Your Environment
Before you can build a RAG system, you need the right tools installed and configured.
In this step, you will create an isolated Python project folder, install the AI and vector database libraries, and get an API key to access Gemini.
In this step, get ready to:
- Create a project folder and Python virtual environment.
- Install the google-genai and chromadb packages.
- Get a Gemini API key and confirm it works.
Create a project folder and virtual environment
A virtual environment gives your project its own isolated set of packages so nothing conflicts with other Python projects on your machine.
- Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Terminal and press Enter to open it.
- Check your Python version by running this command:
python3 --version
You should see a version number like Python 3.x.x printed in your terminal.
✔️ I see version 3.10 or higher
Your Python version is compatible. Move on to the next instruction below.
ⓧ I see a version lower than 3.10
The google-genai SDK requires Python 3.10 or higher. You need to install a newer version.
The simplest way on macOS is to download the official installer from python.org.
- Visit python.org/downloads/macos and download the latest macOS installer.
- Run the installer and follow the prompts.
- Once installed, close and reopen Terminal, then run python3 --version again to confirm the version is 3.10 or higher.
Why the system Python might be old
macOS ships with a Python 3 version that Apple maintains for its own tools. This is often older than what modern packages require. Installing from python.org creates a separate, newer Python that coexists with the system version.
ⓧ Command not found
Python 3 is not installed on your machine. Install it from the official source.
- Visit python.org/downloads/macos and download the latest macOS installer.
- Run the installer and follow the prompts.
- Once installed, close and reopen Terminal, then run python3 --version to confirm a version of 3.10 or higher appears.
- Move to your Desktop by running this command:
cd ~/Desktop
- Create your project folder and move into it by running these commands:
mkdir rag-knowledge-base
cd rag-knowledge-base
- Create and activate a virtual environment by running these commands:
python3 -m venv venv
source venv/bin/activate
You should see (venv) appear at the start of your terminal prompt. This confirms you are working inside the virtual environment.
What does activating a virtual environment do?
When you activate a virtual environment, your terminal uses the Python and pip binaries inside the venv/ folder instead of the system-wide ones. Any packages you install go into this folder, keeping your project self-contained.
Don't see (venv) in your prompt?
Make sure you ran both commands in sequence. If you accidentally closed the terminal, re-open it and run cd ~/Desktop/rag-knowledge-base followed by source venv/bin/activate again.
Help me activate my Python virtual environment
Install dependencies
Your project needs three packages: the Google GenAI SDK for calling Gemini, ChromaDB for storing vectors locally, and python-dotenv for loading environment variables from a file.
- Create a requirements.txt file that pins the exact package versions by running this command:
cat > requirements.txt << 'EOF'
google-genai==2.10.0
chromadb==1.5.9
python-dotenv==1.1.0
EOF
What is requirements.txt?
A requirements.txt file lists every Python package your project needs, pinned to specific versions. This makes your project reproducible. Anyone can run pip install -r requirements.txt and get the exact same setup.
- Confirm the file was created by running:
cat requirements.txt
You should see the three package lines printed in your terminal.
- Install all three packages by running:
pip install -r requirements.txt
This will download and install the packages. It may take a minute or two.
Seeing a permission error or build failure?
Make sure you see (venv) in your terminal prompt. If you don't, run source venv/bin/activate first, then try the install again.
If you see a build error related to chromadb, confirm your Python version is 3.10 or higher with python3 --version.
Help me fix my pip install error
- Verify the packages installed correctly by running:
python3 -c "from google import genai; import chromadb; print('All packages imported successfully')"
You should see All packages imported successfully printed with no errors.
✔️ Awesome, I've got everything!
Great. Your packages are installed and working.
ⓧ I'd like to double check the full code
Your requirements.txt should contain exactly this:
google-genai==2.10.0
chromadb==1.5.9
python-dotenv==1.1.0
Get your Gemini API key
To call the Gemini API, you need an API key. Google AI Studio lets you generate one for free with no credit card required.
- Open your browser and go to aistudio.google.com/api-keys.
- Sign in with your Google account.
- Click Create API key.
- If prompted to select a project, choose the default project that AI Studio creates for you.
- Copy the API key that appears. You will only see the full key once, so keep this tab open.
Is this really free?
Yes. The Gemini API free tier covers both text generation and embedding tokens at no charge. This project stays well within the free tier limits. No credit card is needed.
- Back in your terminal, create a .env file to store your key by running the following command (replace your-api-key-here with the key you just copied):
echo "GEMINI_API_KEY=your-api-key-here" > .env
Why use a .env file?
Storing secrets in a .env file keeps them out of your code. The python-dotenv package loads these values as environment variables at runtime, so your scripts can access the key without it ever appearing in your Python files.
- Test that your API key works by running this command:
python3 << 'EOF'
from google import genai
from dotenv import load_dotenv
import os
load_dotenv()
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
response = client.interactions.create(
model="gemini-2.5-flash",
input="Say hello in one word"
)
print(response.output_text)
EOF
You should see a short response from Gemini (something like Hello) printed in your terminal.
What does this test do?
This script loads your API key from the .env file, creates a Gemini client, and sends a simple prompt. If you get a response, your key is valid and the SDK is working correctly.
Seeing an authentication or API key error?
Double-check that your .env file contains the correct key. Run cat .env to see its contents. Make sure there are no extra spaces or quotes around the key value.
If you see a "model not found" error, confirm you typed gemini-2.5-flash exactly as shown.
Help me fix my Gemini API key error
✔️ Awesome, I've got everything!
Great. Make sure you have saved your .env file.
ⓧ I'd like to double check the full code
Your .env file should look like this (with your actual key in place of the placeholder):
GEMINI_API_KEY=your-api-key-here
- Finally, open VS Code with your project folder. Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Visual Studio Code and press Enter to open it.
- In VS Code, click File in the top menu bar.
- Click Open Folder.
- Navigate to Desktop, select the rag-knowledge-base folder, and click Open.
You should see the rag-knowledge-base folder name in the VS Code title bar, with .env, requirements.txt, and venv/ visible in the Explorer sidebar.
Your development environment is ready. Next up, you will create the knowledge base documents and learn how to split them into chunks that your RAG system can search.
Create Your Knowledge Base and Understand Chunking
Your environment is ready. Now it's time to give your RAG system something to work with.
LLMs can only answer from their training data, which means they don't know about your specific documents. The solution is to feed your documents into a knowledge base that the LLM can search at query time. But there's a catch: raw documents are too long to stuff into a single prompt. You need to split them into smaller passages called chunks.
In this step, get ready to:
- Create sample markdown documents in a docs/ folder.
- Build a document loader and text chunker in Python.
- Test the chunker to verify it produces overlapping passages.
Create your knowledge base documents
Your knowledge base is just a folder of markdown files. Each file covers one Python concept that the RAG system can later retrieve and cite when answering questions.
- Create a docs/ folder inside your rag-knowledge-base project by running this command:
mkdir docs
- Confirm the folder exists by running:
ls
You should see docs listed alongside your existing files.
Now create the three markdown files that will serve as your knowledge base. In VS Code, create a new file called decorators.md inside the docs/ folder.
- In the VS Code Explorer sidebar, right-click the docs folder.
- Click New File.
- Name it decorators.md and paste the following content:
# Python Decorators
A decorator is a function that takes another function as an argument and extends its behavior without modifying it directly. Decorators use the @symbol syntax placed above a function definition.
## How Decorators Work
When you write @my_decorator above a function, Python passes that function to my_decorator() and replaces it with the return value. This is equivalent to writing: my_function = my_decorator(my_function).
## Common Use Cases
Decorators are used for logging, access control, caching, and measuring execution time. The @property decorator turns a method into a getter. The @staticmethod decorator defines a method that does not receive the instance as first argument.
## Writing a Simple Decorator
To write a decorator, define a function that accepts a function as a parameter, defines an inner wrapper function that calls the original, and returns the wrapper. The wrapper can execute code before and after the original function call.
## The functools.wraps Decorator
When writing decorators, use @functools.wraps(func) on your wrapper function. This preserves the original function's name, docstring, and other attributes. Without it, debugging becomes harder because the function appears to be named "wrapper" instead of its real name.
- Create another new file in docs/ called generators.md and paste:
# Python Generators
A generator is a function that returns an iterator using the yield keyword instead of return. Each time yield is called, the function's state is saved and the yielded value is produced. The next call to next() resumes execution from where it left off.
## Why Use Generators
Generators are memory-efficient because they produce values one at a time rather than storing an entire sequence in memory. This makes them ideal for processing large files, infinite sequences, or streaming data.
## Generator Expressions
Similar to list comprehensions, generator expressions use parentheses instead of brackets: (x*2 for x in range(1000000)). This creates a generator object that produces values lazily without allocating memory for all results at once.
## The yield from Statement
Python 3.3 introduced yield from, which delegates to a sub-generator. Instead of writing a loop that yields each item from another iterable, you can write yield from iterable. This simplifies code and properly propagates exceptions.
## Generator-Based Coroutines
Before async/await syntax, generators were used as coroutines via the send() method. While async/await is now preferred, understanding generator coroutines helps when reading older Python code or working with libraries that still use them.
- Create one more file in docs/ called context_managers.md and paste:
# Python Context Managers
A context manager is an object that defines setup and teardown actions for a with statement. The most common example is file handling: with open('file.txt') as f ensures the file is properly closed even if an exception occurs.
## The Protocol
Context managers implement two methods: __enter__() which runs at the start of the with block and returns the resource, and __exit__() which runs when the block ends (whether normally or via exception). The __exit__ method receives exception info and can suppress exceptions by returning True.
## The contextlib Module
The contextlib module provides utilities for working with context managers. The @contextmanager decorator lets you write a context manager as a generator function with a single yield. Code before yield runs on enter, code after yield runs on exit.
## Common Use Cases
Context managers are used for file handling, database connections, locks in threading, temporary changes to state, and resource cleanup. They guarantee cleanup happens even when exceptions occur, preventing resource leaks.
## Nested Context Managers
You can nest multiple context managers in a single with statement using commas: with open('a') as f1, open('b') as f2. For dynamic numbers of context managers, use contextlib.ExitStack which manages a variable number of context managers.
- Save all three files.
- Confirm all files exist by running:
ls docs/
You should see context_managers.md, decorators.md, and generators.md listed.
Why markdown files?
Markdown is a lightweight text format that's easy to read and parse. Many real-world knowledge bases use markdown because it separates content from formatting with minimal overhead. Your RAG system reads the raw text, so the headings and structure help the chunker produce meaningful passages.
Don't see all three files?
Make sure each file is saved inside the docs/ folder, not in the project root. Check that the file extensions are .md (not .txt or .md.txt).
help me debug my docs folder setup
Build the document loader and chunker
Now you need Python code that reads those markdown files and splits them into small, overlapping passages. This is the first half of the ingestion pipeline.
- In VS Code, create a new file in the project root called ingest.py.
- Add the imports and client initialization by pasting this code:
import os
from google import genai
from google.genai import types
import chromadb
from dotenv import load_dotenv
load_dotenv()
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
What does this code do?
These lines set up everything the ingestion script will need. load_dotenv() reads your .env file so the API key is available as an environment variable.
The genai.Client connects to the Gemini API. You won't use types or chromadb until the next step, but importing them now keeps all imports at the top of the file.
- Below the client initialization, add the load_documents() function:
def load_documents(docs_dir="docs"):
documents = []
for filename in os.listdir(docs_dir):
if filename.endswith(".md"):
filepath = os.path.join(docs_dir, filename)
with open(filepath, "r") as f:
content = f.read()
documents.append({"content": content, "source": filename})
return documents
What does this function do?
This function scans the docs/ folder for any file ending in .md. It reads each file's full content and stores it in a list alongside the filename. Tagging each document with its source filename is what makes source citations possible later.
- Below load_documents(), add the chunk_text() function:
def chunk_text(text, chunk_size=500, overlap=100):
chunks = []
start = 0
while start < len(text):
end = start + chunk_size
chunk = text[start:end]
chunks.append(chunk)
start = end - overlap
return chunks
What does this function do?
This splits a long text into passages of roughly 500 characters each. The key detail is the overlap parameter. Each chunk starts 100 characters before the previous chunk ended, so context at the boundary is never lost.
Without overlap, a sentence split across two chunks might not match any search query. Overlap ensures important ideas that fall at the boundary appear in both neighbouring chunks.
- Save ingest.py.
Seeing a red underline on imports?
VS Code may show import warnings if it's not using the Python interpreter from your virtual environment. Click the Python version in the bottom-left status bar and select the interpreter inside your venv/ folder.
help me fix Python import errors in VS Code
✔️ Awesome, I've got everything!
Great. Double check you've saved ingest.py and move on to testing.
ⓧ I'd like to double check the full code
Here is the full ingest.py so far:
import os
from google import genai
from google.genai import types
import chromadb
from dotenv import load_dotenv
load_dotenv()
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
def load_documents(docs_dir="docs"):
documents = []
for filename in os.listdir(docs_dir):
if filename.endswith(".md"):
filepath = os.path.join(docs_dir, filename)
with open(filepath, "r") as f:
content = f.read()
documents.append({"content": content, "source": filename})
return documents
def chunk_text(text, chunk_size=500, overlap=100):
chunks = []
start = 0
while start < len(text):
end = start + chunk_size
chunk = text[start:end]
chunks.append(chunk)
start = end - overlap
return chunks
Test the chunker
Before moving on, confirm that the loader finds your documents and that the chunker produces overlapping passages.
- Run this test script in your terminal to see the chunker in action:
python3 << 'EOF'
from ingest import load_documents, chunk_text
docs = load_documents()
print(f"Loaded {len(docs)} documents")
for doc in docs:
chunks = chunk_text(doc["content"])
print(f" {doc['source']}: {len(chunks)} chunks")
# Show overlap between adjacent chunks
first_doc_chunks = chunk_text(docs[0]["content"])
print(f"\nEnd of chunk 1: ...{first_doc_chunks[0][-50:]}")
print(f"Start of chunk 2: {first_doc_chunks[1][:50]}...")
EOF
What should I see?
You should see output showing 3 documents loaded, each producing 3-4 chunks. The last two lines demonstrate the overlap: the ending of chunk 1 and the beginning of chunk 2 share the same 100-character passage.
This overlap is how your RAG system avoids losing context at chunk boundaries. A relevant sentence that happens to sit at the edge of one chunk will also appear in the next chunk.
Seeing a ModuleNotFoundError?
Make sure your virtual environment is still activated. Your terminal prompt should start with (venv). If it's missing, re-activate by running source venv/bin/activate from your rag-knowledge-base folder.
help me fix ModuleNotFoundError when testing my chunker
Your knowledge base is loaded, chunked, and ready. Next up, you'll turn those text passages into numerical vectors and store them in a searchable database.
Embed and Store Your Documents
In the last step, you built the foundation of your ingestion pipeline. Your code can load markdown files and split them into overlapping chunks.
Now it is time to turn those text chunks into something a computer can search through. You will convert each chunk into a numerical vector (an embedding) and store it in a vector database for fast similarity lookups.
In this step, get ready to:
- Write a function that calls the Gemini embedding API to convert text into vectors.
- Build a main function that orchestrates the full ingestion pipeline.
- Run the script and verify your documents are stored in ChromaDB.
Set up the Gemini embedding client
An embedding is a list of numbers that captures the meaning of a piece of text. Similar texts produce similar number sequences, which is what makes vector search possible.
The Gemini API provides gemini-embedding-001, a model designed specifically for generating these vectors. It accepts a task_type parameter that optimizes the embedding for your use case.
- In ingest.py, find the chunk_text function. Add the following get_embeddings function below it:
def get_embeddings(texts):
result = client.models.embed_content(
model="gemini-embedding-001",
contents=texts,
config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT"),
)
return [e.values for e in result.embeddings]
What does this code do?
- client.models.embed_content() calls the Gemini embedding API with your list of text chunks.
- model="gemini-embedding-001" specifies the embedding model, which outputs a 3072-dimension vector for each input.
- task_type="RETRIEVAL_DOCUMENT" tells the model these texts are documents being stored for later retrieval. When you query later, you will use RETRIEVAL_QUERY instead.
- The return line extracts the raw number arrays from each embedding result.
Why two different task types?
Using RETRIEVAL_DOCUMENT for stored content and RETRIEVAL_QUERY for questions tells the model to optimize the vectors for asymmetric search. A short question like "What is a decorator?" and a long paragraph explaining decorators will land close together in vector space, even though they look very different as text.
- Save ingest.py.
- You will not see a change yet. The next chunk builds the main() function that calls get_embeddings.
Embed all chunks and store in ChromaDB
ChromaDB is a vector database that stores embeddings alongside the original text and metadata. It handles persistence to disk so your data survives between script runs.
- In ingest.py, add the following main() function below get_embeddings:
def main():
print("Loading documents...")
documents = load_documents()
print(f"Found {len(documents)} documents.")
all_chunks = []
all_sources = []
all_ids = []
for doc in documents:
chunks = chunk_text(doc["content"])
for i, chunk in enumerate(chunks):
all_chunks.append(chunk)
all_sources.append(doc["source"])
all_ids.append(f"{doc['source']}_chunk_{i}")
print(f"Created {len(all_chunks)} chunks.")
print("Generating embeddings...")
embeddings = get_embeddings(all_chunks)
print(f"Generated {len(embeddings)} embeddings.")
print("Storing in ChromaDB...")
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(name="knowledge_base")
collection.add(
ids=all_ids,
documents=all_chunks,
embeddings=embeddings,
metadatas=[{"source": source} for source in all_sources],
)
print(f"Stored {collection.count()} chunks in ChromaDB.")
print("Ingestion complete!")
if __name__ == "__main__":
main()
What does this code do?
The main() function orchestrates the entire ingestion pipeline in sequence:
- It loads all markdown files using load_documents(), then chunks each document and tracks the source filename for every chunk.
- Each chunk gets a unique ID like decorators.md_chunk_0 so ChromaDB can identify individual records.
- chromadb.PersistentClient(path="./chroma_db") creates a database that saves to disk in a chroma_db folder.
- get_or_create_collection(name="knowledge_base") creates a named collection (or opens an existing one).
- collection.add() stores the IDs, raw text, embeddings, and metadata together in one call.
The if __name__ == "__main__" guard ensures main() only runs when you execute the file directly.
- Save ingest.py.
✔️ Awesome, I've got everything!
Great. Double check you have saved your file before moving on.
ⓧ I'd like to double check the full code
Here is the complete ingest.py file at this point:
import os
from google import genai
from google.genai import types
import chromadb
from dotenv import load_dotenv
load_dotenv()
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
def load_documents(docs_dir="docs"):
documents = []
for filename in os.listdir(docs_dir):
if filename.endswith(".md"):
filepath = os.path.join(docs_dir, filename)
with open(filepath, "r") as f:
content = f.read()
documents.append({"content": content, "source": filename})
return documents
def chunk_text(text, chunk_size=500, overlap=100):
chunks = []
start = 0
while start < len(text):
end = start + chunk_size
chunk = text[start:end]
chunks.append(chunk)
start = end - overlap
return chunks
def get_embeddings(texts):
result = client.models.embed_content(
model="gemini-embedding-001",
contents=texts,
config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT"),
)
return [e.values for e in result.embeddings]
def main():
print("Loading documents...")
documents = load_documents()
print(f"Found {len(documents)} documents.")
all_chunks = []
all_sources = []
all_ids = []
for doc in documents:
chunks = chunk_text(doc["content"])
for i, chunk in enumerate(chunks):
all_chunks.append(chunk)
all_sources.append(doc["source"])
all_ids.append(f"{doc['source']}_chunk_{i}")
print(f"Created {len(all_chunks)} chunks.")
print("Generating embeddings...")
embeddings = get_embeddings(all_chunks)
print(f"Generated {len(embeddings)} embeddings.")
print("Storing in ChromaDB...")
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(name="knowledge_base")
collection.add(
ids=all_ids,
documents=all_chunks,
embeddings=embeddings,
metadatas=[{"source": source} for source in all_sources],
)
print(f"Stored {collection.count()} chunks in ChromaDB.")
print("Ingestion complete!")
if __name__ == "__main__":
main()
Verify the vector store
Time to run the full pipeline and confirm everything landed in ChromaDB.
- Run the ingestion script by entering this command in your terminal:
python3 ingest.py
You should see output similar to this:
Loading documents...
Found 3 documents.
Created 9 chunks.
Generating embeddings...
Generated 9 embeddings.
Storing in ChromaDB...
Stored 9 chunks in ChromaDB.
Ingestion complete!
What should I see?
The exact chunk count depends on document length, but you should see 3 documents found and the same number for chunks, embeddings, and stored items. The numbers confirm that every chunk made it through the full pipeline: load, chunk, embed, store.
- Confirm the data is queryable by running a quick peek at your collection:
python3 -c "import chromadb; c = chromadb.PersistentClient(path='./chroma_db'); col = c.get_collection('knowledge_base'); print(col.peek(limit=2))"
You should see a dictionary containing the first two stored chunks, their IDs, metadata, and embeddings.
Seeing an error when running the script?
- If you see ModuleNotFoundError: make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt.
- If you see an API key error: check that your .env file has GEMINI_API_KEY set to a valid key from Google AI Studio.
- If the chunk count is 0: confirm your docs/ folder contains .md files and is in the same directory as ingest.py.
Help me debug my ingest.py script
Your documents are now embedded and stored as searchable vectors. Next up, you will build the query pipeline that takes a question, finds the most relevant chunks, and generates a cited answer using Gemini.
Build the Query Pipeline
In the previous step, you embedded all your document chunks and stored them in ChromaDB. Your knowledge base is loaded and ready to be searched.
Now it is time to build the other side of the pipeline. You will take a user's question, convert it into an embedding, find the most relevant chunks via vector similarity, and feed those chunks to Gemini so it can generate a grounded, cited answer.
In this step, get ready to:
- Embed user questions and retrieve the most similar chunks from ChromaDB.
- Generate cited answers using Gemini with retrieved context.
- Build an interactive Q&A loop and test the full RAG pipeline end-to-end.
Write the retrieval function
Your retrieval function needs to do two things: turn a user's question into an embedding, then search ChromaDB for the chunks closest to that embedding.
- Create a new file called query.py in your project folder by running:
touch query.py
- Open query.py in VS Code and add the following imports, client setup, and retrieval functions:
import os
from google import genai
from google.genai import types
import chromadb
from dotenv import load_dotenv
load_dotenv()
# Initialize the Gemini client with your API key
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
def get_query_embedding(query):
# Embed the user's question using RETRIEVAL_QUERY task type
result = client.models.embed_content(
model="gemini-embedding-001",
contents=[query],
config=types.EmbedContentConfig(task_type="RETRIEVAL_QUERY"),
)
return result.embeddings[0].values
def retrieve_relevant_chunks(query, n_results=3):
# Connect to the existing ChromaDB vector store
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_collection(name="knowledge_base")
# Embed the question and search for similar chunks
query_embedding = get_query_embedding(query)
results = collection.query(
query_embeddings=[query_embedding],
n_results=n_results,
)
# Package each result with its text, source file, and distance score
chunks = []
for i in range(len(results["documents"][0])):
chunks.append(
{
"text": results["documents"][0][i],
"source": results["metadatas"][0][i]["source"],
"distance": results["distances"][0][i],
}
)
return chunks
What does this code do?
get_query_embedding() embeds the user's question using RETRIEVAL_QUERY as the task type. This tells the embedding model the text is a search query, which optimizes the embedding for finding relevant documents.
retrieve_relevant_chunks() connects to your existing ChromaDB store, runs a similarity search with the question's embedding, and returns the top 3 closest chunks. Each chunk includes the text, which file it came from, and its distance score (lower means more similar).
- Save query.py.
- You won't see output yet. The next function will use these retrieved chunks to generate an answer.
Write the generation function
The generation function takes the retrieved chunks, formats them as context, and asks Gemini to answer the question using only that context. This is the core of RAG: grounding the LLM's answer in your actual data.
- In query.py, add the following function below retrieve_relevant_chunks:
def generate_answer(query, chunks):
# Format each retrieved chunk with its source label
context = ""
for i, chunk in enumerate(chunks, 1):
context += f"\n[Source {i}: {chunk['source']}]\n{chunk['text']}\n"
# Instruct Gemini to answer ONLY from the context and cite sources
system_instruction = (
"You are a helpful assistant that answers questions based ONLY on the "
"provided context. If the context does not contain enough information "
"to answer the question, say 'I don't have enough information to answer "
"that.' Always cite which source(s) you used in your answer."
)
prompt = f"Context:\n{context}\n\nQuestion: {query}"
# Call Gemini to generate the answer
interaction = client.interactions.create(
model="gemini-2.5-flash",
system_instruction=system_instruction,
input=prompt,
generation_config={"thinking_level": "low"},
)
return interaction.output_text
What does this code do?
The function builds a context string by labeling each chunk with its source filename (e.g. [Source 1: decorators.md]). This gives Gemini something concrete to cite in its answer.
The system_instruction is key to RAG. It constrains Gemini to answer only from the provided context and always cite sources. Without it, the model might hallucinate answers from its training data instead of your documents.
Setting thinking_level to "low" reduces thinking token usage, which keeps your free-tier usage efficient and responses faster.
- Save query.py.
- You won't see output yet. The next chunk wires everything into an interactive loop you can run.
Wire it together into a query loop
The final piece is a loop that accepts questions from the user, retrieves relevant chunks, and prints the generated answer along with which sources were used.
- In query.py, add the following main() function and entry point at the bottom of the file:
def main():
print("RAG Knowledge Base Q&A")
print("Type 'quit' to exit.\n")
while True:
# Wait for the user's question
query = input("Your question: ")
if query.lower() == "quit":
break
# Retrieve the top 3 most relevant chunks
chunks = retrieve_relevant_chunks(query)
# Show which chunks were retrieved and their distances
print("\nRetrieved chunks:")
for i, chunk in enumerate(chunks, 1):
print(f" [{i}] {chunk['source']} (distance: {chunk['distance']:.4f})")
# Generate and display the answer
answer = generate_answer(query, chunks)
print(f"\nAnswer: {answer}\n")
if __name__ == "__main__":
main()
What does this code do?
The main() function creates a simple input loop. It reads a question, retrieves the most relevant chunks, prints each chunk's source filename and distance score, then generates and prints the answer.
The distance scores help you see how relevant each chunk is. Lower distances mean the chunk is more semantically similar to your question.
- Save query.py.
- Run the query pipeline by executing:
python3 query.py
- Try asking a question that IS in your documents, such as:
What is a decorator in Python?
- You should see retrieved chunks from decorators.md with low distance scores, followed by a cited answer.
- Now try asking a question that is NOT in your documents, such as:
How do I set up a Django web server?
- The system should respond that it does not have enough information, because your knowledge base only contains Python concept docs.
- Type quit to exit the loop when you are done testing.
Not getting answers?
If you see an API error, check that your .env file still contains a valid GEMINI_API_KEY and that your virtual environment is activated.
If you see a collection error, make sure you ran ingest.py successfully in the previous step. The chroma_db folder must exist with your stored embeddings.
help me debug my query.py script
✔️ Awesome, I've got everything!
Great. Double check that you have saved query.py and that both in-knowledge-base and out-of-knowledge-base questions produce the expected responses.
ⓧ I'd like to double check the full code
import os
from google import genai
from google.genai import types
import chromadb
from dotenv import load_dotenv
load_dotenv()
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
def get_query_embedding(query):
result = client.models.embed_content(
model="gemini-embedding-001",
contents=[query],
config=types.EmbedContentConfig(task_type="RETRIEVAL_QUERY"),
)
return result.embeddings[0].values
def retrieve_relevant_chunks(query, n_results=3):
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_collection(name="knowledge_base")
query_embedding = get_query_embedding(query)
results = collection.query(
query_embeddings=[query_embedding],
n_results=n_results,
)
chunks = []
for i in range(len(results["documents"][0])):
chunks.append(
{
"text": results["documents"][0][i],
"source": results["metadatas"][0][i]["source"],
"distance": results["distances"][0][i],
}
)
return chunks
def generate_answer(query, chunks):
context = ""
for i, chunk in enumerate(chunks, 1):
context += f"\n[Source {i}: {chunk['source']}]\n{chunk['text']}\n"
system_instruction = (
"You are a helpful assistant that answers questions based ONLY on the "
"provided context. If the context does not contain enough information "
"to answer the question, say 'I don't have enough information to answer "
"that.' Always cite which source(s) you used in your answer."
)
prompt = f"Context:\n{context}\n\nQuestion: {query}"
interaction = client.interactions.create(
model="gemini-2.5-flash",
system_instruction=system_instruction,
input=prompt,
generation_config={"thinking_level": "low"},
)
return interaction.output_text
def main():
print("RAG Knowledge Base Q&A")
print("Type 'quit' to exit.\n")
while True:
query = input("Your question: ")
if query.lower() == "quit":
break
chunks = retrieve_relevant_chunks(query)
print("\nRetrieved chunks:")
for i, chunk in enumerate(chunks, 1):
print(f" [{i}] {chunk['source']} (distance: {chunk['distance']:.4f})")
answer = generate_answer(query, chunks)
print(f"\nAnswer: {answer}\n")
if __name__ == "__main__":
main()
Secret mission
Add Relevance Scoring
Your RAG system always returns an answer, even when the question has nothing to do with your knowledge base. Can you add a relevance threshold so it knows when to say "I don't know"?
Clean Up Your Resources
Clean Up Your Resources
This project runs entirely locally and uses the Gemini API free tier, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.
Resources you used:
- Project folder (rag-knowledge-base) containing all Python scripts, markdown docs, and configuration files.
- Virtual environment (venv/) inside your project folder, taking approximately 200-400MB of disk space.
- ChromaDB database (chroma_db/) inside your project folder, taking a few MB of disk space.
- A .env file containing your Gemini API key.
Keep everything running
No action needed. Choose this if you are still actively building or experimenting with your RAG knowledge base.
- Nothing runs in the background. Your project folder just sits on disk.
- Your API key is stored in .env but the Gemini free tier has no charges, so there is no financial risk.
- You can return any time and run python3 query.py to query your knowledge base again.
Pause - I'll come back to this later
Free up terminal resources while keeping all your files intact for later.
- If your virtual environment is still active, deactivate it by running:
deactivate
- Close the Terminal window. Your project files, virtual environment, and ChromaDB data will all remain on disk.
- When you come back, navigate to your project folder, activate the virtual environment with source venv/bin/activate, and pick up where you left off.
Delete - I don't want to use this again
Remove all project resources and start fresh. Everything is self-contained in one folder, so a single delete removes it all.
- If your virtual environment is still active, deactivate it first by running:
deactivate
- Delete the entire project folder by running:
rm -rf ~/Desktop/rag-knowledge-base
- This removes all Python scripts, the virtual environment, the ChromaDB database, your .env file, and the sample docs.
- Your Google AI Studio API key still exists in your Google account. If you want to revoke it, visit aistudio.google.com, open the API keys section, and delete the key you created for this project.
Nice Work!
Nice Work!
You just built a complete RAG pipeline from scratch. Your system ingests documents, chunks and embeds them into a vector store, retrieves the most relevant passages for any question, and generates cited answers using Gemini.
Here's what you accomplished:
- Built a full RAG pipeline that loads markdown documents, splits them into overlapping chunks, generates embeddings with Gemini, and stores them in a ChromaDB vector database.
- Created a query system that uses vector similarity search to find relevant passages and generates grounded answers with source citations using the Gemini API.
- Engineered prompts with context injection that force the LLM to answer only from retrieved documents and cite its sources.
- Secret Mission: Added relevance scoring with a configurable threshold so the system refuses to answer when no chunk is relevant enough, preventing hallucinated responses.
Ready to quiz yourself?