Build Real-Time Codebase Indexing for AI Coding agents

This post is also on cocoindex.io, with interactive figures. The code uses CocoIndex v1.

Build a live semantic index over a codebase with CocoIndex V1. Point it at a repo, and you get a vector index you can search in natural language ("where do we embed chunks?") that updates itself as you edit. It's the kind of fresh, low-latency code context a coding agent needs, and it's about 100 lines of plain async Python.

CocoIndex has built-in support for codebase chunking, with native Tree-sitter parsing and live updates.

The heavy lifting (incremental processing, change tracking, managed targets) runs in a Rust engine underneath. With a source that watches for changes, like the local filesystem in live mode, the index tracks your edits as they happen: save a file, and only the chunks that changed get re-embedded and re-upserted.

Use case

• Code context for agents: semantic context for Claude, Codex, OpenCode, Factory.
• Code search: natural-language and semantic search over your repo.
• Review & refactor agents: context for code review, security analysis, and large-scale refactoring.

Why CocoIndex for codebase indexing

A codebase is hard to keep indexed well, and it exercises most of what CocoIndex was built for:

• Chunks follow real code structure (functions, classes, blocks) instead of arbitrary line windows. Tree-sitter parsing is built in for every major language.
• Code changes constantly: every commit, save, and rebase mutates a handful of files. CocoIndex re-embeds only the chunks that changed, and with live=True plus cocoindex update -L it keeps watching and applies each edit as it happens.
• The engine is Rust. The run below indexes 622 files in 43 seconds cold; after that, a save costs one file.
• The pipeline is ordinary async Python and your own types. Pick your embedding model, chunking strategy, and vector database; if your agent can write Python, it can extend this flow.

href="https://github.com/cocoindex-io/cocoindex/tree/main/examples/code_embedding"
path="examples/code_embedding"
description="The whole example is ~100 lines of plain Python. Clone it, point it at your repo, and swap in your favorite model or vector store."
/>

Overview

From a high level, these are the steps:

1. Read code files from a local directory.
2. Split each file into syntax-aware chunks with Tree-sitter, then embed every chunk.
3. Store the chunks and their embeddings in Postgres (as target states).

With CocoIndex, you declare the transformation logic with native Python, without worrying about handling updates.

Think: target_state = transformation(source_state).

When your codebase changes, or your processing logic changes (a different embedding model, a new chunk size), only the difference is reprocessed. The Incremental updates section walks through exactly what gets recomputed when.

Setup

In this example we store the indexed rows in Postgres with the pgvector extension. Your metadata and the vector search index both live there. CocoIndex supports many targets, so you can pick the one that works best for you.

1. If you don't have Postgres running, use the Docker Compose file in the cocoindex repo to start one.

   docker compose -f dev/postgres.yaml up -d

1. Point CocoIndex at it

   export POSTGRES_URL="postgres://cocoindex:cocoindex@localhost/cocoindex"

1. Install CocoIndex and the dependencies

   pip install -U "cocoindex[postgres,sentence_transformers]" asyncpg pgvector numpy python-dotenv

Define the App

Apps are the top-level runnable unit in CocoIndex. An app names your pipeline and binds a main function with its parameters.

Define the data and shared resources

This block sets up two things the rest of the code builds on. CodeEmbedding defines one row of the output table: each chunk of code becomes one row, with its text, location, and embedding vector. coco_lifespan provides the shared resources every step needs, the Postgres connection pool and the embedding model, built once at startup instead of per file.

import os
import pathlib
from dataclasses import dataclass
from typing import AsyncIterator, Annotated

import asyncpg
from numpy.typing import NDArray

import cocoindex as coco
from cocoindex.connectors import localfs, postgres
from cocoindex.ops.text import RecursiveSplitter, detect_code_language
from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder
from cocoindex.resources.chunk import Chunk
from cocoindex.resources.file import FileLike, PatternFilePathMatcher
from cocoindex.resources.id import IdGenerator

DATABASE_URL = os.getenv("POSTGRES_URL", "postgres://cocoindex:cocoindex@localhost/cocoindex")
TABLE_NAME = "code_embeddings"
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

PG_DB = coco.ContextKey[asyncpg.Pool]("code_embedding_db")
EMBEDDER = coco.ContextKey[SentenceTransformerEmbedder]("embedder", detect_change=True)

_splitter = RecursiveSplitter()

@dataclass
class CodeEmbedding:
    id: int
    filename: str
    code: str
    embedding: Annotated[NDArray, EMBEDDER]
    start_line: int
    end_line: int

