How to Make an AI Agent Understand an Entire Codebase

The Problem: AI Agents That Don't Know Your Code

You ask an AI agent to refactor a function, but it suggests changes that break three other files. You ask it to find a bug, and it hallucinates a file that doesn't exist. The root cause is simple: the agent never read your codebase. It only knows what you paste into the prompt.

A production-grade agent needs to retrieve relevant code, understand the project structure, and ground its responses in the actual repository. In this guide, you will build a minimal but complete system that ingests a codebase into a vector store, retrieves relevant snippets on demand, and feeds them into an LLM to answer questions or generate code.

You will use Python, the OpenAI API for embeddings and completions, and a local vector store (Chroma). The final result is a script you can run against any repository to get a working codebase-aware agent.

Before You Start

You need Python 3.9 or later, a working git installation, and an OpenAI API key with access to the text-embedding-3-small and gpt-4o-mini models. Set your API key as an environment variable to keep it out of your scripts.

Create a project directory and install the required packages. All commands below assume you are in a Unix-like shell.

mkdir codebase-agent
cd codebase-agent
python -m venv venv
source venv/bin/activate
pip install openai chromadb tiktoken gitpython

Step 1: Clone a Repository to Test With

Use a real, non-trivial repository to test the agent. A good choice is a small but structured library like requests or Flask. For this tutorial, clone the popular requests library.

git clone https://github.com/psf/requests.git sample-repo
cd sample-repo

• Use a repository with multiple files and directories to make retrieval meaningful.
• Avoid huge monorepos on your first run; start with a few hundred files.

Step 2: Write the Codebase Indexer

The first script walks the repository, extracts text from code files, splits it into chunks, embeds each chunk, and stores the vectors in Chroma. This is the foundation of the agent's understanding.

Create a file named index_codebase.py in the project root (outside the sample-repo directory). The script uses a simple file-walking function, a token-based splitter, and the OpenAI embedding API.

import os
import hashlib
from pathlib import Path
import chromadb
from openai import OpenAI
import tiktoken

client = OpenAI()
enc = tiktoken.encoding_for_model("text-embedding-3-small")

# File extensions to index
CODE_EXTENSIONS = {'.py', '.js', '.ts', '.jsx', '.tsx', '.go', '.rs', '.java', '.c', '.cpp', '.h', '.hpp', '.rb', '.php', '.swift', '.kt', '.scala', '.md', '.txt', '.yaml', '.yml', '.json'}

def split_text(text, max_tokens=500):
    tokens = enc.encode(text)
    chunks = []
    for i in range(0, len(tokens), max_tokens):
        chunk = tokens[i:i+max_tokens]
        chunks.append(enc.decode(chunk))
    return chunks

def index_repo(repo_path, collection):
    for root, dirs, files in os.walk(repo_path):
        # Skip hidden and build directories
        dirs[:] = [d for d in dirs if not d.startswith('.') and d not in {'node_modules', 'venv', 'dist', 'build', '__pycache__'}]
        for file in files:
            path = Path(root) / file
            if path.suffix in CODE_EXTENSIONS:
                try:
                    content = path.read_text(encoding='utf-8', errors='ignore')
                except Exception:
                    continue
                if not content.strip():
                    continue
                chunks = split_text(content)
                for i, chunk in enumerate(chunks):
                    # Generate a unique ID from path and chunk index
                    chunk_id = hashlib.md5(f"{path}:{i}".encode()).hexdigest()
                    # Embed the chunk
                    response = client.embeddings.create(
                        model="text-embedding-3-small",
                        input=chunk
                    )
                    embedding = response.data[0].embedding
                    metadata = {"path": str(path), "chunk_index": i}
                    collection.add(
                        ids=[chunk_id],
                        embeddings=[embedding],
                        documents=[chunk],
                        metadatas=[metadata]
                    )

def main():
    chroma_client = chromadb.PersistentClient(path="./chroma_db")
    collection = chroma_client.get_or_create_collection(name="codebase")
    repo_path = "./sample-repo"
    if not os.path.exists(repo_path):
        print("Repository not found. Run the clone step first.")
        return
    index_repo(repo_path, collection)
    print(f"Indexed {collection.count()} chunks")

if __name__ == "__main__":
    main()

Step 3: Run the Indexer

Run the script from the codebase-agent directory. It will take a minute or two depending on the repository size and API latency. The output shows the number of chunks stored.

python index_codebase.py
# Expected output: Indexed 1234 chunks (number varies)

• If you get a rate limit error, add a small sleep between API calls.
• The Chroma database is stored in the ./chroma_db directory; keep it between runs.

Step 4: Build the Retrieval-Augmented Agent

