diff --git a/crates/csp/src/indexing/index.rs b/crates/csp/src/indexing/index.rs index 5570218..61d1d28 100644 --- a/crates/csp/src/indexing/index.rs +++ b/crates/csp/src/indexing/index.rs @@ -255,17 +255,34 @@ impl CspIndex { ) } - /// Find chunks similar to a seed, excluding the seed itself. + /// Find chunks similar to a seed, excluding the seed itself. Candidates are + /// limited to the seed's language when it has one. pub fn find_related(&self, seed: &Chunk, options: &QueryOptions) -> Vec { let top_k = options.top_k.unwrap_or(DEFAULT_TOP_K); if top_k == 0 || self.chunks.is_empty() { return Vec::new(); } + // Mirror upstream (`if target.language`): an absent or empty language is + // falsy there, so the search stays unrestricted. + let selector = seed + .language + .as_ref() + .filter(|lang| !lang.is_empty()) + .and_then(|lang| { + self.build_selector(&QueryOptions { + filter_languages: Some(vec![lang.clone()]), + ..Default::default() + }) + }); + if selector.as_ref().is_some_and(Vec::is_empty) { + return Vec::new(); + } + let query_embedding = self.model.encode(std::slice::from_ref(&seed.content)); let batch = self .semantic_index - .query(&query_embedding, top_k + 1, None) + .query(&query_embedding, top_k + 1, selector.as_deref()) .unwrap_or_default(); let Some(first) = batch.into_iter().next() else { return Vec::new(); diff --git a/crates/csp/src/indexing/index/tests.rs b/crates/csp/src/indexing/index/tests.rs index 97d5377..cf4c178 100644 --- a/crates/csp/src/indexing/index/tests.rs +++ b/crates/csp/src/indexing/index/tests.rs @@ -121,6 +121,96 @@ fn find_related_excludes_seed() { assert!(results.len() <= 5); } +#[test] +fn find_related_restricts_to_seed_language() { + let chunks = vec![ + make_chunk( + "a.rs", + 1, + 3, + Some("rust"), + "fn add_numbers(a: i32, b: i32) -> i32 { a + b }", + ), + make_chunk( + "b.rs", + 1, + 3, + Some("rust"), + "fn add_values(a: i32, b: i32) -> i32 { a + b }", + ), + make_chunk( + "c.py", + 1, + 3, + Some("python"), + "def add_numbers(a, b): return a + b", + ), + make_chunk( + "d.py", + 1, + 3, + Some("python"), + "def add_values(a, b): return a + b", + ), + ]; + let idx = build_index(chunks.clone()); + let opts = QueryOptions { + top_k: Some(5), + ..Default::default() + }; + let results = idx.find_related(&chunks[0], &opts); + assert_eq!(results.len(), 1); + assert_eq!(results[0].chunk.file_path, "b.rs"); +} + +#[test] +fn find_related_without_seed_language_searches_whole_index() { + let chunks = vec![ + make_chunk("a.txt", 1, 3, None, "seed"), + make_chunk("b.rs", 1, 3, Some("rust"), "fn b() {}"), + make_chunk("c.py", 1, 3, Some("python"), "def c(): pass"), + ]; + let idx = build_index(chunks.clone()); + let opts = QueryOptions { + top_k: Some(5), + ..Default::default() + }; + let results = idx.find_related(&chunks[0], &opts); + let mut paths: Vec<_> = results.iter().map(|r| r.chunk.file_path.as_str()).collect(); + paths.sort_unstable(); + assert_eq!(paths, vec!["b.rs", "c.py"]); +} + +#[test] +fn find_related_with_empty_seed_language_searches_whole_index() { + let chunks = vec![ + make_chunk("a.txt", 1, 3, Some(""), "seed"), + make_chunk("b.rs", 1, 3, Some("rust"), "fn b() {}"), + make_chunk("c.py", 1, 3, Some("python"), "def c(): pass"), + ]; + let idx = build_index(chunks.clone()); + let opts = QueryOptions { + top_k: Some(5), + ..Default::default() + }; + let results = idx.find_related(&chunks[0], &opts); + let mut paths: Vec<_> = results.iter().map(|r| r.chunk.file_path.as_str()).collect(); + paths.sort_unstable(); + assert_eq!(paths, vec!["b.rs", "c.py"]); +} + +#[test] +fn find_related_returns_empty_when_seed_language_is_absent_from_index() { + let chunks = vec![make_chunk("a.rs", 1, 3, Some("rust"), "fn a() {}")]; + let idx = build_index(chunks); + let seed = make_chunk("b.py", 1, 3, Some("python"), "def b(): pass"); + let opts = QueryOptions { + top_k: Some(5), + ..Default::default() + }; + assert!(idx.find_related(&seed, &opts).is_empty()); +} + #[test] fn save_load_roundtrip() { let chunks = vec![