Code-aware indexing

This page builds an indexing pipeline that splits source code along syntax boundaries and gives every chunk the context of its file: an outline trimmed to what the chunk uses, questions the code answers, and the symbols it defines and uses. Tree-sitter is a parser library with one grammar per language. Swiftide wraps it in a chunker and four transformers behind the tree-sitter feature flag. The example below also uses OpenAI, a Redis node cache and Qdrant storage:

cargo add swiftide@0.34 --features tree-sitter,openai,redis,qdrant

Reduced context pipeline

This example indexes a Rust codebase. It sets up an OpenAI client, the Redis URL for the node cache and a chunk size. It needs OPENAI_API_KEY, plus Redis and Qdrant running locally:

let openai_client = integrations::openai::OpenAI::builder()
    .default_embed_model("text-embedding-3-small")
    .default_prompt_model("gpt-3.5-turbo")
    .build()?;

let redis_url = std::env::var("REDIS_URL")
    .as_deref()
    .unwrap_or("redis://localhost:6379")
    .to_owned();

let chunk_size = 2048;

Source: examples/index_codebase_reduced_context.rs#L36-L46

The pipeline gives every chunk an outline of its file, trimmed to what that chunk uses:

indexing::Pipeline::from_loader(FileLoader::new(".").with_extensions(&["rs"]))
    .filter_cached(Redis::try_from_url(
        redis_url,
        "swiftide-examples-codebase-reduced-context",
    )?)
    .then(
        indexing::transformers::OutlineCodeTreeSitter::try_for_language(
            "rust",
            Some(chunk_size),
        )?,
    )
    .then(MetadataQACode::new(openai_client.clone()))
    .then_chunk(ChunkCode::try_for_language_and_chunk_size(
        "rust",
        10..chunk_size,
    )?)
    .then(indexing::transformers::CompressCodeOutline::new(
        openai_client.clone(),
    ))
    .then_in_batch(Embed::new(openai_client.clone()).with_batch_size(10))
    .then_store_with(
        Qdrant::builder()
            .batch_size(50)
            .vector_size(1536)
            .collection_name("swiftide-examples-codebase-reduced-context")
            .build()?,
    )
    .run()
    .await?;

Source: examples/index_codebase_reduced_context.rs#L50-L78

The steps run in this order:

1. filter_cached skips files Redis has already seen, keyed on path and content as Cache keys explains.
2. OutlineCodeTreeSitter parses the whole file and stores an outline in the Outline metadata field. An outline is the file with the bodies of functions and methods removed. Signatures, types, imports and comments stay. The second argument skips files smaller than that many bytes. Here it equals the chunk size, because a file that fits in one chunk needs no outline.
3. MetadataQACode asks a large language model (LLM) for questions and answers about the file. Its default prompt includes the outline when the node has one.
4. ChunkCode splits the file into chunks of more than 10 and up to 2048 bytes.
5. CompressCodeOutline asks the LLM to keep only the outline lines the chunk needs and writes the result back to Outline.
6. Embed embeds the chunks in batches of 10, and Qdrant stores them with vectors of 1536 dimensions, the size of text-embedding-3-small.

See Caching and Storage for the node cache and the storage.

Chunking copies the metadata of a file to every chunk, so each chunk starts with the full outline and the questions of its file. The default embed mode embeds metadata together with the chunk. Compressing the outline keeps that embedded context small and specific.

Cost per chunk

MetadataQACode runs before chunking here, so it sends one prompt per file and every chunk shares the answers. Move it after then_chunk to get questions per chunk, at one prompt per chunk. CompressCodeOutline sends one prompt per chunk of a file larger than the outline threshold. Chunks of smaller files have no outline and cost nothing. On a large codebase that is still many requests. Lower the concurrency or wrap the client in LanguageModelWithBackOff to stay within rate limits.

Question prompt

MetadataQACode writes its reply to Questions and Answers (code). The default prompt asks generic questions. The ask-your-codebase showcase replaces it with a prompt about behaviour, which matches the questions developers ask:

examples/ask-your-codebase/src/index.rs

/// Asks for the questions a developer would ask that a chunk answers.
///
/// The default `MetadataQACode` prompt asks generic questions ("What does this code do?").
/// Questions about behaviour, phrased the way a developer would ask them, match real questions
/// much better.
const QUESTIONS_PROMPT: &str = r"
# Task

List the questions a developer could ask about this codebase that the code below answers, each
with a short answer. They are stored with the code to help a search find it.

Focus on behaviour: when and in which order things happen, what starts or ends a process, which
conditions and configuration change the outcome, what happens on errors, and which extension
points (traits, hooks, builders, options) exist and when they apply.

# Constraints