@coco.lifespan
async def coco_lifespan(builder: coco.EnvironmentBuilder) -> AsyncIterator[None]:
    async with asyncpg.create_pool(DATABASE_URL) as pool:
        builder.provide(PG_DB, pool)
        builder.provide(EMBEDDER, SentenceTransformerEmbedder(EMBED_MODEL))
        yield

Define file processing

With the row schema and shared resources in place, we can write the actual work: turning a file into rows.

Process file

@coco.fn(memo=True)
async def process_file(
    file: FileLike,
    table: postgres.TableTarget[CodeEmbedding],
) -> None:
    text = await file.read_text()
    language = detect_code_language(filename=str(file.file_path.path.name))
    chunks = _splitter.split(
        text,
        chunk_size=1000,
        min_chunk_size=300,
        chunk_overlap=300,
        language=language,
    )
    id_gen = IdGenerator()
    await coco.map(process_chunk, chunks, file.file_path.path, id_gen, table)

process_file runs once per file. It reads the file, detects the language so Tree-sitter can parse it, splits the code along the syntax tree, and maps each chunk to process_chunk.

Each chunk is a coherent syntactic unit, so retrieval returns whole functions or blocks, never a fragment cut mid-statement. All major languages are supported (Python, Rust, JavaScript, TypeScript, Java, C++, and more); unknown types fall back to plain text.

Note that @coco.fn with memo=True is what makes this incremental: if a file's content and this function's code are both unchanged, the whole file is skipped on the next run. coco.map fans out to one process_chunk call per chunk.

Here is what chunking produces: each file is split into syntactic chunks, each with its own location and text.

Process chunk

Next, we embed the chunk with the shared embedder and declare the target row.

@coco.fn
async def process_chunk(
    chunk: Chunk,
    filename: pathlib.PurePath,
    id_gen: IdGenerator,
    table: postgres.TableTarget[CodeEmbedding],
) -> None:
    embedding = await coco.use_context(EMBEDDER).embed(chunk.text)
    table.declare_row(
        row=CodeEmbedding(
            id=await id_gen.next_id(chunk.text),
            filename=str(filename),
            code=chunk.text,
            embedding=embedding,
            start_line=chunk.start.line,
            end_line=chunk.end.line,
        ),
    )

We use SentenceTransformerEmbedder with all-MiniLM-L6-v2; there are 12k+ sentence-transformer models on Hugging Face, so swap in whichever you prefer. Note chunk.start.line and chunk.end.line: V1 carries line numbers through, so search results point straight at the lines that matched.

Define the main function

app_main wires the source to the target. It mounts the Postgres table (with a vector index), walks the codebase, and mounts one processing component per file.

@coco.fn
async def app_main(sourcedir: pathlib.Path) -> None:
    target_table = await postgres.mount_table_target(
        PG_DB,
        table_name=TABLE_NAME,
        table_schema=await postgres.TableSchema.from_class(
            CodeEmbedding, primary_key=["id"],
        ),
    )
    target_table.declare_vector_index(column="embedding")

    files = localfs.walk_dir(
        sourcedir,
        recursive=True,
        path_matcher=PatternFilePathMatcher(
            included_patterns=["**/*.py", "**/*.rs", "**/*.toml", "**/*.md", "**/*.mdx"],
            excluded_patterns=["**/.*", "**/target", "**/node_modules"],
        ),
        live=True,  # watch for changes; pass -L to `cocoindex update` to run live
    )
    await coco.mount_each(process_file, files.items(), target_table)

We walk the codebase from sourcedir (point this at the repository you want to index). We index files with the extensions .py, .rs, .toml, .md, .mdx, and skip dot-directories, target, and node_modules.

mount_table_target creates and manages the Postgres table for you: schema, the pgvector index, idempotent upserts, and orphan cleanup when a file disappears. live=True makes the filesystem source watch for changes, and mount_each runs one component per file so the engine can track and update them independently.

Create the App

Bind app_main into a coco.App and point it at the codebase root.

app = coco.App(
    coco.AppConfig(name="CodeEmbeddingV1"),
    app_main,
    sourcedir=pathlib.Path(__file__).parent / ".." / "..",  # index from repo root
)

That is the entire indexing path.

Run the pipeline

Run the cocoindex CLI to set up and update the index. Choose catch-up (scan, sync, exit) or live (catch up, then keep watching):

# Catch-up run
cocoindex update main

# Live run: keep watching for file changes
cocoindex update -L main

You'll see the index update state in the terminal:

✅ app_main: 1 total | 1 added
✅ declare_target_state_with_child: 1 total | 1 added
✅ _MountEachLiveComponent.process: 1 total | 1 added
✅ process_file: 622 total | 622 added
⏳ Elapsed: 43.4s

Query the index

