<img src="https://media.daily.dev/image/upload/s--A9Hm-r70--/f_auto/v1742779758/ugc/content_30133bda-8bd4-407f-8e92-4bb81e0b6e52" alt="cover">
<h2>Overview</h2>
In this blog, we will show you how to index codebase for RAG with CocoIndex. CocoIndex is a tool to help you index and query your data. It is designed to be used as a framework to build your own data pipeline. CocoIndex provides built-in support for code base chunking, with native Tree-sitter support.
<h2>Tree-sitter</h2>
<a href="https://en.wikipedia.org/wiki/Tree-sitter_%28parser_generator%29" target="_blank" rel="noopener nofollow">Tree-sitter</a> is a parser generator tool and an incremental parsing library, it is available in Rust 🦀 - <a href="https://github.com/tree-sitter/tree-sitter" target="_blank" rel="noopener nofollow">GitHub</a>. CocoIndex has built-in Rust integration with Tree-sitter to efficiently parse code and extract syntax trees for various programming languages.
Codebase chunking is the process of breaking down a codebase into smaller, semantically meaningful chunks. CocoIndex leverages Tree-sitter’s capabilities to intelligently chunk code based on the actual syntax structure rather than arbitrary line breaks. These semantically coherent chunks are then used to build a more effective index for RAG systems, enabling more precise code retrieval and better context preservation.
Fast pass 🚀 - you can find the full code <a href="https://github.com/cocoindex-io/cocoindex/tree/main/examples/code_embedding" target="_blank" rel="noopener nofollow">here</a>. Only ~ 50 lines of Python code for the RAG pipeline, check it out 🤗!
Please give <a href="https://github.com/cocoindex-io/cocoindex" target="_blank" rel="noopener nofollow">CocoIndex on Github</a> a star to support us if you like our work. Thank you so much with a warm coconut hug 🥥🤗. <a href="https://github.com/cocoindex-io/cocoindex" target="_blank" rel="noopener nofollow"><img src="https://img.shields.io/github/stars/cocoindex-io/cocoindex?color=5B5BD6" alt="GitHub"></a>
<h2>Prerequisites</h2>
If you don’t have Postgres installed, please refer to the <a href="https://cocoindex.io/docs/getting_started/installation" target="_blank" rel="noopener nofollow">installation guide</a>. CocoIndex uses Postgres to manage the data index, we have it on our roadmap to support other databases, including the in-progress ones. If you are interested in other databases, please let us know by creating a <a href="https://github.com/cocoindex-io/cocoindex/issues" target="_blank" rel="noopener nofollow">GitHub issue</a> or <a href="https://discord.com/invite/zpA9S2DR7s" target="_blank" rel="noopener nofollow">Discord</a>.
<h2>Define cocoIndex Flow</h2>
Let’s define the cocoIndex flow to read from a codebase and index it for RAG.
<img src="https://media.daily.dev/image/upload/s--Y6nUe4Or--/f_auto/v1742779790/ugc/content_61dd6ae1-9a9b-448e-ab34-4bce4d2997d7" alt="flow">
The flow diagram above illustrates how we’ll process our codebase:
<ol>
<li>Read code files from the local filesystem</li>
<li>Extract file extensions</li>
<li>Split code into semantic chunks using Tree-sitter</li>
<li>Generate embeddings for each chunk</li>
<li>Store in a vector database for retrieval</li>
</ol>
Let’s implement this flow step by step.
<h3>1. Add the codebase as a source.</h3>
<pre><code class="language-python">@cocoindex.flow_def(name=&quot;CodeEmbedding&quot;)
def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
 &quot;&quot;&quot;
 Define an example flow that embeds files into a vector database.
 &quot;&quot;&quot;
 data_scope[&quot;files&quot;] = flow_builder.add_source(
 cocoindex.sources.LocalFile(path=&quot;../..&quot;,
 included_patterns=[&quot;*.py&quot;, &quot;*.rs&quot;, &quot;*.toml&quot;, &quot;*.md&quot;, &quot;*.mdx&quot;],
 excluded_patterns=[&quot;.*&quot;, &quot;target&quot;, &quot;**/node_modules&quot;]))
 code_embeddings = data_scope.add_collector()
