Building a RAG System for Internal Codebases

Building a RAG System for Internal Codebases

Retrieval-Augmented Generation (RAG) has become one of the most practical applications of large language models in enterprise environments. When applied to internal codebases, a RAG system allows developers to ask natural-language questions about proprietary code, internal libraries, and undocumented APIs — and get accurate, context-aware answers grounded in the actual source code. This tutorial walks through everything you need to build, deploy, and optimize a RAG system tailored for code.

What Is a Code RAG System?

A code RAG system combines a retrieval mechanism with a generative language model. Instead of relying solely on the model's pre-trained knowledge, the system first retrieves relevant code snippets, documentation, or metadata from your internal repositories, then feeds that context to the LLM to generate an answer. The result is a tool that can explain internal functions, suggest how to use internal libraries, trace dependencies, and even help onboard new engineers.

The core pipeline has four stages: ingestion, embedding, retrieval, and generation. During ingestion, source files are parsed and chunked. During embedding, each chunk is converted into a vector representation. During retrieval, a user query is matched against those vectors to find the most relevant chunks. During generation, the retrieved chunks are inserted into the LLM prompt as context.

Why It Matters

• Onboarding acceleration: New hires can ask questions about the codebase without interrupting senior engineers.
• Reduced hallucination: Grounding the LLM in actual source code dramatically reduces fabricated answers.
• Cross-repository discovery: Developers can find how internal libraries are used across many services in one query.
• Preserved institutional knowledge: Implicit knowledge locked in code becomes queryable.
• Security and privacy: Unlike sending code to public chatbots, an internal RAG system keeps proprietary code within your infrastructure.

Architecture Overview

Before writing code, it helps to understand the components you will build. The system consists of an indexer that walks your repositories, a vector database that stores embeddings, a retriever that performs similarity search, and an LLM client that produces final answers. You will also need a chunking strategy specific to code, since naive character-based chunking often splits functions and classes in unhelpful ways.

Setting Up the Environment

We will use Python with several popular libraries: tree-sitter for code parsing, sentence-transformers for embeddings, chromadb as the vector store, and the openai client for generation. Install them with pip:

pip install tree-sitter tree-sitter-languages sentence-transformers chromadb openai tiktoken

You will also need an OpenAI API key, or you can substitute any compatible LLM endpoint. Set the key as an environment variable:

export OPENAI_API_KEY="sk-your-key-here"

Chunking Code with Tree-Sitter

The most important design decision in a code RAG system is how you chunk files. Splitting by fixed character counts breaks functions mid-body and destroys semantic coherence. A better approach is to parse the abstract syntax tree (AST) and chunk by meaningful units: functions, classes, and methods. Tree-sitter supports many languages and gives you reliable AST nodes.

from tree_sitter_languages import get_parser
from tree_sitter import Node

def extract_code_chunks(source_code: str, language: str, file_path: str):
    parser = get_parser(language)
    tree = parser.parse(bytes(source_code, "utf8"))
    root = tree.root_node

    chunks = []
    # Define node types that represent meaningful code units per language
    target_types = {
        "python": ["function_definition", "class_definition"],
        "javascript": ["function_declaration", "class_declaration", "method_definition"],
        "typescript": ["function_declaration", "class_declaration", "method_definition"],
        "java": ["method_declaration", "class_declaration"],
        "go": ["function_declaration", "method_declaration", "type_declaration"],
    }.get(language, ["function_definition", "class_definition"])

    def walk(node: Node):
        if node.type in target_types:
            snippet = source_code[node.start_byte:node.end_byte]
            name_node = node.child_by_field_name("name")
            name = source_code[name_node.start_byte:name_node.end_byte] if name_node else "anonymous"
            chunks.append({
                "file_path": file_path,
                "name": name,
                "type": node.type,
                "start_line": node.start_point[0] + 1,
                "end_line": node.end_point[0] + 1,
                "code": snippet,
            })
        for child in node.children:
            walk(child)

    walk(root)
    return chunks

