|
| 1 | +--- |
| 2 | +title: Codebase Indexing |
| 3 | +description: Build a real-time codebase index for retrieval-augmented generation (RAG) using CocoIndex and Tree-sitter. Chunk, embed, and search code with semantic understanding. |
| 4 | +sidebar_class_name: hidden |
| 5 | +slug: /examples/code_index |
| 6 | +canonicalUrl: '/examples/code_index' |
| 7 | +--- |
| 8 | + |
| 9 | +import { GitHubButton, YouTubeButton } from '../../../src/components/GitHubButton'; |
| 10 | + |
| 11 | +<GitHubButton url="https://github.com/cocoindex-io/cocoindex/tree/main/examples/code_embedding"/> |
| 12 | +<YouTubeButton url="https://youtu.be/G3WstvhHO24?si=ndYfM0XRs03_hVPR" /> |
| 13 | + |
| 14 | +## Setup |
| 15 | + |
| 16 | +If you don't have Postgres installed, please follow [installation guide](https://cocoindex.io/docs/getting_started/installation). |
| 17 | + |
| 18 | +## Add the codebase as a source. |
| 19 | + |
| 20 | +Ingest files from the CocoIndex codebase root directory. |
| 21 | + |
| 22 | +```python |
| 23 | +@cocoindex.flow_def(name="CodeEmbedding") |
| 24 | +def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope): |
| 25 | + """ |
| 26 | + Define an example flow that embeds files into a vector database. |
| 27 | + """ |
| 28 | + data_scope["files"] = flow_builder.add_source( |
| 29 | + cocoindex.sources.LocalFile(path="../..", |
| 30 | + included_patterns=["*.py", "*.rs", "*.toml", "*.md", "*.mdx"], |
| 31 | + excluded_patterns=[".*", "target", "**/node_modules"])) |
| 32 | + code_embeddings = data_scope.add_collector() |
| 33 | +``` |
| 34 | + |
| 35 | +- Include files with the extensions of `.py`, `.rs`, `.toml`, `.md`, `.mdx` |
| 36 | +- Exclude files and directories starting `.`, `target` in the root and `node_modules` under any directory. |
| 37 | + |
| 38 | +`flow_builder.add_source` will create a table with sub fields (`filename`, `content`). |
| 39 | +See [documentation](https://cocoindex.io/docs/ops/sources) for more details. |
| 40 | + |
| 41 | + |
| 42 | +## Process each file and collect the information. |
| 43 | + |
| 44 | +### Extract the extension of a filename |
| 45 | + |
| 46 | +We need to pass the language (or extension) to Tree-sitter to parse the code. |
| 47 | +Let's define a function to extract the extension of a filename while processing each file. |
| 48 | +You can find the documentation for custom function [here](https://cocoindex.io/docs/core/custom_function). |
| 49 | + |
| 50 | +```python |
| 51 | +@cocoindex.op.function() |
| 52 | +def extract_extension(filename: str) -> str: |
| 53 | + """Extract the extension of a filename.""" |
| 54 | + return os.path.splitext(filename)[1] |
| 55 | +``` |
| 56 | + |
| 57 | +Then we are going to process each file and collect the information. |
| 58 | + |
| 59 | +```python |
| 60 | +with data_scope["files"].row() as file: |
| 61 | + file["extension"] = file["filename"].transform(extract_extension) |
| 62 | +``` |
| 63 | + |
| 64 | +Here we extract the extension of the filename and store it in the `extension` field. |
| 65 | + |
| 66 | + |
| 67 | +### Split the file into chunks |
| 68 | + |
| 69 | +We will chunk the code with Tree-sitter. |
| 70 | +We use the `SplitRecursively` function to split the file into chunks. |
| 71 | +It is integrated with Tree-sitter, so you can pass in the language to the `language` parameter. |
| 72 | +To see all supported language names and extensions, see the documentation [here](https://cocoindex.io/docs/ops/functions#splitrecursively). 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. |
| 73 | + |
| 74 | +```python |
| 75 | +with data_scope["files"].row() as file: |
| 76 | + file["chunks"] = file["content"].transform( |
| 77 | + cocoindex.functions.SplitRecursively(), |
| 78 | + language=file["extension"], chunk_size=1000, chunk_overlap=300) |
| 79 | +``` |
| 80 | + |
| 81 | + |
| 82 | +### Embed the chunks |
| 83 | + |
| 84 | +We use `SentenceTransformerEmbed` to embed the chunks. |
| 85 | +You can refer to the documentation [here](https://cocoindex.io/docs/ops/functions#sentencetransformerembed). |
| 86 | + |
| 87 | +```python |
| 88 | +@cocoindex.transform_flow() |
| 89 | +def code_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]: |
| 90 | + """ |
| 91 | + Embed the text using a SentenceTransformer model. |
| 92 | + """ |
| 93 | + return text.transform( |
| 94 | + cocoindex.functions.SentenceTransformerEmbed( |
| 95 | + model="sentence-transformers/all-MiniLM-L6-v2")) |
| 96 | +``` |
| 97 | + |
| 98 | +Then for each chunk, we will embed it using the `code_to_embedding` function. and collect the embeddings to the `code_embeddings` collector. |
| 99 | + |
| 100 | +`@cocoindex.transform_flow()` is needed to share the transformation across indexing and query. We build a vector index and query against it, |
| 101 | +the embedding computation needs to be consistent between indexing and querying. See [documentation](https://cocoindex.io/docs/query#transform-flow) for more details. |
| 102 | + |
| 103 | + |
| 104 | +```python |
| 105 | +with data_scope["files"].row() as file: |
| 106 | + with file["chunks"].row() as chunk: |
| 107 | + chunk["embedding"] = chunk["text"].call(code_to_embedding) |
| 108 | + code_embeddings.collect(filename=file["filename"], location=chunk["location"], |
| 109 | + code=chunk["text"], embedding=chunk["embedding"]) |
| 110 | +``` |
| 111 | + |
| 112 | + |
| 113 | +### 2.4 Collect the embeddings |
| 114 | + |
| 115 | +Export the embeddings to a table. |
| 116 | + |
| 117 | +```python |
| 118 | +code_embeddings.export( |
| 119 | + "code_embeddings", |
| 120 | + cocoindex.storages.Postgres(), |
| 121 | + primary_key_fields=["filename", "location"], |
| 122 | + vector_indexes=[cocoindex.VectorIndex("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)]) |
| 123 | +``` |
| 124 | + |
| 125 | +We use Consine Similarity to measure the similarity between the query and the indexed data. |
| 126 | +To learn more about Consine Similarity, see [Wiki](https://en.wikipedia.org/wiki/Cosine_similarity). |
| 127 | + |
| 128 | +## Query the index |
| 129 | +We match against user-provided text by a SQL query, reusing the embedding operation in the indexing flow. |
| 130 | + |
| 131 | +```python |
| 132 | +def search(pool: ConnectionPool, query: str, top_k: int = 5): |
| 133 | + # Get the table name, for the export target in the code_embedding_flow above. |
| 134 | + table_name = cocoindex.utils.get_target_storage_default_name(code_embedding_flow, "code_embeddings") |
| 135 | + # Evaluate the transform flow defined above with the input query, to get the embedding. |
| 136 | + query_vector = code_to_embedding.eval(query) |
| 137 | + # Run the query and get the results. |
| 138 | + with pool.connection() as conn: |
| 139 | + with conn.cursor() as cur: |
| 140 | + cur.execute(f""" |
| 141 | + SELECT filename, code, embedding <=> %s::vector AS distance |
| 142 | + FROM {table_name} ORDER BY distance LIMIT %s |
| 143 | + """, (query_vector, top_k)) |
| 144 | + return [ |
| 145 | + {"filename": row[0], "code": row[1], "score": 1.0 - row[2]} |
| 146 | + for row in cur.fetchall() |
| 147 | + ] |
| 148 | +``` |
| 149 | + |
| 150 | +Define a main function to run the query in terminal. |
| 151 | + |
| 152 | +```python |
| 153 | +def main(): |
| 154 | + # Initialize the database connection pool. |
| 155 | + pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL")) |
| 156 | + # Run queries in a loop to demonstrate the query capabilities. |
| 157 | + while True: |
| 158 | + try: |
| 159 | + query = input("Enter search query (or Enter to quit): ") |
| 160 | + if query == '': |
| 161 | + break |
| 162 | + # Run the query function with the database connection pool and the query. |
| 163 | + results = search(pool, query) |
| 164 | + print("\nSearch results:") |
| 165 | + for result in results: |
| 166 | + print(f"[{result['score']:.3f}] {result['filename']}") |
| 167 | + print(f" {result['code']}") |
| 168 | + print("---") |
| 169 | + print() |
| 170 | + except KeyboardInterrupt: |
| 171 | + break |
| 172 | + |
| 173 | +if __name__ == "__main__": |
| 174 | + main() |
| 175 | +``` |
| 176 | + |
| 177 | +## Run the index setup & update |
| 178 | + |
| 179 | +🎉 Now you are all set! |
| 180 | + |
| 181 | +Run following command to setup and update the index. |
| 182 | +```sh |
| 183 | +cocoindex update --setup main.py |
| 184 | +``` |
| 185 | +You'll see the index updates state in the terminal |
| 186 | + |
| 187 | + |
| 188 | +## Test the query |
| 189 | +At this point, you can start the CocoIndex server and develop your RAG runtime against the data. To test your index, you could |
| 190 | + |
| 191 | +``` bash |
| 192 | +python main.py |
| 193 | +``` |
| 194 | + |
| 195 | +When you see the prompt, you can enter your search query. for example: spec. |
| 196 | + |
| 197 | +You can find the search results in the terminal |
| 198 | + |
| 199 | +The returned results - each entry contains score (Cosine Similarity), filename, and the code snippet that get matched. |
0 commit comments