diff --git a/docs/configuration.md b/docs/configuration.md index bcb823a..6de6537 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -63,6 +63,12 @@ exclude = [ # Whether to respect .gitignore files use_gitignore = true +# Line-based chunker (unknown languages): lines per chunk and overlap. +# Chunks are also clamped to the embedding model's token window so dense +# chunks aren't silently truncated at embed time. +chunk_size = 50 +chunk_overlap = 10 + [embedding] # Embedding model (downloaded automatically on first run) model = "all-MiniLM-L6-v2" @@ -266,6 +272,10 @@ Controls which files are indexed. | `include` | string[] | See above | Glob patterns for files to include | | `exclude` | string[] | See above | Additional exclusion patterns (on top of `.gitignore`) | | `use_gitignore` | bool | `true` | Whether to respect `.gitignore` files | +| `chunk_size` | int | `50` | Lines per chunk for the line-based (unknown-language) chunker. | +| `chunk_overlap` | int | `10` | Overlapping lines between consecutive line-based chunks. Capped below `chunk_size`. | + +Line-based chunks are additionally clamped to the embedding model's token window (`max_seq`) so a dense chunk never silently overflows and gets truncated at embed time — chunks are split to fit. The default model (`all-MiniLM-L6-v2`) has a 256-token window; `bge-small-en-v1.5` and `gte-small` allow 512. ### `[embedding]` diff --git a/src/cli/index.rs b/src/cli/index.rs index 191fee8..5c062d2 100644 --- a/src/cli/index.rs +++ b/src/cli/index.rs @@ -210,7 +210,13 @@ pub async fn run(args: IndexArgs, output: OutputConfig) -> Result<()> { let embed = Embedder::from_config(&config.embedding, &model_dir) .context("Failed to load embedding model")?; - let mut parser = Parser::new().context("Failed to initialize parser")?; + let mut parser = Parser::new() + .context("Failed to initialize parser")? + .with_chunking( + config.index.chunk_size, + config.index.chunk_overlap, + embed.max_seq().unwrap_or(0), + ); // When forcing, delete all existing chunks for this repo first to prevent // unbounded LanceDB growth (each --force reindex was duplicating all data). diff --git a/src/cli/watch.rs b/src/cli/watch.rs index 17baca5..dff9bed 100644 --- a/src/cli/watch.rs +++ b/src/cli/watch.rs @@ -105,7 +105,11 @@ pub async fn run(args: WatchArgs, output: OutputConfig) -> Result<()> { let metadata_store = MetadataStore::open(&db_path)?; let mut vector_store = VectorStore::open(&lance_path).await?; let embed = Embedder::load(&model_dir, &config.embedding.model)?; - let mut parser = Parser::new()?; + let mut parser = Parser::new()?.with_chunking( + config.index.chunk_size, + config.index.chunk_overlap, + embed.max_seq().unwrap_or(0), + ); // Precompile include/exclude patterns let include_patterns: Vec = config diff --git a/src/config.rs b/src/config.rs index 424cfdd..475b680 100644 --- a/src/config.rs +++ b/src/config.rs @@ -63,6 +63,11 @@ pub struct IndexConfig { pub exclude: Vec, /// Whether to respect .gitignore pub use_gitignore: bool, + /// Lines per chunk for the line-based (unknown-language) chunker. Default: 50. + pub chunk_size: usize, + /// Overlapping lines between consecutive line-based chunks. Default: 10. + /// Must be smaller than `chunk_size`. + pub chunk_overlap: usize, } impl Default for IndexConfig { @@ -103,6 +108,8 @@ impl Default for IndexConfig { "**/book/book/**".into(), ], use_gitignore: true, + chunk_size: 50, + chunk_overlap: 10, } } } @@ -1014,6 +1021,24 @@ mod tests { } } + #[test] + fn test_index_chunk_defaults_and_override() { + let def = IndexConfig::default(); + assert_eq!(def.chunk_size, 50); + assert_eq!(def.chunk_overlap, 10); + + let toml_str = r#" +[index] +chunk_size = 80 +chunk_overlap = 16 +"#; + let config: Config = toml::from_str(toml_str).unwrap(); + assert_eq!(config.index.chunk_size, 80); + assert_eq!(config.index.chunk_overlap, 16); + // Unspecified index keys still fall back to defaults. + assert!(config.index.use_gitignore); + } + #[test] fn test_parse_legacy_config() { // Old-style config without backend field should still work diff --git a/src/http/handlers/webhook.rs b/src/http/handlers/webhook.rs index a7836cc..3a22797 100644 --- a/src/http/handlers/webhook.rs +++ b/src/http/handlers/webhook.rs @@ -294,7 +294,11 @@ async fn run_incremental_index( let model_dir = Config::model_cache_dir()?; let embedder = Embedder::load(&model_dir, &config.embedding.model)?; - let mut parser = Parser::new()?; + let mut parser = Parser::new()?.with_chunking( + config.index.chunk_size, + config.index.chunk_overlap, + embedder.max_seq().unwrap_or(0), + ); // Collect files let mut walker = WalkBuilder::new(source_dir); diff --git a/src/index/embedder.rs b/src/index/embedder.rs index a29e30d..c48d6f2 100644 --- a/src/index/embedder.rs +++ b/src/index/embedder.rs @@ -291,6 +291,16 @@ impl Embedder { } } + /// Maximum input sequence length (tokens) the model accepts before it + /// silently truncates. `None` for API backends whose window we don't track. + /// Used by the chunker to keep chunks within the embedder window. + pub fn max_seq(&self) -> Option { + match &self.backend { + EmbedderBackend::Onnx(onnx) => Some(onnx.max_seq), + EmbedderBackend::Api(_) => None, + } + } + /// Get the model name pub fn model_name(&self) -> &str { match &self.backend { diff --git a/src/index/parser.rs b/src/index/parser.rs index b5ced4c..3fd5579 100644 --- a/src/index/parser.rs +++ b/src/index/parser.rs @@ -5,6 +5,21 @@ use tree_sitter::{Language, Node}; use crate::types::{Chunk, ChunkEdge, ChunkEdgeType, ChunkType, ImportEdge, RawImport}; +/// Default lines per line-based chunk (unknown languages). Overridable via +/// `[index].chunk_size`. +pub const DEFAULT_CHUNK_SIZE: usize = 50; +/// Default overlapping lines between consecutive line-based chunks. +/// Overridable via `[index].chunk_overlap`. +pub const DEFAULT_CHUNK_OVERLAP: usize = 10; + +/// Estimate the token count of a text fragment for embedder-window clamping. +/// Uses the standard `chars / 4` heuristic (≈4 chars/token) — deterministic, +/// needs no tokenizer files, and is accurate enough to keep chunks under the +/// model window. Matches the budget estimator used in context assembly. +fn estimate_tokens(text: &str) -> usize { + text.chars().count().div_ceil(4) +} + /// Parses source code using tree-sitter to extract semantic chunks pub struct Parser { rust_parser: tree_sitter::Parser, @@ -13,6 +28,14 @@ pub struct Parser { go_parser: tree_sitter::Parser, java_parser: tree_sitter::Parser, cpp_parser: tree_sitter::Parser, + /// Lines per line-based chunk. + chunk_size: usize, + /// Overlapping lines between consecutive line-based chunks. + chunk_overlap: usize, + /// Embedder token window. When > 0, line/markdown chunks are split so none + /// exceeds this many estimated tokens — preventing silent embedding + /// truncation. 0 = no clamp (window unknown / API backend). + max_chunk_tokens: usize, } impl Parser { @@ -25,9 +48,29 @@ impl Parser { go_parser: create_parser(tree_sitter_go::LANGUAGE.into())?, java_parser: create_parser(tree_sitter_java::LANGUAGE.into())?, cpp_parser: create_parser(tree_sitter_cpp::LANGUAGE.into())?, + chunk_size: DEFAULT_CHUNK_SIZE, + chunk_overlap: DEFAULT_CHUNK_OVERLAP, + max_chunk_tokens: 0, }) } + /// Configure line-chunk sizing and the embedder token window for clamping. + /// `max_chunk_tokens` is typically the embedder's `max_seq` (0 to disable + /// clamping). Invalid pairings are sanitized: `chunk_size` is forced to at + /// least 1, and `chunk_overlap` is capped below `chunk_size` to guarantee + /// forward progress. + pub fn with_chunking( + mut self, + chunk_size: usize, + chunk_overlap: usize, + max_chunk_tokens: usize, + ) -> Self { + self.chunk_size = chunk_size.max(1); + self.chunk_overlap = chunk_overlap.min(self.chunk_size - 1); + self.max_chunk_tokens = max_chunk_tokens; + self + } + /// Parse a file and extract semantic chunks pub fn parse_file(&mut self, path: &Path, content: &str) -> Result> { let language = detect_language(path); @@ -681,18 +724,34 @@ impl Parser { chunks } - /// Fall back to line-based chunking for unknown languages + /// Fall back to line-based chunking for unknown languages. + /// + /// Chunk size and overlap come from `[index]` config. When `max_chunk_tokens` + /// is set (the embedder window), each chunk is shrunk to fit so the embedder + /// never silently truncates it — a 50-line dense-code chunk can otherwise + /// blow past a 256/512-token window. fn chunk_by_lines(&self, path: &Path, content: &str) -> Vec { let lines: Vec<&str> = content.lines().collect(); - let chunk_size = 50; // lines per chunk - let overlap = 10; + let overlap = self.chunk_overlap; + let language = detect_language(path).unwrap_or_else(|| "unknown".to_string()); let mut chunks = Vec::new(); let mut start = 0; while start < lines.len() { - let end = (start + chunk_size).min(lines.len()); - let chunk_content = lines[start..end].join("\n"); + let mut end = (start + self.chunk_size).min(lines.len()); + + // Clamp to the embedder token window: shrink the chunk until it fits. + // Always keep at least one line so we make forward progress (a lone + // over-window line can't be split without breaking it; the embedder + // truncates that pathological case as before). + if self.max_chunk_tokens > 0 { + while end > start + 1 + && estimate_tokens(&lines[start..end].join("\n")) > self.max_chunk_tokens + { + end -= 1; + } + } chunks.push(Chunk { id: generate_chunk_id(path, start as u32 + 1, end as u32), @@ -701,15 +760,17 @@ impl Parser { name: None, start_line: start as u32 + 1, end_line: end as u32, - content: chunk_content, - language: detect_language(path).unwrap_or_else(|| "unknown".to_string()), + content: lines[start..end].join("\n"), + language: language.clone(), tags: String::new(), }); if end >= lines.len() { break; } - start = end - overlap; + // Advance with overlap, but never backwards or in place — the clamp + // can make a chunk shorter than `overlap`. + start = end.saturating_sub(overlap).max(start + 1); } chunks @@ -2359,4 +2420,63 @@ fn standalone_function() { let edges = parser.extract_chunk_edges(&path, content, &chunks); assert!(edges.is_empty()); } + + #[test] + fn test_estimate_tokens_heuristic() { + assert_eq!(estimate_tokens(""), 0); + assert_eq!(estimate_tokens("abcd"), 1); // 4 chars => 1 + assert_eq!(estimate_tokens("abcde"), 2); // 5 chars => ceil(5/4) = 2 + } + + #[test] + fn test_with_chunking_sanitizes_overlap() { + // overlap >= size would stall the chunker; it must be capped below size. + let parser = Parser::new().unwrap().with_chunking(5, 10, 0); + assert_eq!(parser.chunk_size, 5); + assert_eq!(parser.chunk_overlap, 4); + // size 0 is forced to at least 1. + let parser = Parser::new().unwrap().with_chunking(0, 0, 0); + assert_eq!(parser.chunk_size, 1); + assert_eq!(parser.chunk_overlap, 0); + } + + #[test] + fn test_chunk_by_lines_configurable_size_overlap() { + // 10-line chunks, 2-line overlap, no token clamp. + let parser = Parser::new().unwrap().with_chunking(10, 2, 0); + let content = (1..=25).map(|i| format!("line{i}")).collect::>().join("\n"); + let path = PathBuf::from("data.unknownext"); + let chunks = parser.chunk_by_lines(&path, &content); + + // start 0..10, then 8..18, then 16..25 => 3 chunks + assert_eq!(chunks.len(), 3); + assert_eq!((chunks[0].start_line, chunks[0].end_line), (1, 10)); + assert_eq!(chunks[1].start_line, 9); // 10 - 2 overlap + 1 (1-based) + assert_eq!(chunks[2].end_line, 25); + } + + #[test] + fn test_chunk_by_lines_clamps_to_token_window() { + // 50-line window but an 8-token clamp: dense lines force smaller chunks + // so nothing exceeds the embedder window (no silent truncation). + let cap = 8; + let parser = Parser::new().unwrap().with_chunking(50, 1, cap); + // 12 lines, 8 chars each (~2 tokens/line). The full 50-line window would + // be ~30+ tokens — must be split. + let content = (0..12).map(|_| "yyyyyyyy".to_string()).collect::>().join("\n"); + let path = PathBuf::from("dense.unknownext"); + let chunks = parser.chunk_by_lines(&path, &content); + + assert!(chunks.len() > 1, "clamp should split the window into multiple chunks"); + for c in &chunks { + assert!( + estimate_tokens(&c.content) <= cap, + "chunk exceeds token window: {} tokens", + estimate_tokens(&c.content) + ); + } + // Coverage: chunks span the whole file start-to-end. + assert_eq!(chunks.first().unwrap().start_line, 1); + assert_eq!(chunks.last().unwrap().end_line, 12); + } }