This function returns a list of dictionaries, each containing the code snippet plus metadata like file path, symbol name, type, and line range. That metadata is critical for retrieval quality because it lets you filter and cite sources in the final answer.

Building the Indexer

The indexer walks a repository directory, identifies source files by extension, parses each file, and collects all chunks. Here is a complete indexer that supports multiple languages:

import os
from pathlib import Path

LANGUAGE_MAP = {
    ".py": "python",
    ".js": "javascript",
    ".ts": "typescript",
    ".java": "java",
    ".go": "go",
}

def index_repository(repo_path: str):
    all_chunks = []
    repo = Path(repo_path)

    for file_path in repo.rglob("*"):
        if not file_path.is_file():
            continue
        ext = file_path.suffix
        if ext not in LANGUAGE_MAP:
            continue
        # Skip common ignored directories
        if any(part in {"node_modules", ".git", "vendor", "__pycache__", "dist", "build"} 
               for part in file_path.parts):
            continue

        language = LANGUAGE_MAP[ext]
        try:
            source = file_path.read_text(encoding="utf-8", errors="ignore")
            chunks = extract_code_chunks(source, language, str(file_path.relative_to(repo)))
            all_chunks.extend(chunks)
        except Exception as e:
            print(f"Failed to parse {file_path}: {e}")

    return all_chunks

Run the indexer on a sample repository to see the output:

chunks = index_repository("./my-service")
print(f"Extracted {len(chunks)} chunks")
for c in chunks[:3]:
    print(f"{c['file_path']}::{c['name']} (lines {c['start_line']}-{c['end_line']})")

Embedding and Storing in a Vector Database

Once you have chunks, you need to embed them and store the vectors. ChromaDB is a lightweight, local-first vector database that works well for development. For each chunk, we create a rich text representation that includes metadata so the embedding captures both the code and its context.

import chromadb
from sentence_transformers import SentenceTransformer

def build_text_for_embedding(chunk: dict) -> str:
    return (
        f"File: {chunk['file_path']}\n"
        f"Symbol: {chunk['name']}\n"
        f"Type: {chunk['type']}\n"
        f"Lines: {chunk['start_line']}-{chunk['end_line']}\n"
        f"\n{chunk['code']}\n"
    )

def create_vector_store(chunks, collection_name="codebase", persist_dir="./chroma_db"):
    client = chromadb.PersistentClient(path=persist_dir)
    collection = client.get_or_create_collection(
        name=collection_name,
        metadata={"hnsw:space": "cosine"}
    )

    embed_model = SentenceTransformer("all-MiniLM-L6-v2")
    texts = [build_text_for_embedding(c) for c in chunks]
    embeddings = embed_model.encode(texts, show_progress_bar=True)

    collection.add(
        ids=[f"chunk_{i}" for i in range(len(chunks))],
        embeddings=embeddings.tolist(),
        documents=texts,
        metadatas=[{
            "file_path": c["file_path"],
            "name": c["name"],
            "type": c["type"],
            "start_line": c["start_line"],
            "end_line": c["end_line"],
        } for c in chunks]
    )
    return collection

The all-MiniLM-L6-v2 model is fast and works reasonably well for code, but for production you may want a code-specific embedding model such as microsoft/codebert-base or a fine-tuned variant. The trade-off is always between latency, cost, and retrieval quality.

Retrieving Relevant Chunks

Retrieval converts the user's question into a vector and finds the most similar chunks. A good retriever also includes metadata filtering and re-ranking. Here is a basic retriever:

def retrieve(collection, query: str, top_k: int = 5, file_filter: str = None):
    embed_model = SentenceTransformer("all-MiniLM-L6-v2")
    query_embedding = embed_model.encode([query]).tolist()

    where_clause = None
    if file_filter:
        where_clause = {"file_path": {"$contains": file_filter}}

    results = collection.query(
        query_embeddings=query_embedding,
        n_results=top_k,
        where=where_clause
    )
    return results

For better results, consider hybrid retrieval: combine vector similarity with keyword search (BM25). Code often contains identifiers and tokens that exact-match search handles better than dense embeddings. Libraries like rank_bm25 can be combined with vector scores for a weighted fusion.