</code></pre>
In this example, we are going to index the cocoindex codebase from the root directory. You can change the path to the codebase you want to index. We will index all the files with the extensions of <code>.py</code>, <code>.rs</code>, <code>.toml</code>, <code>.md</code>, <code>.mdx</code>, and skip directories starting with ., target (in the root) and node_modules (under any directory).
<code>flow_builder.add_source</code> will create a table with the following sub fields, see <a href="https://cocoindex.io/docs/ops/sources" target="_blank" rel="noopener nofollow">documentation</a> here.
<ul>
<li><code>filename</code> (key, type: <code>str</code>): the filename of the file, e.g. <code>dir1/file1.md</code></li>
<li><code>content</code> (type: <code>str</code> if <code>binary</code> is <code>False</code>, otherwise <code>bytes</code>): the content of the file</li>
</ul>
<h3>2. Process each file and collect the information.</h3>
<h4>2.1 Extract the extension of a filename</h4>
First let’s define a function to extract the extension of a filename while processing each file. You can find the documentation for custom function <a href="https://cocoindex.io/docs/core/custom_function" target="_blank" rel="noopener nofollow">here</a>.
<pre><code class="language-python">@cocoindex.op.function()
def extract_extension(filename: str) -&gt; str:
 &quot;&quot;&quot;Extract the extension of a filename.&quot;&quot;&quot;
 return os.path.splitext(filename)[1]
</code></pre>
Then we are going to process each file and collect the information.
<pre><code class="language-python"> # ...
 with data_scope[&quot;files&quot;].row() as file:
 file[&quot;extension&quot;] = file[&quot;filename&quot;].transform(extract_extension)
</code></pre>
Here we extract the extension of the filename and store it in the <code>extension</code> field. for example, if the filename is <code>spec.rs</code>, the <code>extension</code> field will be <code>.rs</code>.
<h4>2.2 Split the file into chunks</h4>
Next, we are going to split the file into chunks. We use the <code>SplitRecursively</code> function to split the file into chunks. You can find the documentation for the function <a href="https://cocoindex.io/docs/core/functions#splitrecursively" target="_blank" rel="noopener nofollow">here</a>.
CocoIndex provides built-in support for Tree-sitter, so you can pass in the language to the <code>language</code> parameter.
To see all supported language names and extensions, see the documentation <a href="https://cocoindex.io/docs/ops/functions#splitrecursively" target="_blank" rel="noopener nofollow">here</a>. All the major languages are supported, e.g., Python, Rust, JavaScript, TypeScript, Java, C++, etc. If it’s unspecified or the specified language is not supported, it will be treated as plain text.
<pre><code class="language-python"> with data_scope[&quot;files&quot;].row() as file:
 # ...
 file[&quot;chunks&quot;] = file[&quot;content&quot;].transform(
 cocoindex.functions.SplitRecursively(),
 language=file[&quot;extension&quot;], chunk_size=1000, chunk_overlap=300) 
</code></pre>
<h4>2.3 Embed the chunks</h4>
We will use the <code>SentenceTransformerEmbed</code> function to embed the chunks. You can find the documentation for the function <a href="https://cocoindex.io/docs/ops/functions#sentencetransformerembed" target="_blank" rel="noopener nofollow">here</a>. There are 12k models supported by 🤗 <a href="https://huggingface.co/models?other=sentence-transformers" target="_blank" rel="noopener nofollow">Hugging Face</a>. You can just pick your favorite model.
<pre><code class="language-python"> def code_to_embedding(text: cocoindex.DataSlice) -&gt; cocoindex.DataSlice:
 &quot;&quot;&quot;
 Embed the text using a SentenceTransformer model.
 &quot;&quot;&quot;
 return text.transform(
 cocoindex.functions.SentenceTransformerEmbed(
 model=&quot;sentence-transformers/all-MiniLM-L6-v2&quot;))
