Back to Blogs

How I Built CodeRAG with Dependency Graph Using Tree-Sitter

Shivam Sahu•

Recent developments in the AI field are quite impressive, which led me to build a cool project for learning about these advancements. I chose to build CodeRAG with a dependency graph using Tree-sitter.

The concept of CodeRAG is extensively being used in many developer tools like Cursor, Windsurf, and Copilot. However, the implementation of these flows are less available as user projects, so I decided to build a CodeRAG system to chat with your codebase via LLMs by learning from research papers, blogs, and other resources. One research paper I found particularly helpful during this development was CAST: Enhancing Code Retrieval-Augmented Generation with Structural Chunking via Abstract Syntax Tree.

The Biggest Challenges I Faced While Building a RAG System for Codebases

Code files are fundamentally different from general files like PDFs and text documents. A random chunking strategy (fixed size) will break the meaning of all chunks.

Code elements like functions, classes, and code blocks have dependencies, and we can’t understand any function completely until we have knowledge about the classes called inside this function.

Imports in codebase files also create inter-file dependencies, and to understand the use of those imports, LLMs need the context of those imports.

How I Tackled Those Challenges:

  1. Used Tree-sitter for context-rich chunking.
  2. Created a dependency graph using Tree-sitter to resolve all types of calls (function, class) and imports.
  3. With the dependency graph, I retrieve the best-match chunks along with their neighbors, which helps LLMs deliver better results with complete context.

Below is a dependency graph in Neo4j for a codebase:

Dependency graph from neo4j

This application is now fully deployed, feel free to explore it.
Live Demo: https://code-rag.vercel.app
GitHub: https://github.com/shivamsahu-tech/coderag-ai

Before diving into the technical details, here are a few images that showcase how the system works.

Front Page of codeRAG AI

Chat Page of codeRAG AI

Detailed Workflow: Server-Side Implementation

This blog post covers Version 2 of my CodeRAG system. I’m continuously improving it by adding other cool features, (like agentic behaviour) by learning new approaches. The workflow consists of two main pipelines: the Ingestion Pipeline (for processing and storing the codebase) and the Retrieval Pipeline (for answering queries).

A. Ingestion Pipeline

1. Cloning the Repository

First, I take the repo URL from the client and clone it locally via Git with a unique session ID.

git.Repo.clone_from(repo_url, local_path)

2. Walking Through the Repository

This walk through the repository file by file, extracting the language for each file via an extension dictionary. If it’s a code file, we send it to a function that extracts each chunk_node from that code file.

for root, dirs, files in os.walk(repo_path):
    if ext in LANGUAGES:
        nodes = extract_nodes_from_file(file_path, language)

3. Building the Parser

I used the tree-sitter-language-pack Python library for building the parser for a specific language to parse the code and build the Abstract Syntax Tree (AST).

# get_parser method is imported from the tree-sitter library
parser = get_parser(language)

4. Extracting the Imports

For any file, one way to extract imports is during tree traversal. If there are nodes of type import_statement or import_declaration (depending on the language), we can store them as the file's imports.

Another approach is to extract all file imports for each language using regex on the source code. I used a specific structure for storing imports with fields like imports_from, module, etc.

5. Extracting Chunks & Creating Intra-File Dependencies

The Abstract Syntax Tree nodes have several fields that help us understand details about any code:

Fields:

  • Type (function_definition, class_declaration, identifier, etc.)
  • is_named
  • text (source code)
  • start_point (row, col of first character of source code)
  • end_point (row, col of last character of source code)
  • parent (property to access the parent node)
  • children (array-like property to access all child nodes)

The children are just sub-components of the source code that the current node represents.

5.1. Tree Traversal (Algorithm)

I traverse the tree using a DFS (Depth-First Search) algorithm. Here’s what happens during this DFS traversal:

Starting from the root node of the Abstract Syntax Tree:

  1. If the code string size is less than MIN_CHUNK_SIZE, then I create a chunk_node for this tree node, because a balanced chunk size is important for better retrieval. Here’s how we create the chunk_node:
    • Create a unique node_id
    • Resolve the name of this source code
    • Store the specific values of the tree node in the chunk node (like file path, node type, start_point, end_point, etc.)
    • Extract the calls from this node (we have 2 methods for this, discussed below)
    • Maintain sibling relationships for each node
    • Store the parent node ID in this chunk node
  2. Else, if the node’s code string size is greater than MIN_CHUNK_SIZE, I break this node further and traverse its children.

Here is the chunk_node structure:

{
    "id": "node_id",
    "name": "name",
    "code_str": "code_str",
    "ast_type": "ast_type",
    "file": "file_path",
    "language": "language",
    "start_line": 0,
    "end_line": 10,
    "start_byte": 0,
    "end_byte": 100,
    "size": 100,
    "relationships": {
        "belongs_to": [],
        "parent": [],
        "sibling": [],
        "function_call": [],
        "class_call": [],
        "implements": [],
        "extends": [],
        "imports_from": []
    },
    "metadata": {
        "depth": 1,
        "calls": [],
        "type_references": [],
        "is_definition": false,
        "definition_type": null
    }
}
5.1.1. MIN_CHUNK_SIZE

The Abstract Syntax Tree can take us to the leaf nodes that contain only literals, operators, and keywords, single tokens that may have no value for understanding the file’s code flow. Storing these as chunks can lead to filling the database with trash chunks.