- Write {{ questions }} questions and answers.
- Phrase questions the way a developer would ask them, e.g. 'When does the agent stop?'.
- Name the relevant types, functions and variants in the answers.
- Only use what can be derived from the code.
- Respond only in this format:

Q1: <question>
A1: <answer>

# Code

```
{{ node.chunk }}
```
";

Source: examples/ask-your-codebase/src/index.rs#L39-L70

The prompt receives questions and the node. The default prompt also renders {{ outline }}, which MetadataQACode sets when the node has an Outline field. The showcase prompt leaves it out, because the showcase builds no outline. Add {{ outline }} to your own prompt when you use it in the pipeline above, or the outline does not reach the LLM. Tree-sitter shows how to pass a prompt template and the number of questions.

References and definitions

MetadataRefsDefsCode adds the symbols each chunk defines and uses, without an LLM:

indexing::Pipeline::from_loader(FileLoader::new("src").with_extensions(&["rs"]))
    .then_chunk(ChunkCode::try_for_language_and_chunk_size(
        "rust",
        10..2048,
    )?)
    .then(MetadataRefsDefsCode::try_from_language("rust")?)
    .then_store_with(MemoryStorage::default())
    .run()
    .await?;

// The same transformer on one node
let refs_defs = MetadataRefsDefsCode::try_from_language("rust")?;
let node = refs_defs
    .transform_node(TextNode::new(r#"fn main() { println!("Hello"); }"#))
    .await?;

assert_eq!(node.metadata.get(NAME_REFERENCES).unwrap(), "println");
assert_eq!(node.metadata.get(NAME_DEFINITIONS).unwrap(), "main");

Source: snippets/src/integrations/tree_sitter.rs#L37-L54

It stores the symbols the chunk defines in Definitions (code) and the symbols it uses but does not define in References (code). Both lists are sorted, without duplicates, and joined with commas. A field is left out when nothing is found. The constants NAME_DEFINITIONS and NAME_REFERENCES in swiftide::indexing::transformers::metadata_refs_defs_code hold the field names. Run it after chunking, so each chunk lists its own symbols.

Languages

ChunkCode, OutlineCodeTreeSitter and MetadataRefsDefsCode take a language. MetadataQACode and CompressCodeOutline take only an LLM client. The constructor for references and definitions is try_from_language; the others use try_for_language. Tree-sitter lists which languages each transformer supports. Perl needs the tree-sitter-perl feature flag.

Code splitter outside a pipeline

CodeSplitter is the splitter inside ChunkCode. Use it directly when you need chunks outside an indexing pipeline, for example in a tool:

// Case-insensitive; "c#", "c-sharp", "cpp" and "c++" parse too
let language: SupportedLanguages = "python".parse()?;
assert_eq!(language.file_extensions(), ["py"]);

// Drop chunks of 100 bytes or less, split anything above 1500 bytes
let splitter = CodeSplitter::builder()
    .try_language(language)?
    .chunk_size(100..1500)
    .build()?;

let chunks = splitter.split(source)?;

// Code with a syntax error comes back whole, as one chunk
let broken = splitter.split("def main(:\n    pass")?;
assert_eq!(broken.len(), 1);

Source: snippets/src/indexing/code.rs#L8-L22

SupportedLanguages parses names without regard to case, and alternative names such as "c#", "c-sharp", "cpp" and "c++" work. Parsing an unknown name with parse() returns a strum parse error that says Matching variant not found. The constructors and try_language return the error Treesitter language not supported instead.

The chunk size is a ChunkSize. A number is a maximum in bytes; ChunkSize::default() is 1500 bytes. A range sets a minimum and a maximum, and a chunk is kept only when it is longer than the minimum. CodeSplitter::new(language) uses the default size.

Gotchas

• The splitter drops chunks at or below the minimum and syntax leaves larger than the maximum, and returns a file with a syntax error as one chunk. Chunking covers these cases.
• OutlineCodeTreeSitter returns an error for a file with a syntax error. The first error aborts run, so add log_errors() and filter_errors() when your codebase contains files that do not parse. See Streams and concurrency.
• OutlineCodeTreeSitter and MetadataRefsDefsCode panic on code in a language they do not handle. In a pipeline the panic happens inside a spawned task, so run fails with a join error.
• Chunk sizes for code are bytes. ChunkMarkdown and ChunkText count characters, see Markdown and text.

Reference

• ChunkCode
• OutlineCodeTreeSitter
• CompressCodeOutline
• MetadataQACode
• MetadataRefsDefsCode and the metadata_refs_defs_code module with NAME_DEFINITIONS and NAME_REFERENCES
• CodeSplitter
• ChunkSize
• SupportedLanguages
• FileLoader
• Embed
• TextNode
• LanguageModelWithBackOff
• Tree-sitter integration and feature flags

Next

How querying works