</code></pre>
Then for each chunk, we will embed it using the <code>code_to_embedding</code> function. and collect the embeddings to the <code>code_embeddings</code> collector.
We extract this code_to_embedding function instead of directly calling transform(cocoindex.functions.SentenceTransformerEmbed(…)) in place.
This is because we want to make this one shared between the indexing flow building and the query handler definition. Alternatively, to make it simpler. It’s also OK to avoid this extra function and directly do things in place - not a big deal to copy paste a little bit, we did this for the <a href="https://cocoindex.io/docs/getting_started/quickstart" target="_blank" rel="noopener nofollow">quickstart</a> project.
<pre><code class="language-python"> with data_scope[&quot;files&quot;].row() as file:
 # ...
 with file[&quot;chunks&quot;].row() as chunk:
 chunk[&quot;embedding&quot;] = chunk[&quot;text&quot;].call(code_to_embedding)
 code_embeddings.collect(filename=file[&quot;filename&quot;], location=chunk[&quot;location&quot;],
 code=chunk[&quot;text&quot;], embedding=chunk[&quot;embedding&quot;])
</code></pre>
<h4>2.4 Collect the embeddings</h4>
Finally, let’s export the embeddings to a table.
<pre><code class="language-python"> code_embeddings.export(
 &quot;code_embeddings&quot;,
 cocoindex.storages.Postgres(),
 primary_key_fields=[&quot;filename&quot;, &quot;location&quot;],
 vector_index=[(&quot;embedding&quot;, cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])

</code></pre>
<h3>3. Setup Query Handler for your index</h3>
We will use the <code>SimpleSemanticsQueryHandler</code> to query the index. Note that we need to pass in the <code>code_to_embedding</code> function to the <code>query_transform_flow</code> parameter. This is because the query handler will use the same embedding model as the one used in the flow.
<pre><code class="language-python">query_handler = cocoindex.query.SimpleSemanticsQueryHandler(
 name=&quot;SemanticsSearch&quot;,
 flow=code_embedding_flow,
 target_name=&quot;code_embeddings&quot;,
 query_transform_flow=code_to_embedding,
 default_similarity_metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)
</code></pre>
Define a main function to run the query handler.
<pre><code class="language-python">@cocoindex.main_fn()
def _run():
 # Run queries in a loop to demonstrate the query capabilities.
 while True:
 try:
 query = input(&quot;Enter search query (or Enter to quit): &quot;)
 if query == &#x27;&#x27;:
 break
 results, _ = query_handler.search(query, 10)
 print(&quot;\nSearch results:&quot;)
 for result in results:
 print(f&quot;[{result.score:.3f}] {result.data[&#x27;filename&#x27;]}&quot;)
 print(f&quot; {result.data[&#x27;code&#x27;]}&quot;)
 print(&quot;---&quot;)
 print()
 except KeyboardInterrupt:
 break

if __name__ == &quot;__main__&quot;:
 load_dotenv(override=True)
 _run()
</code></pre>
The @cocoindex.main_fn() decorator initializes the library with settings loaded from environment variables.
See <a href="https://cocoindex.io/docs/core/initialization" target="_blank" rel="noopener nofollow">documentation for initialization</a> for more details.
<h2>Run the index setup &amp; update</h2>
🎉 Now you are all set!
Run following commands to setup and update the index.
<pre><code>python main.py cocoindex setup
python main.py cocoindex update
</code></pre>
You’ll see the index updates state in the terminal
<img src="https://media.daily.dev/image/upload/s--hAPAOBLF--/f_auto/v1742779810/ugc/content_3605259c-5f83-4955-a39b-ec94bfdca031" alt="terminal">
<h2>Test the query</h2>
At this point, you can start the cocoindex server and develop your RAG runtime agains the data.

The official daily.dev AI community. Led by thought leaders and AI experts. Stay tuned for insights, news and discussions. 

CocoIndex is a tool designed to index and query your codebase, facilitating the construction of a data pipeline. It uses Tree-sitter to intelligently chunk code based on syntax and provides built-in Rust integration. The process involves reading code files, extracting extensions, chunking code, generating embeddings, and storing them in a vector database. CocoIndex supports various languages and allows users to embed code using models from Hugging Face. It leverages Postgres for managing the data index, with plans to support other databases.

Want AI to Actually Understand Your Code? This Tool Says It Can Help