So deciding on a balanced MIN_CHUNK_SIZE is important for storing better context.

Deciding Factors:

  • Your LLM’s token size (directly proportional)
  • Your vector embedding dimensions (directly proportional)
  • Relationship density of your dependency graph (inversely proportional)
5.1.2. Creating a Node ID

I could create a random node ID here, but meaningful instances have better value, so I used this structure for any node ID:

node_id = f"{file_path}:{node.start_point[0]}:{node.type}"
5.1.3. Extracting the Name

Each language has a different node architecture, which leads to different node types and field names. To resolve this, we have some methods:

  • a. Using Normalized Mapping: Create a dictionary mapping fields and their values for each language.
  • b. Use Regex: Use regex on the source code for each language to extract the name.
5.1.4. Extracting the Calls

Here we also have several ways to extract the calls from the node’s source code:

  • a. By Traversing More in Depth: When we traverse to the children or grandchildren of this node, we may find nodes that represent function calls like foo(a, b), with names like call_expression, function_call, etc. By extracting the identifier, we can identify the calls in this codebase.
  • b. By Using Regex: Because source code for any call in any language follows a specific structure and syntax, extracting call names is easy with this method, but less reliable.
5.1.5. Other Considerations

We can extract and store other information here, like class inheritance, interface implementation, etc., as per our requirements, because improvements have no limits.

5.2. Resolving the Calls

Now we have each node’s call names, and to prepare a useful relationship, we have to resolve these calls with the actual node ID of that function or class call. This is required to build the dependency graph.

To resolve the calls of any chunk node, we traverse all chunk nodes and check if there’s any chunk node with the same name as the call.

Example: If any node has a call foo(), I search for the node that has the name foo and add the chunk_node_id in the relationship category.

"relationships": {
    "belongs_to": [],
    "parent": [],
    "sibling": [],
    "function_call": ["file_path.ext:start_point:function_declaration"],
    "class_call": []
}

Now we have each file’s source code as chunked nodes and their resolved intra-file dependencies.

6. Resolving Inter-File Imports

As I already created the complete imports details of each file with module names and extract_from fields, now to resolve those imports with the actual node_ids, we first identify the file from the imports_from field for each language. Then we search in that file's chunk nodes to see whether it has that module. If found, we resolve that import with the chunk_node_id.

At this step, we’ve created a complete dependency graph for all code files of the repository.

7. Handling Non-Code Files

In the repository, there will be lots of files where direct code isn’t written, but they are crucial for understanding the codebase better, like README files, text files, etc.

I used the general methods for chunking as used for documents. Here are some ways to chunk them:

  • Fixed-Size Chunking: The simplest and most basic method, where we split our document into constant predefined sizes.
  • Semantic Chunking: Here, the document is first broken into smaller units like sentences, then embeddings are generated for each unit.
  • Recursive Chunking: For files like READMEs and documentation, this is the best chunking strategy.

8. Storing the Dependency Graph in a Graph Database

I used Neo4j because it also provides the ability to store vector embeddings in nodes and offers functionality to extract nodes based on cosine similarity.

8.1. Code Embedding

We prepare a list of each chunk node’s source code and send it to the embedding model. I used the gemini-embedding-001 model for code embedding.

8.2. Store Chunk Nodes in Neo4j

Add embeddings to each corresponding chunk node and remove any fields that aren’t important. Then create a client for your database and store the nodes in the graph database.

8.3. Store Relationships

Create a list for each relationship with source_id, target_id, and relationship type by traversing all chunk nodes, and store these relationships in the graph database.

B. Retrieval Pipeline

1. Taking User Query with Session ID

When the client creates a request for a user query on this repository, we take both the query and session ID and send them for further processing.

2. Enhancing the User Query

A user might ask “define loginController”. To handle this, we send a message to the LLM with the user_query and custom instructions to enhance this user query.

3. Create Vector Embedding of the Enhanced User Query

Create a vector embedding of the enhanced user query using the same embedding model.

4. Database Nodes Retrieval

  • Retrieve the top K nodes: Using the vector embedding, extract the top K chunk nodes that have the highest cosine similarity with the query embedding.
  • Retrieve all related chunk_nodes: Get all the chunk_nodes that are related to these top K chunk_nodes through relationships.

5. Context Preparation

Now we have the related source code along with their dependencies’ source code. Extract the fields that are helpful as context and concatenate all of them.

6. LLM Response

Now prepare a prompt with the user query, code context, and good instructions. The LLM responds with a beautiful answer to the user query.

Further Improvements

  • Adding Agentic Behaviour: I’m working on implementing agentic behavior where the LLM will decide what context is relevant.
  • Improved Call and Name Extraction: Currently, I’m extracting calls and names via regex, which has limitations. I plan to improve this by extracting details through deeper node traversal in the AST.
  • Implementing a Merging Strategy: I can implement a solution to merge smaller chunks to optimize chunk sizes and improve context quality.
  • Multi-Language Support Expansion: Since Tree-sitter supports many languages, I can extend this to support any language.
  • Routing Specific Queries: For queries like “provide me a summary of this repository”, I plan to create specialized tools for the agent.

Thank You!

If you came till here, thank you for reading! I love working on projects like these that require creative thought and deep technical exploration.