Now create a second script, agent.py, that takes a natural language question, retrieves the most relevant chunks from the vector store, and sends them along with the question to the LLM. The LLM then answers based on the retrieved context.

The retrieval function uses Chroma's query method with the embedded question. We fetch the top 5 chunks and join them into a context string.

import os
import chromadb
from openai import OpenAI

client = OpenAI()

def retrieve(query, collection, n_results=5):
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=query
    )
    query_embedding = response.data[0].embedding
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=n_results
    )
    return results['documents'][0], results['metadatas'][0]

def ask_agent(question, collection):
    docs, metas = retrieve(question, collection)
    context = "\n\n---\n\n".join(docs)
    system_prompt = "You are a senior developer. Answer the question using only the provided code context. If the context is insufficient, say so."
    user_prompt = f"Context:\n{context}\n\nQuestion: {question}"
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    )
    return completion.choices[0].message.content

def main():
    chroma_client = chromadb.PersistentClient(path="./chroma_db")
    collection = chroma_client.get_or_create_collection(name="codebase")
    question = input("Ask about the codebase: ")
    answer = ask_agent(question, collection)
    print("\nAnswer:", answer)

if __name__ == "__main__":
    main()

Step 5: Test the Agent

Run the agent and ask a question about the requests library, such as 'How does the Session class manage cookies?' The agent should produce an answer grounded in the actual source code.

python agent.py
# Ask: How does the Session class manage cookies?
# The answer will reference specific files and functions from the repository.

• The quality of answers depends on chunk size and retrieval count. Tweak the max_tokens and n_results parameters.
• If answers are too generic, increase n_results or use a smaller chunk size.

Verify It Worked

A quick sanity check: ask a question that requires specific file knowledge, like 'Where is the retry logic defined?' The agent should point to requests/adapters.py or similar.

If the agent gives a vague answer, inspect the retrieved chunks by printing them. You can add a debug flag to see which files were used.

# Add to agent.py for debugging
print("Retrieved from:", [m['path'] for m in metas])

Troubleshooting

If the script fails, check the most common issues: missing API key, incorrect model access, or a stale Chroma database. Clear the chroma_db directory and re-index if you change the codebase.

• Error 'openai.AuthenticationError' means your API key is not set or invalid.
• Error 'ModuleNotFoundError' means you missed a pip install.
• If the vector store is empty, run the indexer again.

Scaling to Larger Codebases

For a real project with thousands of files, the simple walk-and-embed approach becomes slow and expensive. Use incremental indexing (only embed changed files), store embeddings in a dedicated database like Pinecone or Weaviate, and use a more efficient chunking strategy that respects code structure.

A common improvement is to use a code-aware splitter that breaks on function and class definitions. Libraries like tree-sitter can parse code into ASTs for precise chunking.

# Example using tree-sitter for Python (simplified)
from tree_sitter import Language, Parser

PY_LANGUAGE = Language('build/my-languages.so', 'python')
parser = Parser()
parser.set_language(PY_LANGUAGE)

tree = parser.parse(bytes(source, 'utf8'))
# Traverse AST to extract function/class nodes and use their text as chunks

What I Would Do in Production

For a production agent, I would not embed every file at once. Instead, I would run a background job that watches the git history and updates the index on each commit. I would also add a permission layer so the agent can only read files, and a logging system to trace which chunks were retrieved for each answer.

Here is a recommended Docker Compose setup that runs the indexer as a service and the agent as a separate container.

version: '3'
services:
  indexer:
    build: .
    command: python index_codebase.py
    volumes:
      - .:/app
      - ./chroma_db:/app/chroma_db
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
  agent:
    build: .
    command: python agent.py
    stdin_open: true
    tty: true
    volumes:
      - .:/app
      - ./chroma_db:/app/chroma_db
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}

FAQ

• Q: How much does this cost to run? A: Embedding a small repo costs a few cents. Each query costs less than a cent with gpt-4o-mini.
• Q: Can I use a local LLM instead of OpenAI? A: Yes, replace the client with something like Ollama or llama.cpp, but you must also use a local embedding model.
• Q: What if the codebase has multiple languages? A: The indexer already handles common extensions; ensure you include the ones you need.
• Q: How do I keep the index up to date? A: Run the indexer after every pull, or implement a git hook that triggers re-indexing.

Next Steps

Your agent now understands your codebase. The next step is to wire it into your editor or CI pipeline. Start by running the agent on your own repository and asking it to explain a module you wrote months ago.

If you want a quick way to test retrieval without building the whole pipeline, use the code.live JSON Path Tester to validate any JSON responses you get from your agent's API calls.

python agent.py