We match user text against the index with a plain SQL query, reusing the same embedder from the indexing flow so indexing and querying stay consistent.

async def query_once(pool, embedder, query: str, *, top_k: int = 5) -> None:
    query_vec = await embedder.embed(query)
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            f"""
            SELECT filename, code, embedding <=> $1 AS distance, start_line, end_line
            FROM "{TABLE_NAME}"
            ORDER BY distance ASC
            LIMIT $2
            """,
            query_vec, top_k,
        )
    for r in rows:
        score = 1.0 - float(r["distance"])
        print(f"[{score:.3f}] {r['filename']} (L{r['start_line']}-L{r['end_line']})")
        print(f"    {r['code']}")
        print("---")

The <=> operator is pgvector's cosine distance. We turn it into a similarity score and print the filename, the matched line range, and the code snippet. Run it against your index:

python main.py "where do we embed chunks"

It prints the top five chunks, each with its similarity score, file, and line range. Snippets are cut to their first lines here (…):

[0.527] rust/ops_text/src/split/by_separators.rs (L72-L94)
            let bytes = text.as_bytes();

        // Collect raw chunks (byte ranges)
        struct RawChunk {
            start: usize,
            end: usize,
    …
---
[0.459] rust/ops_text/src/split/by_separators.rs (L116-L138)
            let mut positions: Vec<Position> = raw_chunks
            .iter()
            .flat_map(|c| vec![Position::new(c.start), Position::new(c.end)])
            .collect();

        set_output_positions(text, positions.iter_mut());
    …
---
[0.448] rust/ops_text/src/split/recursive.rs (L560-L578)
        fn split_root_chunk(&self, kind: ChunkKind<'t>) -> Vec<ChunkOutput> {
        let mut atom_collector = AtomChunksCollector {
            full_text: self.full_text,
            min_level: 0,
            curr_level: 0,
            atom_chunks: Vec::new(),
    …
---
[0.447] rust/sdk/cocoindex/tests/code.rs (L300-L314)
    #[test]
fn code_ast_split_chunk_positions_are_accurate() {
    // Three-line Python source; verify start/end positions of first chunk.
    let src = "a = 1\nb = 2\nc = 3\n";
    let ast = CodeAst::new(src, "python").unwrap();
    let chunks = ast.split(200, None, None).unwrap(); // one big chunk
    …
---
[0.444] rust/sdk/cocoindex/src/ops/code.rs (L315-L341)
            let config = RecursiveChunkConfig {
            chunk_size,
            min_chunk_size,
            chunk_overlap,
        };
        let raw = chunker.split(&self.source, config);
    …
---

Incremental updates

You never compute a diff or write update logic: you change something, and CocoIndex works out what to embed, upsert, and delete. @coco.fn(memo=True) decides what to recompute: a file is skipped when its content and the function's code are both unchanged, and an embedding is reused when its chunk text is unchanged. mount_table_target decides what to write: each row's id is derived from its chunk's content, so only rows that actually changed are upserted, and rows whose source is gone are deleted.

The same machinery covers two kinds of change: changes to your data (the code being indexed) and changes to your logic (the pipeline itself).

Data changes.

• A file is added: only that file is chunked and embedded, and its rows are inserted. The rest of the repo is untouched.
• A file is deleted: its rows are removed from the target.
• A file is edited: the file is re-chunked, and the new chunks usually overlap the old ones. Chunks whose text is unchanged keep their id and embedding, so they are left as-is; genuinely new chunks are embedded and inserted; chunks that no longer exist are deleted. Edit one function and only that function's chunks are re-embedded, even though the whole file was re-read.

Logic changes. A pipeline change is reconciled the same way: CocoIndex compares the new output against what is already in Postgres and applies only the difference.

• Change the file patterns (included_patterns / excluded_patterns): files that now match are added automatically; files that no longer match have their rows deleted automatically.
• Tune the chunking (chunk size, overlap): files are re-chunked, producing the same partial-overlap diff shown above: unchanged chunks are no-ops, new chunks are embedded and inserted, dropped chunks are deleted.
• Swap the embedding model: the vectors genuinely change, so all embeddings are recomputed, but row identity is stable: it is an in-place update of the embedding column, with no rows added or removed.

A catch-up run (cocoindex update main) does this once and exits; live mode (cocoindex update -L main) keeps watching and applies each change as you code.

Next steps

The full, runnable example is in the CocoIndex repo: examples/code_embedding, with a walkthrough in the index-codebase example.

• Point sourcedir at your own repository.
• Swap the embedding model or the target store; the pipeline is plain Python.
• Read the CocoIndex V1 docs for more on incremental processing, components, and targets.

Issues and contributions are welcome on GitHub.