Generating Answers with Context

The final step sends the retrieved chunks to the LLM along with a carefully constructed prompt. The prompt should instruct the model to use only the provided context, cite file paths and line numbers, and admit when it does not know the answer.

import os
from openai import OpenAI

client = OpenAI()

SYSTEM_PROMPT = """You are a senior software engineer helping developers understand an internal codebase.
Answer questions using ONLY the provided code context.
When referencing code, cite the file path and line numbers.
If the context does not contain enough information, say you don't know.
Be concise but complete."""

def generate_answer(query: str, retrieved_results: dict) -> str:
    context_parts = []
    for doc, meta in zip(retrieved_results["documents"][0], retrieved_results["metadatas"][0]):
        context_parts.append(
            f"--- {meta['file_path']} ({meta['name']}, lines {meta['start_line']}-{meta['end_line']}) ---\n{doc}"
        )
    context = "\n\n".join(context_parts)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
        ],
        temperature=0.1,
        max_tokens=1024,
    )
    return response.choices[0].message.content

Putting It All Together

Now combine every component into a single query function and test it end to end:

def ask_codebase(query: str, collection, file_filter: str = None) -> str:
    results = retrieve(collection, query, top_k=5, file_filter=file_filter)
    answer = generate_answer(query, results)
    return answer

# Full pipeline
if __name__ == "__main__":
    # Step 1: Index
    chunks = index_repository("./my-service")
    print(f"Indexed {len(chunks)} code chunks")

    # Step 2: Store
    collection = create_vector_store(chunks)

    # Step 3: Query
    question = "How does the authentication middleware validate JWT tokens?"
    answer = ask_codebase(question, collection)
    print(answer)

Best Practices

Chunk Granularity

Aim for chunks between 50 and 300 lines of code. Too small and you lose context; too large and retrieval precision drops. For very large functions, consider splitting at logical boundaries such as nested blocks, and always include the function signature in each sub-chunk so the model knows what it is looking at.

Include Surrounding Context

When you retrieve a function, also retrieve its imports, class definition, and any docstrings in the same file. A simple way to do this is to store a "file summary" chunk for every file and always include it when any chunk from that file is retrieved.

Keep the Index Fresh

Code changes constantly. Build a CI pipeline that re-indexes on every merge to the main branch, or use a file-watcher service that incrementally updates the vector store. Track file hashes so you only re-embed changed files.

Add Metadata Filters

Let users filter by repository, language, file path, or symbol type. This dramatically improves precision. For example, a query about "database connection pool" can be scoped to .go files in the infra repository.

Handle Multi-File Reasoning

Some questions require understanding relationships across files. Consider building a dependency graph alongside the vector index. When a user asks about a function, you can retrieve not just that function but also its callers and callees by traversing the graph.

Evaluate Retrieval Quality

Build a small evaluation set of question-answer pairs with known relevant files. Measure recall@k and precision@k for your retriever. Iterate on chunking strategy and embedding models using these metrics rather than subjective judgment alone.

Security Considerations

• Run the LLM and embedding model on private infrastructure if your code is highly sensitive.
• Implement access control so users only retrieve code from repositories they are authorized to see.
• Log all queries for audit purposes.
• Avoid embedding secrets, API keys, or credentials — strip them during ingestion.

Cost and Latency Optimization

For large codebases, embedding millions of chunks can be expensive. Use a smaller, faster embedding model for initial retrieval and a larger model for re-ranking the top candidates. Cache frequent queries. Consider quantized embeddings to reduce storage and memory.

Conclusion

Building a RAG system for internal codebases is one of the highest-ROI investments a engineering team can make. By combining tree-sitter-based chunking, a vector database, and a grounded LLM prompt, you create a tool that turns scattered source code into a queryable knowledge base. Start simple with the pipeline in this tutorial, measure retrieval quality on real questions from your team, and iterate on chunking, metadata, and hybrid retrieval until the answers are consistently useful. Over time, add dependency graphs, freshness pipelines, and access controls to scale the system across your entire engineering organization.