From a71aac8090dcd473e3cde98500d76d056948c4f4 Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 11:41:21 +0200 Subject: [PATCH 01/10] [#336] Avoid UTF-8 truncation panic in embedding chunks String::truncate panics when MAX_CHUNK_CHARS falls inside a multi-byte character, which a real embedding run hit on files with non-ASCII comments or string literals. Cut back to the nearest char boundary instead. --- src/embeddings/chunk.rs | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/src/embeddings/chunk.rs b/src/embeddings/chunk.rs index ad512a43..3be7de50 100644 --- a/src/embeddings/chunk.rs +++ b/src/embeddings/chunk.rs @@ -66,6 +66,21 @@ fn slice_lines(content: &str, start: usize, end: usize) -> String { /// headroom under that default rather than requiring users to raise it. const MAX_CHUNK_CHARS: usize = 1_600; +/// Truncate `text` to at most `limit` bytes without splitting a UTF-8 code +/// point. [`String::truncate`] panics when the boundary falls inside a +/// multi-byte character, which is easy to hit in files containing comments or +/// string literals outside ASCII. +fn truncate_on_char_boundary(text: &mut String, limit: usize) { + if text.len() <= limit { + return; + } + let mut cut = limit; + while cut > 0 && !text.is_char_boundary(cut) { + cut -= 1; + } + text.truncate(cut); +} + /// Extract embeddable chunks from one file's content. Returns an empty vector /// for unsupported languages or files with no extractable symbols. pub fn chunk_file( @@ -83,7 +98,7 @@ pub fn chunk_file( .filter_map(|node| { let (start, end) = parse_location(&node.source_location)?; let mut body = slice_lines(content, start, end); - body.truncate(MAX_CHUNK_CHARS); + truncate_on_char_boundary(&mut body, MAX_CHUNK_CHARS); // Prefix the symbol name so the embedding captures intent even when // the body is truncated or terse. let text = format!("{} in {}\n{}", node.label, node.source_file, body); @@ -116,4 +131,17 @@ mod tests { assert_eq!(slice_lines(src, 2, 4), "b\nc\nd"); assert_eq!(slice_lines(src, 1, 1), "a"); } + + #[test] + fn truncate_on_char_boundary_never_splits_a_code_point() { + // "é" is two bytes, so a 1_600-byte cut can land inside it. + let mut text = "a".repeat(1_599) + "é" + "tail"; + truncate_on_char_boundary(&mut text, 1_600); + assert_eq!(text.len(), 1_599); + assert!(text.is_char_boundary(text.len())); + + let mut short = "abc".to_string(); + truncate_on_char_boundary(&mut short, 1_600); + assert_eq!(short, "abc"); + } } From a9213acd9acffbbad9e4be360d4f6092b7cc4afa Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 11:41:21 +0200 Subject: [PATCH 02/10] [#336] Share workspace scan and hashing between search indexes Walk the workspace and compute per-file SHA-256 once per search in a shared source module. The embedding index consumes that scan instead of performing its own walk and hash pass; the lexical sidecar will use the same scan. --- src/embeddings/index.rs | 72 +++++++++------------- src/embeddings/mod.rs | 1 + src/embeddings/source.rs | 125 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 43 deletions(-) create mode 100644 src/embeddings/source.rs diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index 7e7da26b..c2da1baa 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -31,11 +31,11 @@ use std::sync::atomic::{AtomicBool, Ordering}; use anyhow::Result; use rayon::prelude::*; use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; use super::chunk::chunk_file; use super::client::EmbeddingClient; -use crate::graph::extract::{GraphExtractor, SupportedLanguage}; +use super::source::{SourceFile, scan_workspace}; +use crate::graph::extract::GraphExtractor; use crate::graph::store::GraphStore; const INDEX_VERSION: u32 = 1; @@ -221,6 +221,29 @@ impl EmbeddingIndex { compile_workers: usize, progress: &std::sync::atomic::AtomicU64, eta: &std::sync::atomic::AtomicU64, + ) -> Result { + let scan = scan_workspace(workspace); + Self::build_or_update_from_scan( + workspace, + &scan, + client, + cancel, + compile_workers, + progress, + eta, + ) + .await + } + + /// [`Self::build_or_update`] over a scan shared with the lexical sidecar. + pub async fn build_or_update_from_scan( + workspace: &Path, + scan: &[SourceFile], + client: &EmbeddingClient, + cancel: &AtomicBool, + compile_workers: usize, + progress: &std::sync::atomic::AtomicU64, + eta: &std::sync::atomic::AtomicU64, ) -> Result { let started = std::time::Instant::now(); let mut previous = Self::load(workspace); @@ -229,7 +252,6 @@ impl EmbeddingIndex { // (so the saved `meta.json` lists every file once the local phase is // done); an unchanged file that still has cached chunks is reused, the // rest are queued for (re-)embedding. - let files = collect_source_files(workspace); let mut new_hashes: HashMap = HashMap::new(); let mut kept: Vec = Vec::new(); let mut stale_files: Vec<(PathBuf, String, String)> = Vec::new(); @@ -243,16 +265,9 @@ impl EmbeddingIndex { .push(chunk); } - for path in files { - let rel = path - .strip_prefix(workspace) - .unwrap_or(&path) - .to_string_lossy() - .replace('\\', "/"); - let Ok(content) = std::fs::read_to_string(&path) else { - continue; - }; - let hash = sha256(&content); + for file in scan { + let rel = file.rel.clone(); + let hash = file.hash.clone(); new_hashes.insert(rel.clone(), hash.clone()); // Reuse only when unchanged AND the chunks are actually cached — a // file that was analysed but not embedded (no chunks) is re-embedded. @@ -261,7 +276,7 @@ impl EmbeddingIndex { { kept.extend(chunks.iter().cloned()); } else { - stale_files.push((path, rel, hash)); + stale_files.push((file.path.clone(), rel, hash)); } } @@ -708,35 +723,6 @@ fn append_manifest<'a>(path: &Path, files: impl Iterator) { } } -/// Compute the sha256 hex digest of `content`. -fn sha256(content: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(content.as_bytes()); - hasher - .finalize() - .iter() - .map(|b| format!("{b:02x}")) - .collect() -} - -/// Collect the workspace's source files, honouring `.gitignore`, restricted to -/// languages the Tree-sitter extractor understands. -fn collect_source_files(workspace: &Path) -> Vec { - ignore::WalkBuilder::new(workspace) - .hidden(true) - .git_ignore(true) - .build() - .filter_map(|entry| { - let path = entry.ok()?.into_path(); - if !path.is_file() { - return None; - } - let ext = path.extension().and_then(|e| e.to_str())?; - SupportedLanguage::from_extension(ext).map(|_| path) - }) - .collect() -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/embeddings/mod.rs b/src/embeddings/mod.rs index 74d76d86..8ede1073 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -29,6 +29,7 @@ pub mod chunk; pub mod client; pub mod index; pub mod sketch; +pub mod source; use crate::config::LlmConfiguration; use crate::llm::normalized_openai_endpoint; diff --git a/src/embeddings/source.rs b/src/embeddings/source.rs new file mode 100644 index 00000000..13a990db --- /dev/null +++ b/src/embeddings/source.rs @@ -0,0 +1,125 @@ +// Copyright (C) 2026 The orangu community +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +//! Shared workspace source scan. +//! +//! The embedding index and the lexical sidecar need the same file set and the +//! same per-file SHA-256 map. This module computes both once per search so the +//! workspace is not walked and hashed twice. + +use std::path::{Path, PathBuf}; + +use sha2::{Digest, Sha256}; + +use crate::graph::extract::SupportedLanguage; + +/// One readable source file, with its workspace-relative path and content hash. +#[derive(Debug, Clone)] +pub struct SourceFile { + pub path: PathBuf, + pub rel: String, + pub hash: String, +} + +/// Walk the workspace and hash every supported source file once. +/// +/// Files that cannot be read as UTF-8 are skipped, matching the previous +/// per-index behaviour. Discovery honours `.gitignore` and is restricted to the +/// languages the Tree-sitter extractor understands. +pub fn scan_workspace(workspace: &Path) -> Vec { + collect_source_files(workspace) + .into_iter() + .filter_map(|path| { + let rel = path + .strip_prefix(workspace) + .unwrap_or(&path) + .to_string_lossy() + .replace('\\', "/"); + let content = std::fs::read_to_string(&path).ok()?; + Some(SourceFile { + path, + rel, + hash: sha256(&content), + }) + }) + .collect() +} + +/// Compute the sha256 hex digest of `content`. +pub fn sha256(content: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(content.as_bytes()); + hasher + .finalize() + .iter() + .map(|b| format!("{b:02x}")) + .collect() +} + +/// Collect the workspace's source files, honouring `.gitignore`, restricted to +/// languages the Tree-sitter extractor understands. +fn collect_source_files(workspace: &Path) -> Vec { + ignore::WalkBuilder::new(workspace) + .hidden(true) + .git_ignore(true) + .build() + .filter_map(|entry| { + let path = entry.ok()?.into_path(); + if !path.is_file() { + return None; + } + let ext = path.extension().and_then(|e| e.to_str())?; + SupportedLanguage::from_extension(ext).map(|_| path) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sha256_is_stable_and_content_sensitive() { + assert_eq!(sha256("abc"), sha256("abc")); + assert_ne!(sha256("abc"), sha256("abcd")); + } + + #[test] + fn scan_hashes_supported_files_with_relative_paths() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(workspace.path().join("src")).unwrap(); + std::fs::write(workspace.path().join("src/lib.rs"), "fn one() {}\n").unwrap(); + std::fs::write(workspace.path().join("notes.txt"), "not source\n").unwrap(); + + let scan = scan_workspace(workspace.path()); + assert_eq!(scan.len(), 1); + assert_eq!(scan[0].rel, "src/lib.rs"); + assert_eq!(scan[0].hash, sha256("fn one() {}\n")); + } + + #[test] + fn scan_honours_gitignore() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::create_dir(workspace.path().join(".git")).unwrap(); + std::fs::write(workspace.path().join(".gitignore"), "ignored.rs\n").unwrap(); + std::fs::write(workspace.path().join("kept.rs"), "fn kept() {}\n").unwrap(); + std::fs::write(workspace.path().join("ignored.rs"), "fn ignored() {}\n").unwrap(); + + let scan = scan_workspace(workspace.path()); + let files: Vec<_> = scan.iter().map(|file| file.rel.as_str()).collect(); + assert!(files.contains(&"kept.rs")); + assert!(!files.contains(&"ignored.rs")); + } +} From 3de2b94afadaeaa409b91929d13a247975a17e5b Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 13:07:01 +0200 Subject: [PATCH 03/10] [#336] Add versioned lexical postings sidecar Index the full extracted Tree-sitter symbol range into a separately versioned lexical.json sidecar next to the embedding cache. Persist postings and symbol metadata only, keep the embedding truncation unchanged, and update incrementally from the shared workspace scan. A file hash is recorded only after that file parses successfully. --- src/embeddings/chunk.rs | 42 ++- src/embeddings/lexical.rs | 568 ++++++++++++++++++++++++++++++++++++++ src/embeddings/mod.rs | 1 + 3 files changed, 606 insertions(+), 5 deletions(-) create mode 100644 src/embeddings/lexical.rs diff --git a/src/embeddings/chunk.rs b/src/embeddings/chunk.rs index 3be7de50..de3c5ce0 100644 --- a/src/embeddings/chunk.rs +++ b/src/embeddings/chunk.rs @@ -36,8 +36,12 @@ pub struct Chunk { pub start_line: usize, /// 1-indexed inclusive end line. pub end_line: usize, - /// The source text embedded for this chunk (symbol name + body). + /// The source text embedded for this chunk (symbol name + truncated body). pub text: String, + /// The full, untruncated symbol body. Kept separate from [`Self::text`] so + /// the lexical index can cover the whole extracted symbol range without + /// changing what gets embedded. + pub body: String, } /// Parse the `L-L` location string the extractor records into a @@ -82,12 +86,36 @@ fn truncate_on_char_boundary(text: &mut String, limit: usize) { } /// Extract embeddable chunks from one file's content. Returns an empty vector -/// for unsupported languages or files with no extractable symbols. +/// for unsupported languages or files with no extractable symbols. Bodies are +/// truncated to [`MAX_CHUNK_CHARS`] for embedding; [`lexical_chunks`] keeps the +/// full extracted symbol range for lexical indexing. pub fn chunk_file( extractor: &GraphExtractor, path: &std::path::Path, rel_path: &str, content: &str, +) -> Vec { + chunks_with_body_limit(extractor, path, rel_path, content, Some(MAX_CHUNK_CHARS)) +} + +/// Extract chunks with the full symbol body, for the lexical index. The +/// embedding index intentionally keeps its own truncation, so the two caches +/// share extraction but not the indexed text boundary. +pub fn lexical_chunks( + extractor: &GraphExtractor, + path: &std::path::Path, + rel_path: &str, + content: &str, +) -> Vec { + chunks_with_body_limit(extractor, path, rel_path, content, None) +} + +fn chunks_with_body_limit( + extractor: &GraphExtractor, + path: &std::path::Path, + rel_path: &str, + content: &str, + body_limit: Option, ) -> Vec { let Ok((nodes, _edges)) = extractor.extract_from_file(path, rel_path, content) else { return Vec::new(); @@ -97,11 +125,14 @@ pub fn chunk_file( .into_iter() .filter_map(|node| { let (start, end) = parse_location(&node.source_location)?; - let mut body = slice_lines(content, start, end); - truncate_on_char_boundary(&mut body, MAX_CHUNK_CHARS); + let body = slice_lines(content, start, end); + let mut embedded_body = body.clone(); + if let Some(limit) = body_limit { + truncate_on_char_boundary(&mut embedded_body, limit); + } // Prefix the symbol name so the embedding captures intent even when // the body is truncated or terse. - let text = format!("{} in {}\n{}", node.label, node.source_file, body); + let text = format!("{} in {}\n{}", node.label, node.source_file, embedded_body); Some(Chunk { id: node.id, symbol: node.label, @@ -109,6 +140,7 @@ pub fn chunk_file( start_line: start, end_line: end, text, + body, }) }) .collect() diff --git a/src/embeddings/lexical.rs b/src/embeddings/lexical.rs new file mode 100644 index 00000000..2a40251e --- /dev/null +++ b/src/embeddings/lexical.rs @@ -0,0 +1,568 @@ +// Copyright (C) 2026 The orangu community +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +//! The lexical BM25 sidecar used by hybrid search. +//! +//! The sidecar lives next to the embedding cache as `lexical.json` and has its +//! own version, so tokenizer or postings changes never force a re-embedding. +//! It indexes the full extracted Tree-sitter symbol range (the embedding index +//! keeps its own 1,600-character truncation) and shares file discovery and +//! SHA-256 computation with [`super::index::EmbeddingIndex`], while tracking +//! which files *it* has successfully processed independently. +//! +//! Postings and per-document lengths are persisted rather than source text. + +use std::collections::{HashMap, HashSet}; +use std::path::{Path, PathBuf}; + +use anyhow::Result; +use rayon::prelude::*; +use serde::{Deserialize, Serialize}; + +use super::chunk::lexical_chunks; +use super::index::EmbeddingIndex; +use super::source::{SourceFile, scan_workspace}; +use crate::graph::extract::GraphExtractor; + +/// Independent sidecar version. Bump when the tokenizer, postings format, or +/// BM25 fields change — never when the embedding format changes. +pub const LEXICAL_INDEX_VERSION: u32 = 1; + +/// Token repetition weights approximate field weighting without a multi-field +/// BM25 implementation: a symbol name counts more than a file path, which in +/// turn counts more than the body. +const NAME_WEIGHT: usize = 3; +const PATH_WEIGHT: usize = 2; + +/// One indexed symbol. The line range lets query-time verification inspect only +/// the symbol's source rather than the whole file. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LexicalDoc { + pub id: String, + pub symbol: String, + pub file: String, + pub start_line: usize, + pub end_line: usize, + /// Weighted token count, used for BM25 length normalisation. + pub length: u32, +} + +/// One term occurrence in one document. Serialised as `[doc, tf]` to keep the +/// sidecar compact. +#[derive(Debug, Clone, Copy, Serialize, Deserialize)] +pub struct Posting(pub u32, pub u32); + +/// The persisted lexical index: version, per-file hashes for incremental +/// updates, document metadata, and the term postings. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct LexicalIndex { + #[serde(default)] + version: u32, + #[serde(default)] + file_hashes: HashMap, + #[serde(default)] + docs: Vec, + #[serde(default)] + postings: HashMap>, + #[serde(default)] + avg_length: f32, +} + +/// One document together with its term frequencies, used while building. +type DocWithTerms = (LexicalDoc, Vec<(String, u32)>); + +/// One parsed file's documents. +type ParsedFile = Vec; + +impl LexicalIndex { + /// The lexical sidecar path, next to the embedding index: + /// + /// ```text + /// ~/.orangu/workspace//embeddings/lexical.json + /// ``` + pub fn cache_path(workspace: &Path) -> PathBuf { + EmbeddingIndex::cache_dir(workspace).join("lexical.json") + } + + /// Load the sidecar, or an empty index when it is missing or stale. + pub fn load(workspace: &Path) -> Self { + let Ok(text) = std::fs::read_to_string(Self::cache_path(workspace)) else { + return Self::default(); + }; + let Ok(index) = serde_json::from_str::(&text) else { + return Self::default(); + }; + if index.version != LEXICAL_INDEX_VERSION { + return Self::default(); + } + index + } + + /// Number of indexed symbols. + pub fn len(&self) -> usize { + self.docs.len() + } + + /// Whether the index holds no symbols. + pub fn is_empty(&self) -> bool { + self.docs.is_empty() + } + + /// Build or incrementally update the lexical index for `workspace`. + /// + /// Only files whose content hash changed are re-parsed and re-tokenised; + /// unchanged documents and their term frequencies are carried over. The + /// sidecar is then rewritten in full because postings are positional and + /// corpus statistics change whenever any file changes. + pub fn build_or_update(workspace: &Path, compile_workers: usize) -> Result { + let scan = scan_workspace(workspace); + Self::build_or_update_from_scan(workspace, &scan, compile_workers, None) + } + + /// [`Self::build_or_update`] over a scan shared with the embedding index, + /// optionally publishing progress (permille) for the lexical-only path, + /// where no embedding build follows to report it. + pub fn build_or_update_with_progress( + workspace: &Path, + compile_workers: usize, + progress: Option<&std::sync::atomic::AtomicU64>, + ) -> Result { + let scan = scan_workspace(workspace); + Self::build_or_update_from_scan(workspace, &scan, compile_workers, progress) + } + + /// Build or update the sidecar from an existing workspace scan. + pub fn build_or_update_from_scan( + workspace: &Path, + scan: &[SourceFile], + compile_workers: usize, + progress: Option<&std::sync::atomic::AtomicU64>, + ) -> Result { + use std::sync::atomic::Ordering; + let report = |permille: u64| { + if let Some(progress) = progress { + progress.store(permille.min(1000), Ordering::Relaxed); + } + }; + report(0); + + let previous = Self::load(workspace); + let previous_terms = previous.terms_by_index(); + let mut previous_by_file: HashMap<&str, Vec> = HashMap::new(); + for (index, doc) in previous.docs.iter().enumerate() { + previous_by_file + .entry(doc.file.as_str()) + .or_default() + .push(index); + } + + let current_files: HashSet<&str> = scan.iter().map(|file| file.rel.as_str()).collect(); + let mut new_hashes: HashMap = HashMap::new(); + let mut kept: Vec = Vec::new(); + let mut stale_files: Vec = Vec::new(); + + for file in scan { + let unchanged = previous.file_hashes.get(&file.rel) == Some(&file.hash); + if unchanged { + // Keep every document previously emitted for this file. A file + // with no symbols legitimately keeps none. Terms are carried by + // document *position*, not id: one file can emit two symbols + // with the same id (repeated method names), and keying by id + // would give both copies the same term frequencies. + new_hashes.insert(file.rel.clone(), file.hash.clone()); + if let Some(indices) = previous_by_file.get(file.rel.as_str()) { + for &index in indices { + let doc = previous.docs[index].clone(); + let terms = previous_terms[index].clone(); + kept.push((doc, terms)); + } + } + } else { + stale_files.push(file.clone()); + } + } + report(250); + + let removed_files = previous + .file_hashes + .keys() + .any(|file| !current_files.contains(file.as_str())); + if stale_files.is_empty() && !removed_files && previous.version == LEXICAL_INDEX_VERSION { + report(1000); + return Ok(previous); + } + + let threads = if compile_workers > 0 { + compile_workers + } else { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + }; + let extractor = GraphExtractor::new()?; + let analysed = std::sync::atomic::AtomicUsize::new(0); + let total_files = stale_files.len().max(1); + let parse_one = |file: &SourceFile| { + let content = std::fs::read_to_string(&file.path).ok()?; + let chunks = lexical_chunks(&extractor, &file.path, &file.rel, &content); + let docs: ParsedFile = chunks + .into_iter() + .map(|chunk| { + let (terms, length) = doc_terms(&chunk.symbol, &chunk.file, &chunk.body); + ( + LexicalDoc { + id: chunk.id, + symbol: chunk.symbol, + file: chunk.file, + start_line: chunk.start_line, + end_line: chunk.end_line, + length, + }, + terms, + ) + }) + .collect(); + let done = analysed.fetch_add(1, Ordering::Relaxed) + 1; + report(250 + (done as u64 * 650 / total_files as u64).min(650)); + Some((file.rel.clone(), file.hash.clone(), docs)) + }; + let parsed: Vec<(String, String, ParsedFile)> = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build()? + .install(|| stale_files.par_iter().filter_map(parse_one).collect()); + + // Record a hash only after that file parsed successfully, so a failed + // parse is retried on the next run instead of being treated as done. + for (rel, hash, _) in &parsed { + new_hashes.insert(rel.clone(), hash.clone()); + } + + let mut docs_with_terms = kept; + for (_, _, file_docs) in parsed { + docs_with_terms.extend(file_docs); + } + + let mut docs = Vec::with_capacity(docs_with_terms.len()); + let mut postings: HashMap> = HashMap::new(); + let mut total_length: u64 = 0; + for (doc_index, (doc, terms)) in docs_with_terms.into_iter().enumerate() { + total_length += doc.length as u64; + for (term, tf) in terms { + postings + .entry(term) + .or_default() + .push(Posting(doc_index as u32, tf)); + } + docs.push(doc); + } + // Deterministic sidecar output: postings order is already stable for a + // given document order, but sorting each list by doc makes it explicit. + for list in postings.values_mut() { + list.sort_by_key(|posting| posting.0); + } + + let index = Self { + version: LEXICAL_INDEX_VERSION, + file_hashes: new_hashes, + avg_length: if docs.is_empty() { + 0.0 + } else { + total_length as f32 / docs.len() as f32 + }, + docs, + postings, + }; + index.save(workspace)?; + report(1000); + Ok(index) + } + + /// Rewrite the sidecar in full. + fn save(&self, workspace: &Path) -> Result<()> { + let path = Self::cache_path(workspace); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let json = serde_json::to_string(self)?; + std::fs::write(path, json)?; + Ok(()) + } + + /// Reconstruct per-document term frequencies from the postings so unchanged + /// files can be carried over without re-tokenising them. Indexed by document + /// position, because ids are not unique within a file. + fn terms_by_index(&self) -> Vec> { + let mut terms: Vec> = vec![Vec::new(); self.docs.len()]; + for (term, postings) in &self.postings { + for posting in postings { + if let Some(slot) = terms.get_mut(posting.0 as usize) { + slot.push((term.clone(), posting.1)); + } + } + } + terms + } +} + +/// Term frequencies for one document, with symbol name and path weighted above +/// the body. The full body is indexed, not the embedding truncation. +fn doc_terms(symbol: &str, file: &str, body: &str) -> (Vec<(String, u32)>, u32) { + let mut weighted = Vec::new(); + push_weighted(&mut weighted, &tokenize(symbol), NAME_WEIGHT); + push_weighted(&mut weighted, &tokenize(file), PATH_WEIGHT); + push_weighted(&mut weighted, &tokenize(body), 1); + + let length = weighted.len() as u32; + let mut counts: HashMap = HashMap::new(); + for token in weighted { + *counts.entry(token).or_default() += 1; + } + let mut terms: Vec<(String, u32)> = counts.into_iter().collect(); + terms.sort_by(|a, b| a.0.cmp(&b.0)); + (terms, length) +} + +fn push_weighted(out: &mut Vec, tokens: &[String], weight: usize) { + for _ in 0..weight { + out.extend(tokens.iter().cloned()); + } +} + +/// Tokenise text for BM25: keep each original identifier-like span (so an exact +/// identifier is its own token) and add lowercase separator/camel-case parts. +fn tokenize(text: &str) -> Vec { + let mut tokens = Vec::new(); + for span in spans(text) { + tokens.extend(span_tokens(&span)); + } + tokens +} + +fn spans(text: &str) -> Vec { + let mut spans = Vec::new(); + let mut current = String::new(); + for character in text.chars() { + if character.is_alphanumeric() + || matches!(character, '_' | ':' | '.' | '/' | '\\' | '$' | '-') + { + current.push(character); + } else { + push_span(&mut spans, &mut current); + } + } + push_span(&mut spans, &mut current); + spans +} + +fn push_span(spans: &mut Vec, current: &mut String) { + let trimmed = current.trim_matches(|c| matches!(c, '.' | '/' | ':' | '-' | '\\')); + if !trimmed.is_empty() { + spans.push(trimmed.to_string()); + } + current.clear(); +} + +fn span_tokens(span: &str) -> Vec { + let mut tokens = Vec::new(); + let mut seen = HashSet::new(); + push_unique(&mut tokens, &mut seen, span.to_lowercase()); + for segment in split_segments(span) { + push_unique(&mut tokens, &mut seen, segment.to_lowercase()); + for part in camel_parts(&segment) { + push_unique(&mut tokens, &mut seen, part); + } + } + tokens +} + +fn push_unique(tokens: &mut Vec, seen: &mut HashSet, token: String) { + if !token.is_empty() && seen.insert(token.clone()) { + tokens.push(token); + } +} + +fn split_segments(span: &str) -> Vec { + span.split(['_', ':', '.', '/', '\\', '-']) + .filter(|segment| !segment.is_empty()) + .map(str::to_string) + .collect() +} + +/// Split `GraphStore` into `graph`, `store` and `HTTP2Server` into `http`, `2`, +/// `server`. +fn camel_parts(segment: &str) -> Vec { + let chars: Vec = segment.chars().collect(); + if chars.is_empty() { + return Vec::new(); + } + let mut parts = Vec::new(); + let mut start = 0; + for index in 1..chars.len() { + let previous = chars[index - 1]; + let current = chars[index]; + let next = chars.get(index + 1); + let boundary = (previous.is_lowercase() && current.is_uppercase()) + || (previous.is_uppercase() + && current.is_uppercase() + && next.is_some_and(|c| c.is_lowercase())) + || (previous.is_alphabetic() && current.is_numeric()) + || (previous.is_numeric() && current.is_alphabetic()); + if boundary { + parts.push( + chars[start..index] + .iter() + .collect::() + .to_lowercase(), + ); + start = index; + } + } + parts.push(chars[start..].iter().collect::().to_lowercase()); + parts.retain(|part| !part.is_empty()); + parts +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tokenizer_keeps_full_identifier_and_parts() { + let tokens = tokenize("EMBED_BATCH_TOKEN_BUDGET"); + assert!(tokens.contains(&"embed_batch_token_budget".to_string())); + assert!(tokens.contains(&"embed".to_string())); + assert!(tokens.contains(&"budget".to_string())); + } + + #[test] + fn tokenizer_splits_camel_case_and_paths() { + let tokens = tokenize("GraphStore::lookup src/embeddings/index.rs"); + assert!(tokens.contains(&"graphstore::lookup".to_string())); + assert!(tokens.contains(&"graph".to_string())); + assert!(tokens.contains(&"store".to_string())); + assert!(tokens.contains(&"embeddings".to_string())); + assert!(tokens.contains(&"index".to_string())); + } + + #[test] + fn sidecar_round_trips_and_persists_no_source_text() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn retry_backoff() {\n let marker = \"secret_body_text\";\n}\n", + ) + .unwrap(); + + let built = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + assert_eq!(built.len(), 1); + assert_eq!(built.docs[0].symbol, "retry_backoff"); + assert!(built.postings.contains_key("retry_backoff")); + + let loaded = LexicalIndex::load(workspace.path()); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded.docs[0].id, built.docs[0].id); + + let raw = std::fs::read_to_string(LexicalIndex::cache_path(workspace.path())).unwrap(); + // Postings store individual terms, never the source line or function body. + assert!(!raw.contains("let marker")); + assert!(!raw.contains("pub fn retry_backoff")); + } + + #[test] + fn version_mismatch_invalidates_the_sidecar() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(EmbeddingIndex::cache_dir(workspace.path())).unwrap(); + std::fs::write( + LexicalIndex::cache_path(workspace.path()), + r#"{"version":999,"file_hashes":{},"docs":[],"postings":{},"avg_length":0.0}"#, + ) + .unwrap(); + assert!(LexicalIndex::load(workspace.path()).is_empty()); + } + + #[test] + fn changed_and_removed_files_update_only_their_documents() { + let workspace = tempfile::tempdir().unwrap(); + let first = workspace.path().join("a.rs"); + let second = workspace.path().join("b.rs"); + std::fs::write(&first, "pub fn one() {}\n").unwrap(); + std::fs::write(&second, "pub fn two() {}\n").unwrap(); + + let built = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + assert_eq!(built.len(), 2); + + std::fs::write(&first, "pub fn three() {}\n").unwrap(); + let updated = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let symbols: Vec<_> = updated.docs.iter().map(|doc| doc.symbol.as_str()).collect(); + assert!(symbols.contains(&"three")); + assert!(symbols.contains(&"two")); + assert!(!symbols.contains(&"one")); + + std::fs::remove_file(&second).unwrap(); + let updated = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + assert_eq!(updated.len(), 1); + assert_eq!(updated.docs[0].symbol, "three"); + } + + #[test] + fn warm_rebuild_keeps_duplicate_id_documents_distinct() { + let workspace = tempfile::tempdir().unwrap(); + // Two methods with the same name in one file share an id, but not a + // body, so they must keep their own weights across a warm rebuild. + std::fs::write( + workspace.path().join("a.rs"), + "impl A {\n pub fn run(&self) {\n one();\n }\n}\n\n\ + impl B {\n pub fn run(&self) {\n two(); three(); four(); five(); six();\n }\n}\n", + ) + .unwrap(); + let cold = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let runs = cold.docs.iter().filter(|doc| doc.symbol == "run").count(); + assert_eq!(runs, 2, "both `run` methods should be indexed"); + let mut cold_lengths: Vec = cold.docs.iter().map(|doc| doc.length).collect(); + cold_lengths.sort_unstable(); + + let warm = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let mut warm_lengths: Vec = warm.docs.iter().map(|doc| doc.length).collect(); + warm_lengths.sort_unstable(); + assert_eq!(cold_lengths, warm_lengths); + } + + #[test] + fn unchanged_file_reuses_cached_documents() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write(workspace.path().join("a.rs"), "pub fn one() {}\n").unwrap(); + let built = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + + let updated = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + assert_eq!(updated.len(), built.len()); + assert_eq!(updated.docs[0].id, built.docs[0].id); + assert_eq!(updated.docs[0].length, built.docs[0].length); + } + + #[test] + fn lexical_index_covers_the_full_symbol_range() { + let workspace = tempfile::tempdir().unwrap(); + let body = format!( + "pub fn long() {{\n let _ = \"{}\";\n}}\n", + "word ".repeat(2_000) + ); + std::fs::write(workspace.path().join("a.rs"), body).unwrap(); + + let index = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + assert_eq!(index.len(), 1); + assert!(index.docs[0].length > 500); + assert!(index.postings.contains_key("word")); + } +} diff --git a/src/embeddings/mod.rs b/src/embeddings/mod.rs index 8ede1073..ab2d67ca 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -28,6 +28,7 @@ pub mod chunk; pub mod client; pub mod index; +pub mod lexical; pub mod sketch; pub mod source; From 3e7fd72e33b03063c5660bc223523210b2606043 Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 13:13:29 +0200 Subject: [PATCH 04/10] [#336] Add BM25 lexical-only retrieval Rank the lexical sidecar with a code-aware tokenizer and BM25, expand graph neighbours, and apply the existing token budget. With no embeddings-capable server, /search and semantic_search now return ranked lexical results instead of an error. --- src/bin/orangu/dispatch.rs | 66 +++++++------ src/embeddings/index.rs | 48 +++++++++- src/embeddings/lexical.rs | 189 ++++++++++++++++++++++++++++++++++++- src/embeddings/mod.rs | 53 ++++++++++- src/tools.rs | 66 +++++++++---- 5 files changed, 363 insertions(+), 59 deletions(-) diff --git a/src/bin/orangu/dispatch.rs b/src/bin/orangu/dispatch.rs index 8def9036..5dedf724 100644 --- a/src/bin/orangu/dispatch.rs +++ b/src/bin/orangu/dispatch.rs @@ -2149,42 +2149,40 @@ pub(crate) fn handle_command( "Usage: /search ".to_string(), )), LocalCommand::Search(Some(query)) => { - // Semantic search is enabled only when an embeddings-capable endpoint - // was detected at startup; `embeddings_server` holds that server's - // name, or is empty when none responded. - if embeddings_server.is_empty() { - return Ok(CommandOutcome::OutputError( - "Semantic /search is unavailable: no embeddings-capable server was detected \ - at startup. Give a server section `role = embeddings` (or load an embedding \ - model on your default server) and restart. Meanwhile, use /grep or the \ - knowledge graph." - .to_string(), - )); - } - let Some(profile) = llms.get(embeddings_server) else { - return Ok(CommandOutcome::OutputError(format!( - "The embeddings server '{embeddings_server}' is not a configured server \ - section." - ))); - }; - // Under a confirmed coordinator, it alone owns model/role - // decisions: force `.model = "embeddings"` on the active - // connection rather than sending whatever it's normally - // configured with (e.g. "all"), which — since a coordinator - // matches an exact model/role name before it ever considers the - // request path — could otherwise route this to the wrong - // backend entirely. - let search = match orangu::embeddings::SemanticSearch::from_profile( - profile, - is_coordinator, - compile_workers, - semantic_budget_tokens, - ) { - Ok(search) => search, - Err(err) => { + // `embeddings_server` holds the embeddings-capable server detected + // at startup, or is empty when none responded. Either way `/search` + // works: without embeddings it ranks with the lexical BM25 sidecar. + let search = if embeddings_server.is_empty() { + orangu::embeddings::SemanticSearch::lexical_only( + compile_workers, + semantic_budget_tokens, + ) + } else { + let Some(profile) = llms.get(embeddings_server) else { return Ok(CommandOutcome::OutputError(format!( - "Could not initialise embeddings client: {err:#}" + "The embeddings server '{embeddings_server}' is not a configured server \ + section." ))); + }; + // Under a confirmed coordinator, it alone owns model/role + // decisions: force `.model = "embeddings"` on the active + // connection rather than sending whatever it's normally + // configured with (e.g. "all"), which — since a coordinator + // matches an exact model/role name before it ever considers the + // request path — could otherwise route this to the wrong + // backend entirely. + match orangu::embeddings::SemanticSearch::from_profile( + profile, + is_coordinator, + compile_workers, + semantic_budget_tokens, + ) { + Ok(search) => search, + Err(err) => { + return Ok(CommandOutcome::OutputError(format!( + "Could not initialise embeddings client: {err:#}" + ))); + } } }; let workspace = workspace.to_path_buf(); diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index c2da1baa..a448ce0d 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -34,6 +34,7 @@ use serde::{Deserialize, Serialize}; use super::chunk::chunk_file; use super::client::EmbeddingClient; +use super::lexical::LexicalDoc; use super::source::{SourceFile, scan_workspace}; use crate::graph::extract::GraphExtractor; use crate::graph::store::GraphStore; @@ -65,7 +66,7 @@ const EMBED_CONCURRENCY: usize = 8; /// Score multiplier applied to a chunk pulled in only because it neighbours a /// strong semantic hit in the knowledge graph. Below `1.0` so a graph-expanded /// result never outranks the semantic hit that surfaced it. -const EXPANSION_DECAY: f32 = 0.6; +pub(super) const EXPANSION_DECAY: f32 = 0.6; // ── On-disk format ──────────────────────────────────────────────────────────── @@ -604,7 +605,7 @@ impl EmbeddingIndex { /// Insert `hit` into `best` unless an equal-or-better score for the same id is /// already recorded. -fn insert_hit(best: &mut HashMap, hit: SearchHit) { +pub(super) fn insert_hit(best: &mut HashMap, hit: SearchHit) { match best.get(&hit.id) { Some(existing) if existing.score >= hit.score => {} _ => { @@ -625,6 +626,48 @@ fn hit_from(chunk: &EmbeddedChunk, score: f32, expanded_from: Option) -> } } +/// Build a hit from a lexical document, for chunks that have no embedding +/// vector (lexical-only search) or for graph neighbours pulled in by id. +pub(super) fn hit_from_doc( + doc: &LexicalDoc, + score: f32, + expanded_from: Option, +) -> SearchHit { + SearchHit { + id: doc.id.clone(), + symbol: doc.symbol.clone(), + file: doc.file.clone(), + start_line: doc.start_line, + end_line: doc.end_line, + score, + expanded_from, + } +} + +/// Apply the token budget to an already-ranked list, always keeping at least +/// one hit so a single large symbol is still returned. +pub(super) fn budget_cut( + hits: Vec, + top_k: usize, + semantic_budget_tokens: usize, +) -> Vec { + let mut total_tokens = 0; + let mut final_hits = Vec::new(); + for hit in hits { + let chunk_lines = hit.end_line.saturating_sub(hit.start_line) + 1; + let chunk_tokens = chunk_lines * 10; + if total_tokens + chunk_tokens > semantic_budget_tokens && !final_hits.is_empty() { + break; + } + total_tokens += chunk_tokens; + final_hits.push(hit); + if final_hits.len() >= top_k { + break; + } + } + final_hits +} + /// Cosine similarity of two vectors. Returns `0.0` for a length mismatch or a /// zero-magnitude vector. fn cosine(a: &[f32], b: &[f32]) -> f32 { @@ -735,6 +778,7 @@ mod tests { start_line: 1, end_line: 1, text: "x".repeat(chars), + body: "x".repeat(chars), } } diff --git a/src/embeddings/lexical.rs b/src/embeddings/lexical.rs index 2a40251e..186931bd 100644 --- a/src/embeddings/lexical.rs +++ b/src/embeddings/lexical.rs @@ -32,14 +32,18 @@ use rayon::prelude::*; use serde::{Deserialize, Serialize}; use super::chunk::lexical_chunks; -use super::index::EmbeddingIndex; +use super::index::{EmbeddingIndex, SearchHit, budget_cut, hit_from_doc, insert_hit}; use super::source::{SourceFile, scan_workspace}; use crate::graph::extract::GraphExtractor; +use crate::graph::store::GraphStore; /// Independent sidecar version. Bump when the tokenizer, postings format, or /// BM25 fields change — never when the embedding format changes. pub const LEXICAL_INDEX_VERSION: u32 = 1; +const BM25_K1: f32 = 1.2; +const BM25_B: f32 = 0.75; + /// Token repetition weights approximate field weighting without a multi-field /// BM25 implementation: a symbol name counts more than a file path, which in /// turn counts more than the body. @@ -80,6 +84,31 @@ pub struct LexicalIndex { avg_length: f32, } +/// One ranked lexical result, before graph expansion and budgeting. +#[derive(Debug, Clone)] +pub struct LexicalHit { + pub id: String, + pub symbol: String, + pub file: String, + pub start_line: usize, + pub end_line: usize, + pub score: f32, +} + +impl LexicalHit { + fn to_search_hit(&self, expanded_from: Option) -> SearchHit { + SearchHit { + id: self.id.clone(), + symbol: self.symbol.clone(), + file: self.file.clone(), + start_line: self.start_line, + end_line: self.end_line, + score: self.score, + expanded_from, + } + } +} + /// One document together with its term frequencies, used while building. type DocWithTerms = (LexicalDoc, Vec<(String, u32)>); @@ -314,6 +343,133 @@ impl LexicalIndex { } terms } + + /// Rank documents for `query` with BM25. + pub fn rank(&self, _workspace: &Path, query: &str) -> Vec { + if self.docs.is_empty() { + return Vec::new(); + } + let terms = query_terms(query); + if terms.is_empty() { + return Vec::new(); + } + let scores = self.bm25(&terms); + + let mut hits = Vec::new(); + for (index, score) in scores.into_iter().enumerate() { + if score <= 0.0 { + continue; + } + let doc = &self.docs[index]; + hits.push(LexicalHit { + id: doc.id.clone(), + symbol: doc.symbol.clone(), + file: doc.file.clone(), + start_line: doc.start_line, + end_line: doc.end_line, + score, + }); + } + hits.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.id.cmp(&b.id))); + hits + } + + /// Lexical-only retrieval: rank, expand graph neighbours, and apply the + /// token budget. Used when no embeddings server is available. + pub fn search( + &self, + workspace: &Path, + query: &str, + graph: Option<&GraphStore>, + top_k: usize, + semantic_budget_tokens: usize, + ) -> Vec { + let ranked = self.rank(workspace, query); + if ranked.is_empty() { + return Vec::new(); + } + let id_to_index: HashMap<&str, usize> = self + .docs + .iter() + .enumerate() + .map(|(index, doc)| (doc.id.as_str(), index)) + .collect(); + + let mut best: HashMap = HashMap::new(); + let seed_count = top_k.saturating_mul(2).max(top_k); + for hit in ranked.iter().take(seed_count) { + insert_hit(&mut best, hit.to_search_hit(None)); + } + + if let Some(graph) = graph { + for hit in ranked.iter().take(seed_count) { + for result in graph.lookup(&hit.symbol) { + if result.node.id != hit.id { + continue; + } + let neighbours = result.callers.iter().chain(result.callees.iter()); + for edge in neighbours { + if let Some(&index) = id_to_index.get(edge.node_id.as_str()) { + let doc = &self.docs[index]; + insert_hit( + &mut best, + hit_from_doc( + doc, + hit.score * super::index::EXPANSION_DECAY, + Some(hit.symbol.clone()), + ), + ); + } + } + } + } + } + + let mut hits: Vec = best.into_values().collect(); + hits.sort_by(|a, b| b.score.total_cmp(&a.score)); + budget_cut(hits, top_k, semantic_budget_tokens) + } + + /// BM25 score for every document. Returns a zero for documents containing + /// none of the query terms. + fn bm25(&self, query_terms: &[String]) -> Vec { + let mut scores = vec![0.0f32; self.docs.len()]; + let n = self.docs.len() as f32; + let avg_length = self.avg_length.max(1.0); + let mut seen: HashSet<&str> = HashSet::new(); + + for term in query_terms { + if !seen.insert(term.as_str()) { + continue; + } + let Some(postings) = self.postings.get(term) else { + continue; + }; + let df = postings.len() as f32; + let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln(); + for posting in postings { + let Some(doc) = self.docs.get(posting.0 as usize) else { + continue; + }; + let tf = posting.1 as f32; + let length = doc.length as f32; + let denominator = tf + BM25_K1 * (1.0 - BM25_B + BM25_B * length / avg_length); + if denominator > 0.0 { + scores[posting.0 as usize] += idf * (tf * (BM25_K1 + 1.0)) / denominator; + } + } + } + scores + } +} + +/// Tokenise a plain query into unique BM25 terms. +fn query_terms(query: &str) -> Vec { + let mut seen = HashSet::new(); + tokenize(query) + .into_iter() + .filter(|token| seen.insert(token.clone())) + .collect() } /// Term frequencies for one document, with symbol name and path weighted above @@ -565,4 +721,35 @@ mod tests { assert!(index.docs[0].length > 500); assert!(index.postings.contains_key("word")); } + + #[test] + fn bm25_ranks_exact_identifier_above_unrelated_docs() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn retry() {}\n\ + pub fn retry_backoff() {\n retry();\n}\n\ + pub fn other() {}\n", + ) + .unwrap(); + let index = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let ranked = index.rank(workspace.path(), "retry_backoff"); + assert_eq!(ranked.first().unwrap().symbol, "retry_backoff"); + assert!(!ranked.iter().any(|hit| hit.symbol == "other")); + } + + #[test] + fn lexical_search_returns_ranked_hits_without_vectors() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn retry_backoff() {\n sleep(1);\n}\n", + ) + .unwrap(); + let index = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let hits = index.search(workspace.path(), "retry_backoff", None, 5, 16_384); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].symbol, "retry_backoff"); + assert!(hits[0].score > 0.0); + } } diff --git a/src/embeddings/mod.rs b/src/embeddings/mod.rs index ab2d67ca..c4bed06a 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -42,7 +42,9 @@ pub use sketch::{EllipsoidSketch, SketchMatch}; /// Shared refresh-before-search implementation for `/search` and LLM tools. #[derive(Clone)] pub struct SemanticSearch { - client: EmbeddingClient, + /// `None` means lexical-only mode: no embeddings-capable server was + /// detected at startup, or the embedding half failed and was degraded. + client: Option, compile_workers: usize, budget_tokens: usize, } @@ -51,6 +53,9 @@ pub struct SemanticSearchResult { pub indexed_symbols: usize, pub hits: Vec, pub cancelled: bool, + /// Whether the ranking came from the hybrid/semantic path. `false` means + /// lexical-only fallback. + pub used_embeddings: bool, } impl SemanticSearch { @@ -65,12 +70,22 @@ impl SemanticSearch { profile.model = "embeddings".to_string(); } Ok(Self { - client: EmbeddingClient::from_profile(&profile)?, + client: Some(EmbeddingClient::from_profile(&profile)?), compile_workers, budget_tokens, }) } + /// Lexical-only search, for when no embeddings-capable server is available. + /// BM25 and graph expansion do not need an embeddings server. + pub fn lexical_only(compile_workers: usize, budget_tokens: usize) -> Self { + Self { + client: None, + compile_workers, + budget_tokens, + } + } + pub async fn search( &self, workspace: &std::path::Path, @@ -85,9 +100,38 @@ impl SemanticSearch { !query.trim().is_empty(), "semantic search query must not be empty" ); + let Some(client) = &self.client else { + // No embeddings-capable server was detected at startup. BM25 and + // graph expansion do not need one, so rank lexically instead of + // failing the search. + let lexical = crate::embeddings::lexical::LexicalIndex::build_or_update_with_progress( + workspace, + self.compile_workers, + Some(progress), + )?; + let mut result = SemanticSearchResult { + indexed_symbols: lexical.len(), + hits: Vec::new(), + cancelled: cancel.load(Ordering::Relaxed), + used_embeddings: false, + }; + if result.cancelled || lexical.is_empty() { + return Ok(result); + } + result.cancelled = cancel.load(Ordering::Relaxed); + if !result.cancelled { + let guard = graph + .lock() + .map_err(|_| anyhow::anyhow!("graph_store mutex poisoned"))?; + result.hits = + lexical.search(workspace, query, guard.as_ref(), 10, self.budget_tokens); + } + return Ok(result); + }; + let index = EmbeddingIndex::build_or_update( workspace, - &self.client, + client, cancel, self.compile_workers, progress, @@ -98,11 +142,12 @@ impl SemanticSearch { indexed_symbols: index.len(), hits: Vec::new(), cancelled: cancel.load(Ordering::Relaxed), + used_embeddings: true, }; if result.cancelled || index.is_empty() { return Ok(result); } - let query_vector = self.client.embed_one(query).await?; + let query_vector = client.embed_one(query).await?; result.cancelled = cancel.load(Ordering::Relaxed); if !result.cancelled { let guard = graph diff --git a/src/tools.rs b/src/tools.rs index 92e4c35b..a7f83fb9 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -429,16 +429,18 @@ impl ToolExecutor { compile_workers: usize, budget_tokens: usize, ) -> Result<()> { - let search = profile - .map(|profile| { - crate::embeddings::SemanticSearch::from_profile( - profile, - is_coordinator, - compile_workers, - budget_tokens, - ) - }) - .transpose()?; + let search = match profile { + Some(profile) => Some(crate::embeddings::SemanticSearch::from_profile( + profile, + is_coordinator, + compile_workers, + budget_tokens, + )?), + None => Some(crate::embeddings::SemanticSearch::lexical_only( + compile_workers, + budget_tokens, + )), + }; *self .semantic_search .lock() @@ -753,7 +755,7 @@ impl ToolExecutor { defs.push(tool( "semantic_search", - "Search workspace code by meaning. Refreshes the embedding index before searching; requires an embeddings server.", + "Search workspace code by meaning and exact tokens. Uses lexical ranking when no embeddings server is available.", json!({ "type": "object", "properties": { "query": { "type": "string", "minLength": 1 } }, @@ -4006,6 +4008,34 @@ mod semantic_search_tests { } } + #[tokio::test] + async fn semantic_search_works_without_an_embeddings_server() { + let workspace = tempfile::tempdir().unwrap(); + let _cache = CacheCleanup(crate::embeddings::EmbeddingIndex::cache_dir( + workspace.path(), + )); + fs::write( + workspace.path().join("retry.rs"), + "pub fn retry_backoff() {}\n", + ) + .unwrap(); + let tools = ToolExecutor::new_read_only(workspace.path()); + tools + .configure_semantic_search(None, false, 1, 16384) + .unwrap(); + let answer = tools + .execute( + "semantic_search", + json!({"query": "retry_backoff"}).as_object().unwrap(), + ) + .await + .unwrap(); + assert!( + answer.contains("retry_backoff") && answer.contains("retry.rs:1"), + "{answer}" + ); + } + #[tokio::test] async fn semantic_search_reports_unavailable_and_rejects_invalid_queries() { let workspace = tempfile::tempdir().unwrap(); @@ -4117,13 +4147,13 @@ mod semantic_search_tests { tools .configure_semantic_search(None, false, 1, 16384) .unwrap(); - assert!( - tools - .clone() - .execute("semantic_search", arguments.as_object().unwrap()) - .await - .is_err() - ); + // With no embeddings server, search degrades to lexical-only rather + // than failing outright. + let answer = tools + .execute("semantic_search", arguments.as_object().unwrap()) + .await + .unwrap(); + assert!(answer.contains("No semantic matches"), "{answer}"); server.abort(); } } From cb67e00027301ebe1b9e94247556927d2b60718f Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 13:15:56 +0200 Subject: [PATCH 05/10] [#336] Verify quoted phrases and exact identifiers against source ranges Intersect posting lists to narrow candidates, group them by file, read each file once, and verify phrase order/adjacency or identifier token boundaries only inside the symbol's line range. Verified matches are marked exact and ranked first. Verification is a reranking step, not a recall guarantee. --- src/embeddings/index.rs | 5 + src/embeddings/lexical.rs | 333 ++++++++++++++++++++++++++++++++++++-- src/embeddings/mod.rs | 4 +- 3 files changed, 330 insertions(+), 12 deletions(-) diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index a448ce0d..deae58a3 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -107,6 +107,8 @@ pub struct SearchHit { /// `None` for a direct semantic hit, or `Some(symbol)` when this chunk was /// pulled in because it neighbours that semantic hit in the graph. pub expanded_from: Option, + /// Verified as an exact identifier or quoted-phrase match in the source. + pub exact: bool, } /// A completed file's embedding result: its workspace-relative path, content @@ -623,6 +625,7 @@ fn hit_from(chunk: &EmbeddedChunk, score: f32, expanded_from: Option) -> end_line: chunk.end_line, score, expanded_from, + exact: false, } } @@ -632,6 +635,7 @@ pub(super) fn hit_from_doc( doc: &LexicalDoc, score: f32, expanded_from: Option, + exact: bool, ) -> SearchHit { SearchHit { id: doc.id.clone(), @@ -641,6 +645,7 @@ pub(super) fn hit_from_doc( end_line: doc.end_line, score, expanded_from, + exact, } } diff --git a/src/embeddings/lexical.rs b/src/embeddings/lexical.rs index 186931bd..2011c49d 100644 --- a/src/embeddings/lexical.rs +++ b/src/embeddings/lexical.rs @@ -44,6 +44,11 @@ pub const LEXICAL_INDEX_VERSION: u32 = 1; const BM25_K1: f32 = 1.2; const BM25_B: f32 = 0.75; +/// How many BM25 candidates to source-verify for a quoted phrase or explicit +/// identifier. Common phrases can match many documents; verifying a larger set +/// is cheap because each candidate file is read at most once per query. +const MAX_VERIFY_CANDIDATES: usize = 5_000; + /// Token repetition weights approximate field weighting without a multi-field /// BM25 implementation: a symbol name counts more than a file path, which in /// turn counts more than the body. @@ -93,6 +98,8 @@ pub struct LexicalHit { pub start_line: usize, pub end_line: usize, pub score: f32, + /// Verified as an exact identifier or quoted-phrase match in the source. + pub exact: bool, } impl LexicalHit { @@ -105,6 +112,7 @@ impl LexicalHit { end_line: self.end_line, score: self.score, expanded_from, + exact: self.exact, } } } @@ -344,20 +352,32 @@ impl LexicalIndex { terms } - /// Rank documents for `query` with BM25. - pub fn rank(&self, _workspace: &Path, query: &str) -> Vec { + /// Rank documents for `query` with BM25, giving verified exact identifier + /// and quoted-phrase matches priority. + pub fn rank(&self, workspace: &Path, query: &str) -> Vec { if self.docs.is_empty() { return Vec::new(); } - let terms = query_terms(query); - if terms.is_empty() { + let parsed = parse_query(query); + if parsed.terms.is_empty() { return Vec::new(); } - let scores = self.bm25(&terms); + let scores = self.bm25(&parsed.terms); + + let mut exact: HashSet = HashSet::new(); + if let Some(identifier) = &parsed.identifier { + let candidates = self.identifier_candidates(identifier, &scores); + exact.extend(self.verify(workspace, &candidates, &[], Some(identifier))); + } + for phrase in &parsed.phrases { + let candidates = self.phrase_candidates(phrase, &scores); + exact.extend(self.verify(workspace, &candidates, std::slice::from_ref(phrase), None)); + } let mut hits = Vec::new(); for (index, score) in scores.into_iter().enumerate() { - if score <= 0.0 { + let is_exact = exact.contains(&index); + if score <= 0.0 && !is_exact { continue; } let doc = &self.docs[index]; @@ -368,9 +388,15 @@ impl LexicalIndex { start_line: doc.start_line, end_line: doc.end_line, score, + exact: is_exact, }); } - hits.sort_by(|a, b| b.score.total_cmp(&a.score).then_with(|| a.id.cmp(&b.id))); + hits.sort_by(|a, b| { + b.exact + .cmp(&a.exact) + .then_with(|| b.score.total_cmp(&a.score)) + .then_with(|| a.id.cmp(&b.id)) + }); hits } @@ -417,6 +443,7 @@ impl LexicalIndex { doc, hit.score * super::index::EXPANSION_DECAY, Some(hit.symbol.clone()), + false, ), ); } @@ -426,7 +453,11 @@ impl LexicalIndex { } let mut hits: Vec = best.into_values().collect(); - hits.sort_by(|a, b| b.score.total_cmp(&a.score)); + hits.sort_by(|a, b| { + b.exact + .cmp(&a.exact) + .then_with(|| b.score.total_cmp(&a.score)) + }); budget_cut(hits, top_k, semantic_budget_tokens) } @@ -461,17 +492,233 @@ impl LexicalIndex { } scores } + + /// Candidate documents for an explicit identifier: its exact-token postings + /// plus symbol-name matches. These are already narrow, so unlike phrase + /// candidates they do not fall back to the broader BM25 ranking. + fn identifier_candidates(&self, identifier: &str, scores: &[f32]) -> Vec { + let tokens = unique(tokenize(identifier)); + let mut candidates: HashSet = HashSet::new(); + + if let Some(full) = tokens.first() + && let Some(postings) = self.postings.get(full) + { + candidates.extend(postings.iter().map(|posting| posting.0 as usize)); + } + if candidates.is_empty() { + candidates.extend(self.intersect_postings(&tokens[1..])); + } + for (index, doc) in self.docs.iter().enumerate() { + if doc.symbol.eq_ignore_ascii_case(identifier) { + candidates.insert(index); + } + } + self.cap_candidates(candidates, scores) + } + + /// Candidate documents for a quoted phrase: documents containing every + /// phrase term, capped to the strongest BM25 candidates. + fn phrase_candidates(&self, phrase: &[String], scores: &[f32]) -> Vec { + let candidates = self.intersect_postings(phrase); + self.cap_candidates(candidates, scores) + } + + fn intersect_postings(&self, terms: &[String]) -> HashSet { + let mut unique_terms = unique(terms.to_vec()); + if unique_terms.is_empty() { + return HashSet::new(); + } + let Some(first) = self.postings.get(&unique_terms[0]) else { + return HashSet::new(); + }; + let mut result: HashSet = first.iter().map(|posting| posting.0 as usize).collect(); + for term in unique_terms.drain(1..) { + let Some(postings) = self.postings.get(&term) else { + return HashSet::new(); + }; + let docs: HashSet = postings.iter().map(|posting| posting.0 as usize).collect(); + result.retain(|doc| docs.contains(doc)); + if result.is_empty() { + break; + } + } + result + } + + fn cap_candidates(&self, candidates: HashSet, scores: &[f32]) -> Vec { + let mut candidates: Vec = candidates.into_iter().collect(); + candidates.sort_by(|a, b| { + scores[*b] + .total_cmp(&scores[*a]) + .then_with(|| self.docs[*a].id.cmp(&self.docs[*b].id)) + }); + candidates.truncate(MAX_VERIFY_CANDIDATES); + candidates + } + + /// Verify candidate documents against the current source. Reads each file at + /// most once per query and checks only the symbol's line range. + fn verify( + &self, + workspace: &Path, + candidates: &[usize], + phrases: &[Vec], + identifier: Option<&str>, + ) -> HashSet { + let mut by_file: HashMap<&str, Vec> = HashMap::new(); + for &index in candidates { + if let Some(doc) = self.docs.get(index) { + by_file.entry(doc.file.as_str()).or_default().push(index); + } + } + + let mut verified = HashSet::new(); + for (file, indices) in by_file { + let Ok(content) = std::fs::read_to_string(workspace.join(file)) else { + continue; + }; + let lines: Vec<&str> = content.lines().collect(); + for index in indices { + let doc = &self.docs[index]; + let Some(slice) = slice_lines(&lines, doc.start_line, doc.end_line) else { + continue; + }; + if let Some(identifier) = identifier { + let symbol_match = doc.symbol.eq_ignore_ascii_case(identifier); + let source_match = contains_identifier(&slice, identifier); + if symbol_match || source_match { + verified.insert(index); + continue; + } + } + let haystack = phrase_tokens(&slice); + if phrases + .iter() + .any(|phrase| contains_sequence(&haystack, phrase)) + { + verified.insert(index); + } + } + } + verified + } +} + +/// One parsed query: BM25 terms, quoted phrases, and an explicit identifier. +struct ParsedQuery { + terms: Vec, + phrases: Vec>, + identifier: Option, +} + +fn parse_query(query: &str) -> ParsedQuery { + let mut phrases = Vec::new(); + let mut outside = String::new(); + let mut inside = String::new(); + let mut in_quote = false; + + for character in query.chars() { + if character == '"' { + if in_quote { + let phrase = phrase_tokens(&inside); + if !phrase.is_empty() { + phrases.push(phrase); + } + inside.clear(); + } + in_quote = !in_quote; + } else if in_quote { + inside.push(character); + } else { + outside.push(character); + } + } + // An unmatched quote is treated as a phrase rather than dropped. + if in_quote { + let phrase = phrase_tokens(&inside); + if !phrase.is_empty() { + phrases.push(phrase); + } + } + + let mut terms = tokenize(&outside); + for phrase in &phrases { + terms.extend(phrase.iter().cloned()); + } + let terms = unique(terms); + + let trimmed = query.trim(); + let identifier = if phrases.is_empty() && looks_like_identifier(trimmed) { + Some(trimmed.to_string()) + } else { + None + }; + + ParsedQuery { + terms, + phrases, + identifier, + } } -/// Tokenise a plain query into unique BM25 terms. -fn query_terms(query: &str) -> Vec { +/// A single identifier-like query with no spaces: `MAX_CHUNK_CHARS`, +/// `GraphStore::lookup`, `src/embeddings/index.rs`. +fn looks_like_identifier(query: &str) -> bool { + !query.is_empty() + && !query.chars().any(char::is_whitespace) + && !query.contains('"') + && (query.contains('_') + || query.contains("::") + || query.contains('.') + || query.contains('/') + || (query.chars().any(|c| c.is_uppercase()) && query.chars().any(|c| c.is_lowercase()))) +} + +/// Tokenise text for phrase verification: only lowercase separator/camel-case +/// parts, so `graph_store` and `GraphStore` compare as the same two words. +fn phrase_tokens(text: &str) -> Vec { + let mut tokens = Vec::new(); + for span in spans(text) { + for segment in split_segments(&span) { + tokens.extend(camel_parts(&segment)); + } + } + tokens +} + +fn unique(tokens: Vec) -> Vec { let mut seen = HashSet::new(); - tokenize(query) + tokens .into_iter() .filter(|token| seen.insert(token.clone())) .collect() } +/// Whether `haystack` contains `needle` as a whole identifier-like token. +fn contains_identifier(haystack: &str, needle: &str) -> bool { + let needle = needle.to_lowercase(); + spans(haystack) + .iter() + .any(|span| span.to_lowercase() == needle) +} + +fn slice_lines(lines: &[&str], start: usize, end: usize) -> Option { + if start == 0 || end < start || start > lines.len() { + return None; + } + let end = end.min(lines.len()); + Some(lines[start - 1..end].join("\n")) +} + +fn contains_sequence(haystack: &[String], needle: &[String]) -> bool { + if needle.is_empty() || needle.len() > haystack.len() { + return false; + } + haystack + .windows(needle.len()) + .any(|window| window == needle) +} + /// Term frequencies for one document, with symbol name and path weighted above /// the body. The full body is indexed, not the embedding truncation. fn doc_terms(symbol: &str, file: &str, body: &str) -> (Vec<(String, u32)>, u32) { @@ -612,6 +859,70 @@ mod tests { assert!(tokens.contains(&"index".to_string())); } + #[test] + fn phrase_tokens_normalise_separators_and_case() { + assert_eq!( + phrase_tokens("graph_store mutex poisoned"), + vec!["graph", "store", "mutex", "poisoned"] + ); + assert_eq!( + phrase_tokens("GraphStore mutex poisoned"), + vec!["graph", "store", "mutex", "poisoned"] + ); + } + + #[test] + fn quoted_phrase_is_verified_against_source_and_prioritised() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "fn first() {\n log::error!(\"graph_store mutex poisoned\");\n}\n", + ) + .unwrap(); + std::fs::write( + workspace.path().join("b.rs"), + "fn second() {\n // graph store and mutex poisoned in prose\n}\n", + ) + .unwrap(); + let index = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let ranked = index.rank(workspace.path(), "\"graph_store mutex poisoned\""); + assert_eq!(ranked.first().unwrap().symbol, "first"); + assert!(ranked.first().unwrap().exact); + } + + #[test] + fn explicit_identifier_is_verified_and_prioritised() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "fn lookup() {\n store.lookup(\"x\");\n}\n\ + fn lookup_extra() {\n store.lookup_extra(\"y\");\n}\n", + ) + .unwrap(); + let index = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let ranked = index.rank(workspace.path(), "store.lookup"); + assert_eq!(ranked.first().unwrap().symbol, "lookup"); + assert!(ranked.first().unwrap().exact); + } + + #[test] + fn identifier_verification_uses_token_boundaries() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn rrf_k() {}\npub fn rrf_k_extra() {}\n", + ) + .unwrap(); + let index = LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let ranked = index.rank(workspace.path(), "rrf_k"); + let exact: Vec<&str> = ranked + .iter() + .filter(|hit| hit.exact) + .map(|hit| hit.symbol.as_str()) + .collect(); + assert_eq!(exact, vec!["rrf_k"]); + } + #[test] fn sidecar_round_trips_and_persists_no_source_text() { let workspace = tempfile::tempdir().unwrap(); diff --git a/src/embeddings/mod.rs b/src/embeddings/mod.rs index c4bed06a..13c45efc 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -215,13 +215,15 @@ pub fn format_hits(query: &str, hits: &[SearchHit]) -> String { Some(sym) => format!(" (via {sym})"), None => String::new(), }; + let exact = if hit.exact { " (exact)" } else { "" }; out.push_str(&format!( - "{:>2}. {} — [{}:{}] {:.3}{}\n", + "{:>2}. {} — [{}:{}] {:.3}{}{}\n", i + 1, hit.symbol, hit.file, hit.start_line, hit.score, + exact, via, )); } From 47674aae2f2d5327c08af8f38df3fc57bd115a59 Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 13:19:00 +0200 Subject: [PATCH 06/10] [#336] Fuse lexical and semantic rankings with RRF Combine the semantic ranking (existing sketch filter and diversity penalty) with the BM25 ranking using plain Reciprocal Rank Fusion (k = 60) before graph expansion and the token budget. Verified exact matches are promoted ahead of fused results; equal fused scores prefer the semantic candidate. --- src/embeddings/index.rs | 344 ++++++++++++++++++++++++++++++++------ src/embeddings/lexical.rs | 5 + src/embeddings/mod.rs | 24 ++- 3 files changed, 321 insertions(+), 52 deletions(-) diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index deae58a3..0610e7e3 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -34,7 +34,7 @@ use serde::{Deserialize, Serialize}; use super::chunk::chunk_file; use super::client::EmbeddingClient; -use super::lexical::LexicalDoc; +use super::lexical::{LexicalDoc, LexicalIndex}; use super::source::{SourceFile, scan_workspace}; use crate::graph::extract::GraphExtractor; use crate::graph::store::GraphStore; @@ -68,6 +68,11 @@ const EMBED_CONCURRENCY: usize = 8; /// result never outranks the semantic hit that surfaced it. pub(super) const EXPANSION_DECAY: f32 = 0.6; +/// Reciprocal Rank Fusion constant. The standard default (60) damps the +/// influence of very high individual ranks while keeping the fusion stable +/// across lexical and semantic lists. +const RRF_K: f32 = 60.0; + // ── On-disk format ──────────────────────────────────────────────────────────── /// One embedded code chunk: a symbol plus its vector. @@ -467,9 +472,9 @@ impl EmbeddingIndex { Ok(()) } - /// Hybrid search: rank chunks by cosine similarity to `query_vector`, then, - /// when a graph is available, expand the strongest hits along their call - /// edges. Returns up to `top_k` results, most relevant first. + /// Semantic-only search: rank chunks by cosine similarity to + /// `query_vector`, then, when a graph is available, expand the strongest + /// hits along their call edges. Returns up to `top_k` results. pub fn search( &self, query_vector: &[f32], @@ -481,6 +486,51 @@ impl EmbeddingIndex { return Vec::new(); } + let penalized_scores = self.semantic_ranked(query_vector); + let id_to_index: HashMap<&str, usize> = self + .chunks + .iter() + .enumerate() + .map(|(i, c)| (c.id.as_str(), i)) + .collect(); + + // Take the strongest semantic hits as seeds for expansion. + let seed_count = top_k.saturating_mul(2).max(top_k); + let mut best: HashMap = HashMap::new(); + for &(idx, score) in penalized_scores.iter().take(seed_count) { + let chunk = &self.chunks[idx]; + insert_hit(&mut best, hit_from(chunk, score, None, false)); + + // Structural expansion: pull the graph neighbours of this symbol. + if let Some(graph) = graph { + for result in graph.lookup(&chunk.symbol) { + if result.node.id != chunk.id { + continue; + } + let neighbours = result.callers.iter().chain(result.callees.iter()); + for edge in neighbours { + if let Some(&nidx) = id_to_index.get(edge.node_id.as_str()) { + let n = &self.chunks[nidx]; + let boosted = score * EXPANSION_DECAY; + insert_hit( + &mut best, + hit_from(n, boosted, Some(chunk.symbol.clone()), false), + ); + } + } + } + } + } + + let mut hits: Vec = best.into_values().collect(); + hits.sort_by(|a, b| b.score.total_cmp(&a.score)); + budget_cut(hits, top_k, semantic_budget_tokens) + } + + /// Score every chunk against `query_vector`, apply the existing + /// occupancy-aware sketch filter and diversity penalty, and return the + /// chunk indices in descending score order. + fn semantic_ranked(&self, query_vector: &[f32]) -> Vec<(usize, f32)> { // Occupancy-Aware Sketch Filtering // Determine candidate modules that are close enough to the query. // Recall safety net: always force-include the 3 most recently modified files @@ -506,10 +556,6 @@ impl EmbeddingIndex { } } - // Track divergence/recall in debug builds - #[cfg(debug_assertions)] - let _unfiltered_count = self.chunks.len(); - let filter_candidates = !candidate_files.is_empty(); // Semantic pass: cosine of the query against every chunk. @@ -522,18 +568,7 @@ impl EmbeddingIndex { .collect(); scored.sort_by(|a, b| b.1.total_cmp(&a.1)); - // Best score seen per chunk id, so expansion never demotes a direct hit. - let mut best: HashMap = HashMap::new(); - let id_to_index: HashMap<&str, usize> = self - .chunks - .iter() - .enumerate() - .map(|(i, c)| (c.id.as_str(), i)) - .collect(); - - // Take the strongest semantic hits as seeds for expansion. // Pluribus Diversity-Aware Planner: track selected seeds to penalize redundant semantics - let seed_count = top_k.saturating_mul(2).max(top_k); let mut selected_seed_indices: Vec = Vec::new(); const DIVERSITY_THRESHOLD: f32 = 0.85; @@ -560,23 +595,137 @@ impl EmbeddingIndex { } penalized_scores.sort_by(|a, b| b.1.total_cmp(&a.1)); + penalized_scores + } - for &(idx, score) in penalized_scores.iter().take(seed_count) { - let chunk = &self.chunks[idx]; - insert_hit(&mut best, hit_from(chunk, score, None)); + /// Hybrid search: fuse the semantic ranking with the lexical BM25 ranking + /// using Reciprocal Rank Fusion, giving verified exact identifier/phrase + /// matches priority, then apply the existing graph expansion and token + /// budget. Falls back to lexical-only when the embedding index is empty. + #[allow(clippy::too_many_arguments)] + pub fn hybrid_search( + &self, + workspace: &Path, + query: &str, + query_vector: &[f32], + lexical: &LexicalIndex, + graph: Option<&GraphStore>, + top_k: usize, + semantic_budget_tokens: usize, + ) -> Vec { + if self.chunks.is_empty() || query_vector.is_empty() { + return lexical.search(workspace, query, graph, top_k, semantic_budget_tokens); + } - // Structural expansion: pull the graph neighbours of this symbol. - if let Some(graph) = graph { - for result in graph.lookup(&chunk.symbol) { - if result.node.id != chunk.id { + let semantic = self.semantic_ranked(query_vector); + let lexical_hits = lexical.rank(workspace, query); + if lexical_hits.is_empty() { + return self.search(query_vector, graph, top_k, semantic_budget_tokens); + } + + // Reciprocal Rank Fusion: rank-based, so BM25 scores and cosine + // similarities never need a shared scale. + let mut fused: HashMap<&str, f32> = HashMap::new(); + let mut from_semantic: std::collections::HashSet<&str> = std::collections::HashSet::new(); + for (rank, (idx, score)) in semantic.iter().enumerate() { + // Zero/negative cosine chunks are not semantic candidates at all; + // giving them a rank would let an unrelated lexical hit win. + if *score <= 0.0 { + continue; + } + let id = self.chunks[*idx].id.as_str(); + from_semantic.insert(id); + *fused.entry(id).or_default() += 1.0 / (RRF_K + rank as f32 + 1.0); + } + for (rank, hit) in lexical_hits.iter().enumerate() { + *fused.entry(hit.id.as_str()).or_default() += 1.0 / (RRF_K + rank as f32 + 1.0); + } + + let exact: std::collections::HashSet<&str> = lexical_hits + .iter() + .filter(|hit| hit.exact) + .map(|hit| hit.id.as_str()) + .collect(); + let chunk_by_id: HashMap<&str, &EmbeddedChunk> = self + .chunks + .iter() + .map(|chunk| (chunk.id.as_str(), chunk)) + .collect(); + let lexical_by_id = lexical.doc_by_id(); + + let mut candidates: Vec = fused + .into_iter() + .map(|(id, score)| { + let is_exact = exact.contains(id); + match chunk_by_id.get(id) { + Some(chunk) => hit_from(chunk, score, None, is_exact), + None => match lexical_by_id.get(id) { + Some(doc) => hit_from_doc(doc, score, None, is_exact), + // Should not happen: every lexical id belongs to a doc. + None => SearchHit { + id: id.to_string(), + symbol: String::new(), + file: String::new(), + start_line: 0, + end_line: 0, + score, + expanded_from: None, + exact: is_exact, + }, + }, + } + }) + .collect(); + // Plain RRF can tie a lexical-only top hit with the semantic top hit. + // Prefer the semantic candidate on equal scores, then the document id. + candidates.sort_by(|a, b| { + b.exact + .cmp(&a.exact) + .then_with(|| b.score.total_cmp(&a.score)) + .then_with(|| { + from_semantic + .contains(b.id.as_str()) + .cmp(&from_semantic.contains(a.id.as_str())) + }) + .then_with(|| a.id.cmp(&b.id)) + }); + + let mut best: HashMap = HashMap::new(); + for hit in &candidates { + insert_hit(&mut best, hit.clone()); + } + + // Structural expansion from the strongest fused seeds, so lexical hits + // also pull in their callers/callees. + let seed_count = top_k.saturating_mul(2).max(top_k); + if let Some(graph) = graph { + for hit in candidates.iter().take(seed_count) { + for result in graph.lookup(&hit.symbol) { + if result.node.id != hit.id { continue; } let neighbours = result.callers.iter().chain(result.callees.iter()); for edge in neighbours { - if let Some(&nidx) = id_to_index.get(edge.node_id.as_str()) { - let n = &self.chunks[nidx]; - let boosted = score * EXPANSION_DECAY; - insert_hit(&mut best, hit_from(n, boosted, Some(chunk.symbol.clone()))); + if let Some(chunk) = chunk_by_id.get(edge.node_id.as_str()) { + insert_hit( + &mut best, + hit_from( + chunk, + hit.score * EXPANSION_DECAY, + Some(hit.symbol.clone()), + false, + ), + ); + } else if let Some(doc) = lexical_by_id.get(edge.node_id.as_str()) { + insert_hit( + &mut best, + hit_from_doc( + doc, + hit.score * EXPANSION_DECAY, + Some(hit.symbol.clone()), + false, + ), + ); } } } @@ -584,24 +733,18 @@ impl EmbeddingIndex { } let mut hits: Vec = best.into_values().collect(); - hits.sort_by(|a, b| b.score.total_cmp(&a.score)); - - let mut total_tokens = 0; - let mut final_hits = Vec::new(); - for hit in hits { - let chunk_lines = hit.end_line.saturating_sub(hit.start_line) + 1; - let chunk_tokens = chunk_lines * 10; - if total_tokens + chunk_tokens > semantic_budget_tokens && !final_hits.is_empty() { - break; - } - total_tokens += chunk_tokens; - final_hits.push(hit); - if final_hits.len() >= top_k { - break; - } - } - - final_hits + hits.sort_by(|a, b| { + b.exact + .cmp(&a.exact) + .then_with(|| b.score.total_cmp(&a.score)) + .then_with(|| { + from_semantic + .contains(b.id.as_str()) + .cmp(&from_semantic.contains(a.id.as_str())) + }) + .then_with(|| a.id.cmp(&b.id)) + }); + budget_cut(hits, top_k, semantic_budget_tokens) } } @@ -616,7 +759,12 @@ pub(super) fn insert_hit(best: &mut HashMap, hit: SearchHit) } } -fn hit_from(chunk: &EmbeddedChunk, score: f32, expanded_from: Option) -> SearchHit { +fn hit_from( + chunk: &EmbeddedChunk, + score: f32, + expanded_from: Option, + exact: bool, +) -> SearchHit { SearchHit { id: chunk.id.clone(), symbol: chunk.symbol.clone(), @@ -625,7 +773,7 @@ fn hit_from(chunk: &EmbeddedChunk, score: f32, expanded_from: Option) -> end_line: chunk.end_line, score, expanded_from, - exact: false, + exact, } } @@ -920,4 +1068,100 @@ mod tests { // "three" should beat "two" due to the diversity penalty assert_eq!(hits[1].symbol, "three"); } + + #[test] + fn hybrid_search_prioritises_verified_exact_identifier() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn target_fn() {}\npub fn other_fn() {}\n", + ) + .unwrap(); + let lexical = + crate::embeddings::lexical::LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let index = EmbeddingIndex { + version: INDEX_VERSION, + file_hashes: HashMap::new(), + module_sketches: HashMap::new(), + chunks: vec![ + EmbeddedChunk { + id: "a::target_fn".into(), + symbol: "target_fn".into(), + file: "a.rs".into(), + start_line: 1, + end_line: 1, + vector: vec![0.0, 1.0], + }, + EmbeddedChunk { + id: "a::other_fn".into(), + symbol: "other_fn".into(), + file: "a.rs".into(), + start_line: 2, + end_line: 2, + vector: vec![1.0, 0.0], + }, + ], + }; + // The semantic ranking prefers `other_fn`, but `target_fn` is a verified + // exact identifier and must be promoted ahead of it. + let hits = index.hybrid_search( + workspace.path(), + "target_fn", + &[1.0, 0.0], + &lexical, + None, + 5, + 16384, + ); + assert_eq!(hits.first().unwrap().symbol, "target_fn"); + assert!(hits.first().unwrap().exact); + } + + #[test] + fn hybrid_search_keeps_semantic_lead_for_a_conceptual_query() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn semantic() {}\n\ + pub fn lexical() {\n retry();\n backoff();\n timeout();\n}\n", + ) + .unwrap(); + let lexical = + crate::embeddings::lexical::LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let index = EmbeddingIndex { + version: INDEX_VERSION, + file_hashes: HashMap::new(), + module_sketches: HashMap::new(), + chunks: vec![ + EmbeddedChunk { + id: "a::semantic".into(), + symbol: "semantic".into(), + file: "a.rs".into(), + start_line: 1, + end_line: 1, + vector: vec![1.0, 0.0], + }, + EmbeddedChunk { + id: "a::lexical".into(), + symbol: "lexical".into(), + file: "a.rs".into(), + start_line: 2, + end_line: 6, + vector: vec![0.0, 1.0], + }, + ], + }; + // BM25 strongly prefers `lexical`, but the query is conceptual (no + // verified exact match), so the semantic list must stay on top. + let hits = index.hybrid_search( + workspace.path(), + "retry backoff timeout", + &[1.0, 0.0], + &lexical, + None, + 5, + 16384, + ); + assert_eq!(hits.first().unwrap().symbol, "semantic"); + } } diff --git a/src/embeddings/lexical.rs b/src/embeddings/lexical.rs index 2011c49d..f4e01516 100644 --- a/src/embeddings/lexical.rs +++ b/src/embeddings/lexical.rs @@ -157,6 +157,11 @@ impl LexicalIndex { self.docs.is_empty() } + /// A map from document id to its metadata, for graph expansion by id. + pub fn doc_by_id(&self) -> HashMap<&str, &LexicalDoc> { + self.docs.iter().map(|doc| (doc.id.as_str(), doc)).collect() + } + /// Build or incrementally update the lexical index for `workspace`. /// /// Only files whose content hash changed are re-parsed and re-tokenised; diff --git a/src/embeddings/mod.rs b/src/embeddings/mod.rs index 13c45efc..1b51ea76 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -129,8 +129,17 @@ impl SemanticSearch { return Ok(result); }; - let index = EmbeddingIndex::build_or_update( + let scan = crate::embeddings::source::scan_workspace(workspace); + let lexical = crate::embeddings::lexical::LexicalIndex::build_or_update_from_scan( workspace, + &scan, + self.compile_workers, + None, + ) + .ok(); + let index = EmbeddingIndex::build_or_update_from_scan( + workspace, + &scan, client, cancel, self.compile_workers, @@ -153,7 +162,18 @@ impl SemanticSearch { let guard = graph .lock() .map_err(|_| anyhow::anyhow!("graph_store mutex poisoned"))?; - result.hits = index.search(&query_vector, guard.as_ref(), 10, self.budget_tokens); + result.hits = match &lexical { + Some(lexical) => index.hybrid_search( + workspace, + query, + &query_vector, + lexical, + guard.as_ref(), + 10, + self.budget_tokens, + ), + None => index.search(&query_vector, guard.as_ref(), 10, self.budget_tokens), + }; } Ok(result) } From b22f61a61becd24542d4044f3c89ead049cfae5f Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 13:20:21 +0200 Subject: [PATCH 07/10] [#336] Harden lexical and semantic fallback paths Keep the lexical index independent of the embeddings server: when the embedding build or query embedding fails after startup, fall back to BM25 results instead of returning an error. When lexical retrieval is unavailable but embeddings work, keep semantic-only behavior. Add a regression test for the server disappearing after startup. --- src/bin/orangu/dispatch.rs | 8 +++- src/embeddings/mod.rs | 89 ++++++++++++++++++++++++++------------ src/tools.rs | 67 ++++++++++++++++++++++++++-- 3 files changed, 132 insertions(+), 32 deletions(-) diff --git a/src/bin/orangu/dispatch.rs b/src/bin/orangu/dispatch.rs index 5dedf724..38c8b8d1 100644 --- a/src/bin/orangu/dispatch.rs +++ b/src/bin/orangu/dispatch.rs @@ -2220,7 +2220,7 @@ pub(crate) fn handle_command( } if result.indexed_symbols == 0 { let _ = sink.send( - "The embedding index is empty — no supported source files found to \ + "The search index is empty — no supported source files found to \ search." .to_string(), ); @@ -2231,7 +2231,11 @@ pub(crate) fn handle_command( result.indexed_symbols, format_status_duration(started.elapsed()) )); - let _ = sink.send(orangu::embeddings::format_hits(&query, &result.hits)); + let _ = sink.send(orangu::embeddings::format_hits( + &query, + &result.hits, + result.used_embeddings, + )); Ok(()) }), Some(control), diff --git a/src/embeddings/mod.rs b/src/embeddings/mod.rs index 1b51ea76..b414f22b 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -137,7 +137,12 @@ impl SemanticSearch { None, ) .ok(); - let index = EmbeddingIndex::build_or_update_from_scan( + // Try the embedding half. A server that was available at startup can + // disappear later; in that case fall through to lexical-only rather + // than failing the whole search. + let mut semantic = None; + let mut embedding_error = None; + match EmbeddingIndex::build_or_update_from_scan( workspace, &scan, client, @@ -146,34 +151,58 @@ impl SemanticSearch { progress, eta, ) - .await?; + .await + { + Ok(index) if !index.is_empty() => match client.embed_one(query).await { + Ok(vector) => semantic = Some((index, vector)), + Err(err) => embedding_error = Some(err), + }, + Ok(_) => {} + Err(err) => embedding_error = Some(err), + } + let mut result = SemanticSearchResult { - indexed_symbols: index.len(), + indexed_symbols: lexical.as_ref().map_or(0, |index| index.len()), hits: Vec::new(), cancelled: cancel.load(Ordering::Relaxed), - used_embeddings: true, + used_embeddings: false, }; - if result.cancelled || index.is_empty() { + if result.cancelled { return Ok(result); } - let query_vector = client.embed_one(query).await?; - result.cancelled = cancel.load(Ordering::Relaxed); - if !result.cancelled { - let guard = graph - .lock() - .map_err(|_| anyhow::anyhow!("graph_store mutex poisoned"))?; - result.hits = match &lexical { - Some(lexical) => index.hybrid_search( - workspace, - query, - &query_vector, - lexical, - guard.as_ref(), - 10, - self.budget_tokens, - ), - None => index.search(&query_vector, guard.as_ref(), 10, self.budget_tokens), - }; + + let guard = graph + .lock() + .map_err(|_| anyhow::anyhow!("graph_store mutex poisoned"))?; + + match semantic { + Some((index, query_vector)) => { + result.indexed_symbols = index.len(); + result.used_embeddings = true; + result.hits = match &lexical { + Some(lexical) => index.hybrid_search( + workspace, + query, + &query_vector, + lexical, + guard.as_ref(), + 10, + self.budget_tokens, + ), + None => index.search(&query_vector, guard.as_ref(), 10, self.budget_tokens), + }; + } + None => { + if let Some(lexical) = &lexical { + result.indexed_symbols = lexical.len(); + result.hits = + lexical.search(workspace, query, guard.as_ref(), 10, self.budget_tokens); + } else if let Some(err) = embedding_error { + return Err(err); + } else { + anyhow::bail!("lexical search index could not be built for this workspace"); + } + } } Ok(result) } @@ -224,12 +253,18 @@ pub async fn probe_endpoint(profile: &LlmConfiguration) -> Result<(), String> { } /// Render search hits as a human-readable block for the output window, with each -/// result anchored to its `file:line` so it is clickable. -pub fn format_hits(query: &str, hits: &[SearchHit]) -> String { +/// result anchored to its `file:line` so it is clickable. `used_embeddings` +/// selects the header so a lexical-only fallback does not claim to be semantic. +pub fn format_hits(query: &str, hits: &[SearchHit], used_embeddings: bool) -> String { + let label = if used_embeddings { + "Semantic" + } else { + "Lexical" + }; if hits.is_empty() { - return format!("No semantic matches for \"{query}\"."); + return format!("No matches for \"{query}\"."); } - let mut out = format!("Semantic matches for \"{query}\":\n"); + let mut out = format!("{label} matches for \"{query}\":\n"); for (i, hit) in hits.iter().enumerate() { let via = match &hit.expanded_from { Some(sym) => format!(" (via {sym})"), diff --git a/src/tools.rs b/src/tools.rs index a7f83fb9..08841b67 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -1344,7 +1344,11 @@ impl ToolExecutor { &AtomicU64::new(0), &AtomicU64::new(0), ))?; - Ok(crate::embeddings::format_hits(&query, &result.hits)) + Ok(crate::embeddings::format_hits( + &query, + &result.hits, + result.used_embeddings, + )) }) .await? } @@ -4137,7 +4141,7 @@ mod semantic_search_tests { .execute("semantic_search", arguments.as_object().unwrap()) .await .unwrap(); - assert!(answer.contains("No semantic matches"), "{answer}"); + assert!(answer.contains("No matches"), "{answer}"); assert_eq!( requests.lock().unwrap().len(), 5, @@ -4153,7 +4157,64 @@ mod semantic_search_tests { .execute("semantic_search", arguments.as_object().unwrap()) .await .unwrap(); - assert!(answer.contains("No semantic matches"), "{answer}"); + assert!(answer.contains("No matches"), "{answer}"); + server.abort(); + } + + #[tokio::test] + async fn semantic_search_falls_back_to_lexical_when_embeddings_go_down_after_startup() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let requests = Requests::default(); + let app = Router::new() + .route("/v1/embeddings", post(embed)) + .with_state(requests); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let workspace = tempfile::tempdir().unwrap(); + let _cache = CacheCleanup(crate::embeddings::EmbeddingIndex::cache_dir( + workspace.path(), + )); + fs::write( + workspace.path().join("retry.rs"), + "pub fn retry_backoff() {}\n", + ) + .unwrap(); + let tools = ToolExecutor::new_read_only(workspace.path()); + let profile = crate::config::LlmConfiguration { + endpoint, + model: "chat-model".into(), + role: "all".into(), + api_key: None, + request_timeout_seconds: 1, + max_tool_rounds: 1, + review_max_tokens: 0, + code_max_tokens: 0, + system_prompt: String::new(), + model_verbosity: None, + review_confidence_threshold: 0, + security_confidence_threshold: 0, + review_sanitize: Default::default(), + review_sanitize_timeout: 1_800, + review_sanitize_alloc_failures: 8, + }; + tools + .configure_semantic_search(Some(&profile), true, 1, 16384) + .unwrap(); + let arguments = json!({"query": "retry_backoff"}); + let answer = tools + .execute("semantic_search", arguments.as_object().unwrap()) + .await + .unwrap(); + assert!(answer.contains("retry_backoff"), "{answer}"); + + // The server disappears after startup. The next search must still + // return ranked lexical results instead of failing. server.abort(); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + let answer = tools + .execute("semantic_search", arguments.as_object().unwrap()) + .await + .expect("lexical fallback should not fail"); + assert!(answer.contains("retry_backoff"), "{answer}"); } } From 3bc9dacb68780759dd5acbdd6b0b6b0f114c86ab Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 18:36:24 +0200 Subject: [PATCH 08/10] [#336] Bound the diversity pass to the result head MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The diversity penalty compared every chunk with every seed selected before it, so one query cost O(chunks^2 * dimensions). On this workspace (~21k chunks at 768 dimensions) that is roughly six minutes per query, which makes the benchmark — and `/search` on any repository this size — unusable. Only the head of the ranking can reach the result set: `search` reads at most `top_k * 2` entries, and RRF fuses the head of the semantic ranking. Restrict the pass to the top `DIVERSITY_CANDIDATE_LIMIT` (256) candidates. Query time drops to milliseconds and ordering is unchanged for everything that can actually be returned. Found by the Phase 7 benchmark once it ran against a full-size index. --- src/embeddings/index.rs | 46 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index 0610e7e3..499ed89b 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -73,6 +73,11 @@ pub(super) const EXPANSION_DECAY: f32 = 0.6; /// across lexical and semantic lists. const RRF_K: f32 = 60.0; +/// How many top-scoring chunks the diversity pass may consider. Scoring every +/// chunk against every seed is quadratic, so this bounds the pass to the part +/// of the ranking that can still reach the result set. +const DIVERSITY_CANDIDATE_LIMIT: usize = 256; + // ── On-disk format ──────────────────────────────────────────────────────────── /// One embedded code chunk: a symbol plus its vector. @@ -568,6 +573,13 @@ impl EmbeddingIndex { .collect(); scored.sort_by(|a, b| b.1.total_cmp(&a.1)); + // The diversity pass below is quadratic in the number of candidates + // (each chunk is compared with every seed selected before it), which is + // minutes per query on a 20k-symbol workspace. Only the head of the + // ranking can reach the result set, so bound the pool: `search` reads at + // most `top_k * 2` entries and RRF fuses the head of this ranking. + scored.truncate(DIVERSITY_CANDIDATE_LIMIT); + // Pluribus Diversity-Aware Planner: track selected seeds to penalize redundant semantics let mut selected_seed_indices: Vec = Vec::new(); @@ -1069,6 +1081,40 @@ mod tests { assert_eq!(hits[1].symbol, "three"); } + #[test] + fn diversity_pass_only_scores_the_head_of_the_ranking() { + // Every chunk here is on the unit circle, evenly spread, so most pairs + // are *not* near-duplicates and the seed list keeps growing — the + // quadratic worst case. The pass must stop at the bounded pool instead + // of comparing every chunk against every seed. + let total = DIVERSITY_CANDIDATE_LIMIT * 2; + let chunks: Vec = (0..total) + .map(|i| { + let angle = std::f32::consts::TAU * i as f32 / total as f32; + EmbeddedChunk { + id: format!("a::c{i}"), + symbol: format!("c{i}"), + file: "a.rs".into(), + start_line: i + 1, + end_line: i + 1, + vector: vec![angle.cos(), angle.sin()], + } + }) + .collect(); + let index = EmbeddingIndex { + version: INDEX_VERSION, + file_hashes: HashMap::new(), + module_sketches: HashMap::new(), + chunks, + }; + + let ranked = index.semantic_ranked(&[1.0, 0.0]); + assert_eq!(ranked.len(), DIVERSITY_CANDIDATE_LIMIT); + // Bound the head, not the tail: the closest chunk is still first. + assert_eq!(ranked[0].0, 0); + assert!(ranked[0].1 > 0.999); + } + #[test] fn hybrid_search_prioritises_verified_exact_identifier() { let workspace = tempfile::tempdir().unwrap(); From ef17c2d84640304ab502f18554360b27d6e1c5a8 Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Fri, 13 Nov 2026 13:23:12 +0200 Subject: [PATCH 09/10] [#336] Add hybrid search benchmark Add a committed fixture and binary that compares lexical-only, semantic-only, and hybrid ranking, reporting MRR, Recall@5, latency, indexing time, and sidecar size. Standalone constants that the current Rust extractor does not emit are tracked as coverage gaps. Expectations name unique ids where a symbol name is ambiguous. --verbose prints the lexical, semantic, and hybrid hits. Record the results in the README. --- Cargo.toml | 6 + src/bin/orangu-search-bench/README.md | 56 ++++ src/bin/orangu-search-bench/main.rs | 381 +++++++++++++++++++++++ src/bin/orangu-search-bench/queries.json | 154 +++++++++ 4 files changed, 597 insertions(+) create mode 100644 src/bin/orangu-search-bench/README.md create mode 100644 src/bin/orangu-search-bench/main.rs create mode 100644 src/bin/orangu-search-bench/queries.json diff --git a/Cargo.toml b/Cargo.toml index de64ca4a..4fc50a31 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,12 @@ path = "src/bin/orangu-server/main.rs" name = "orangu-bench" path = "src/bin/orangu-bench/main.rs" +# Developer retrieval benchmark for the hybrid lexical/semantic search PoC. +# See src/bin/orangu-search-bench/README.md. +[[bin]] +name = "orangu-search-bench" +path = "src/bin/orangu-search-bench/main.rs" + # Builds a model file: trains one from a manifest of permissively-licensed # repositories, or rewrites an existing one at another weight format. # Documented in doc/manual/en/77-gguf.md. diff --git a/src/bin/orangu-search-bench/README.md b/src/bin/orangu-search-bench/README.md new file mode 100644 index 00000000..838eaa9e --- /dev/null +++ b/src/bin/orangu-search-bench/README.md @@ -0,0 +1,56 @@ +# orangu-search-bench + +A small retrieval benchmark for the hybrid lexical/semantic search proof of +concept. It runs the committed query fixture (`queries.json`) against: + +- **lexical-only** — BM25 over the lexical sidecar +- **semantic-only** — cosine similarity over the embedding index +- **hybrid** — Reciprocal Rank Fusion of both + +and reports MRR, Recall@5, mean latency, indexing time, and the lexical +sidecar's on-disk size. + +```sh +# Lexical-only (no embeddings server required) +cargo run --release --bin orangu-search-bench -- --workspace . + +# Semantic-only and hybrid +cargo run --release --bin orangu-search-bench -- --workspace . \ + --endpoint http://127.0.0.1:8080 --model embeddings +``` + +`--verbose` prints the top lexical, semantic and hybrid hits per query. The fixture marks +standalone constants that the current Rust extractor does not emit (for example +`MAX_CHUNK_CHARS`) as `coverage_gap`; those are reported separately rather than +counted as ranking failures. + +Queries are evaluated against the workspace as it exists, so test code that +contains the same identifiers or strings also participates in the ranking. +The fixture avoids phrases that only exist in test code. + +## Results + +Run on the `hybrid-search-phases` branch, 2026-11-13: 22,492 indexed symbols, +18 scored queries (10 conceptual, 7 exact, 1 phrase). Embeddings: +`ggml-org/embeddinggemma-300M-GGUF:Q8_0` (768 dims, CPU). + +| metric | semantic-only | lexical-only | hybrid | +|---|---|---|---| +| overall MRR | 0.294 | 0.536 | 0.483 | +| overall Recall@5 | 0.444 | 0.667 | 0.667 | +| conceptual MRR (n=10) | 0.278 | 0.165 | 0.220 | +| conceptual Recall@5 (n=10) | 0.500 | 0.400 | 0.500 | +| exact MRR / Recall@5 (n=7) | 0.357 / 0.429 | 1.000 / 1.000 | 0.857 / 1.000 | +| phrase rank | miss | 1 (verified exact) | 2 | +| mean query latency | 230 ms | 17 ms | 276 ms | + +Semantic and hybrid latency include one query embedding request (~200 ms). + +| index | cold build | warm build | on disk | +|---|---|---|---| +| lexical sidecar | 10.8 s | 2.8 s | 43,921 KiB (42.9 MiB) | +| semantic index | ~4.8 h | 26 s | 210 MB | + +Standalone constants (`MAX_CHUNK_CHARS`, `EMBED_BATCH_TOKEN_BUDGET`, `RRF_K`) +are reported as extraction coverage gaps and excluded from the metrics. + diff --git a/src/bin/orangu-search-bench/main.rs b/src/bin/orangu-search-bench/main.rs new file mode 100644 index 00000000..5af9334f --- /dev/null +++ b/src/bin/orangu-search-bench/main.rs @@ -0,0 +1,381 @@ +// Copyright (C) 2026 The orangu community +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +//! A small, committed retrieval benchmark for the hybrid-search PoC. +//! +//! Runs the same query fixture against lexical-only, semantic-only, and hybrid +//! ranking, reporting MRR, Recall@5, per-query latency, indexing time, and the +//! lexical sidecar's on-disk size. The fixture targets symbols that the current +//! Rust extractor actually emits; standalone constants that are not extracted +//! are marked `coverage_gap` and reported separately. +//! +//! ```text +//! cargo run --release --bin orangu-search-bench -- --workspace . +//! cargo run --release --bin orangu-search-bench -- --workspace . \ +//! --endpoint http://127.0.0.1:8080 --model embeddings +//! ``` + +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::time::{Duration, Instant}; + +use orangu::config::LlmConfiguration; +use orangu::embeddings::{EmbeddingClient, EmbeddingIndex, SearchHit, lexical::LexicalIndex}; +use serde::Deserialize; + +#[derive(Deserialize)] +struct Fixture { + queries: Vec, +} + +#[derive(Deserialize)] +struct QueryCase { + query: String, + #[serde(default)] + kind: String, + expected: Vec, + #[serde(default)] + coverage_gap: bool, +} + +#[derive(Default)] +struct Scores { + mrr: f64, + recall: f64, + latency: Duration, + count: usize, +} + +impl Scores { + fn add(&mut self, mrr: f64, recall: f64, latency: Duration) { + self.mrr += mrr; + self.recall += recall; + self.latency += latency; + self.count += 1; + } + + fn summary(&self) -> String { + if self.count == 0 { + return "-".to_string(); + } + format!( + "{:.3}/{:.3} {}ms", + self.mrr / self.count as f64, + self.recall / self.count as f64, + self.latency.as_millis() / self.count as u128 + ) + } +} + +fn main() -> anyhow::Result<()> { + let options = Options::parse()?; + let fixture: Fixture = serde_json::from_str(include_str!("queries.json"))?; + let workspace = options.workspace.clone(); + let runtime = tokio::runtime::Runtime::new()?; + + let lexical_started = Instant::now(); + let lexical = LexicalIndex::build_or_update(&workspace, options.compile_workers)?; + let lexical_build = lexical_started.elapsed(); + let sidecar_bytes = std::fs::metadata(LexicalIndex::cache_path(&workspace)) + .map(|meta| meta.len()) + .unwrap_or(0); + + println!("workspace: {}", workspace.display()); + println!( + "lexical: {} symbols, build {}ms, sidecar {:.1} KiB", + lexical.len(), + lexical_build.as_millis(), + sidecar_bytes as f64 / 1024.0 + ); + + let embedding = if let Some(endpoint) = &options.endpoint { + let profile = LlmConfiguration { + endpoint: endpoint.clone(), + model: options.model.clone(), + role: "embeddings".into(), + api_key: options.api_key.clone(), + request_timeout_seconds: 120, + max_tool_rounds: 1, + review_max_tokens: 0, + code_max_tokens: 0, + system_prompt: String::new(), + model_verbosity: None, + review_confidence_threshold: 0, + security_confidence_threshold: 0, + review_sanitize: Default::default(), + review_sanitize_timeout: 1_800, + review_sanitize_alloc_failures: 8, + }; + let client = EmbeddingClient::from_profile(&profile)?; + let started = Instant::now(); + let index = runtime.block_on(EmbeddingIndex::build_or_update( + &workspace, + &client, + &AtomicBool::new(false), + options.compile_workers, + &AtomicU64::new(0), + &AtomicU64::new(0), + ))?; + println!( + "semantic: {} symbols, build {}ms", + index.len(), + started.elapsed().as_millis() + ); + Some((client, index)) + } else { + println!("semantic: skipped (pass --endpoint and --model to enable)"); + None + }; + + let mut lexical_scores = Scores::default(); + let mut semantic_scores = Scores::default(); + let mut hybrid_scores = Scores::default(); + let mut gaps = 0usize; + + println!(); + println!( + "{:<48} {:<10} {:>18} {:>18} {:>18}", + "query", "kind", "lexical", "semantic", "hybrid" + ); + + for case in fixture.queries.iter().filter(|case| !case.coverage_gap) { + let started = Instant::now(); + let lexical_hits = lexical.search( + &workspace, + &case.query, + None, + options.top_k, + options.budget_tokens, + ); + let lexical_latency = started.elapsed(); + let (lexical_mrr, lexical_recall) = metrics(&lexical_hits, &case.expected, options.top_k); + lexical_scores.add(lexical_mrr, lexical_recall, lexical_latency); + + let mut semantic_hits: Vec = Vec::new(); + let mut hybrid_hits: Vec = Vec::new(); + let (semantic_cell, hybrid_cell) = match &embedding { + Some((client, index)) => { + let started = Instant::now(); + let vector = runtime.block_on(client.embed_one(&case.query))?; + let embed_latency = started.elapsed(); + + let started = Instant::now(); + semantic_hits = index.search(&vector, None, options.top_k, options.budget_tokens); + let semantic_latency = embed_latency + started.elapsed(); + let (semantic_mrr, semantic_recall) = + metrics(&semantic_hits, &case.expected, options.top_k); + semantic_scores.add(semantic_mrr, semantic_recall, semantic_latency); + + let started = Instant::now(); + hybrid_hits = index.hybrid_search( + &workspace, + &case.query, + &vector, + &lexical, + None, + options.top_k, + options.budget_tokens, + ); + let hybrid_latency = embed_latency + started.elapsed(); + let (hybrid_mrr, hybrid_recall) = + metrics(&hybrid_hits, &case.expected, options.top_k); + hybrid_scores.add(hybrid_mrr, hybrid_recall, hybrid_latency); + + ( + format!("{semantic_mrr:.3}/{semantic_recall:.3}"), + format!("{hybrid_mrr:.3}/{hybrid_recall:.3}"), + ) + } + None => ("-".to_string(), "-".to_string()), + }; + + println!( + "{:<48} {:<10} {:>18} {:>18} {:>18}", + truncate(&case.query, 48), + truncate(&case.kind, 10), + format!("{lexical_mrr:.3}/{lexical_recall:.3}"), + semantic_cell, + hybrid_cell + ); + if options.verbose { + for (label, hits) in [ + ("lexical", &lexical_hits), + ("semantic", &semantic_hits), + ("hybrid", &hybrid_hits), + ] { + if hits.is_empty() { + continue; + } + println!(" {label} hits:"); + for hit in hits.iter().take(options.top_k) { + println!( + " {:<24} {:<32} {:.4}{}{}", + hit.symbol, + format!("{}:{}", hit.file, hit.start_line), + hit.score, + if hit.exact { " exact" } else { "" }, + match &hit.expanded_from { + Some(from) => format!(" via {from}"), + None => String::new(), + } + ); + } + } + } + } + + println!(); + println!( + "{:<48} {:<10} {:>18} {:>18} {:>18}", + "summary (MRR/Recall@5, mean latency)", + "", + lexical_scores.summary(), + semantic_scores.summary(), + hybrid_scores.summary() + ); + + for case in fixture.queries.iter().filter(|case| case.coverage_gap) { + gaps += 1; + println!( + "coverage gap: {:<40} expected {} (not emitted by the current extractor)", + truncate(&case.query, 40), + case.expected.join(", ") + ); + } + if gaps > 0 { + println!("{gaps} coverage-gap queries are excluded from ranking metrics."); + } + + Ok(()) +} + +struct Options { + workspace: PathBuf, + endpoint: Option, + model: String, + api_key: Option, + top_k: usize, + budget_tokens: usize, + compile_workers: usize, + verbose: bool, +} + +impl Options { + fn parse() -> anyhow::Result { + let mut options = Self { + workspace: std::env::current_dir()?, + endpoint: None, + model: "embeddings".to_string(), + api_key: std::env::var("ORANGU_EMBED_API_KEY").ok(), + top_k: 5, + budget_tokens: 16_384, + compile_workers: 0, + verbose: false, + }; + let mut args = std::env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--workspace" => { + options.workspace = PathBuf::from( + args.next() + .ok_or_else(|| anyhow::anyhow!("--workspace requires a path"))?, + ); + } + "--endpoint" => { + options.endpoint = Some( + args.next() + .ok_or_else(|| anyhow::anyhow!("--endpoint requires a URL"))?, + ); + } + "--model" => { + options.model = args + .next() + .ok_or_else(|| anyhow::anyhow!("--model requires a name"))?; + } + "--api-key" => { + options.api_key = Some( + args.next() + .ok_or_else(|| anyhow::anyhow!("--api-key requires a value"))?, + ); + } + "--top-k" => { + options.top_k = args + .next() + .ok_or_else(|| anyhow::anyhow!("--top-k requires a number"))? + .parse()?; + } + "--budget-tokens" => { + options.budget_tokens = args + .next() + .ok_or_else(|| anyhow::anyhow!("--budget-tokens requires a number"))? + .parse()?; + } + "--compile-workers" => { + options.compile_workers = args + .next() + .ok_or_else(|| anyhow::anyhow!("--compile-workers requires a number"))? + .parse()?; + } + "--verbose" => options.verbose = true, + "--help" | "-h" => { + println!( + "usage: orangu-search-bench [--workspace PATH] [--endpoint URL] \ + [--model NAME] [--api-key KEY] [--top-k N] [--budget-tokens N] \ + [--compile-workers N]" + ); + std::process::exit(0); + } + other => anyhow::bail!("unknown argument: {other}"), + } + } + Ok(options) + } +} + +/// MRR and Recall@k for one query. A hit is relevant when its symbol, id, or +/// id suffix matches one of the expected names. +fn metrics(hits: &[SearchHit], expected: &[String], k: usize) -> (f64, f64) { + let reciprocal_rank = hits + .iter() + .position(|hit| is_relevant(hit, expected)) + .map(|position| 1.0 / (position + 1) as f64) + .unwrap_or(0.0); + let recall = if hits.iter().take(k).any(|hit| is_relevant(hit, expected)) { + 1.0 + } else { + 0.0 + }; + (reciprocal_rank, recall) +} + +fn is_relevant(hit: &SearchHit, expected: &[String]) -> bool { + expected.iter().any(|name| { + hit.symbol.eq_ignore_ascii_case(name) + || hit.id.eq_ignore_ascii_case(name) + || hit + .id + .to_ascii_lowercase() + .ends_with(&format!("::{}", name.to_ascii_lowercase())) + }) +} + +fn truncate(value: &str, width: usize) -> String { + if value.chars().count() <= width { + return value.to_string(); + } + let mut truncated: String = value.chars().take(width.saturating_sub(1)).collect(); + truncated.push('…'); + truncated +} diff --git a/src/bin/orangu-search-bench/queries.json b/src/bin/orangu-search-bench/queries.json new file mode 100644 index 00000000..1d32cfe6 --- /dev/null +++ b/src/bin/orangu-search-bench/queries.json @@ -0,0 +1,154 @@ +{ + "queries": [ + { + "query": "where is the embedding batch size capped", + "kind": "conceptual", + "expected": [ + "batches_within_budget" + ] + }, + { + "query": "how are near-duplicate search results penalised", + "kind": "conceptual", + "expected": [ + "semantic_ranked" + ] + }, + { + "query": "grep that respects gitignore", + "kind": "conceptual", + "expected": [ + "search_text" + ] + }, + { + "query": "how does the knowledge graph find callers", + "kind": "conceptual", + "expected": [ + "src_graph_store::lookup" + ] + }, + { + "query": "walk the workspace source files", + "kind": "conceptual", + "expected": [ + "src_embeddings_source::collect_source_files" + ] + }, + { + "query": "probe the embeddings endpoint with a timeout", + "kind": "conceptual", + "expected": [ + "probe_endpoint" + ] + }, + { + "query": "split camel case identifiers", + "kind": "conceptual", + "expected": [ + "camel_parts" + ] + }, + { + "query": "apply the semantic token budget", + "kind": "conceptual", + "expected": [ + "budget_cut" + ] + }, + { + "query": "fuse lexical and semantic rankings", + "kind": "conceptual", + "expected": [ + "hybrid_search" + ] + }, + { + "query": "incrementally update the lexical sidecar", + "kind": "conceptual", + "expected": [ + "src_embeddings_lexical::build_or_update" + ] + }, + { + "query": "batches_within_budget", + "kind": "exact", + "expected": [ + "batches_within_budget" + ] + }, + { + "query": "probe_endpoint", + "kind": "exact", + "expected": [ + "probe_endpoint" + ] + }, + { + "query": "collect_source_files", + "kind": "exact", + "expected": [ + "collect_source_files" + ] + }, + { + "query": "camel_parts", + "kind": "exact", + "expected": [ + "camel_parts" + ] + }, + { + "query": "budget_cut", + "kind": "exact", + "expected": [ + "budget_cut" + ] + }, + { + "query": "hybrid_search", + "kind": "exact", + "expected": [ + "hybrid_search" + ] + }, + { + "query": "lexical_chunks", + "kind": "exact", + "expected": [ + "lexical_chunks" + ] + }, + { + "query": "\"lexical search index could not be built\"", + "kind": "phrase", + "expected": [ + "src_embeddings_mod::search" + ] + }, + { + "query": "MAX_CHUNK_CHARS", + "kind": "exact", + "expected": [ + "MAX_CHUNK_CHARS" + ], + "coverage_gap": true + }, + { + "query": "EMBED_BATCH_TOKEN_BUDGET", + "kind": "exact", + "expected": [ + "EMBED_BATCH_TOKEN_BUDGET" + ], + "coverage_gap": true + }, + { + "query": "RRF_K", + "kind": "exact", + "expected": [ + "RRF_K" + ], + "coverage_gap": true + } + ] +} From 1d57789dface784bab86b3b33df13fc0813fd031 Mon Sep 17 00:00:00 2001 From: youssef-joe Date: Sat, 14 Nov 2026 00:30:07 +0200 Subject: [PATCH 10/10] [#336] Keep lexical order among verified exact hits Hybrid search placed verified exact hits first but then ordered them by their RRF score. RRF is rank-only: lexical #1 and #2 differ by ~0.0003, while appearing anywhere in the semantic list adds >= ~0.003, so a test that calls an identifier outranked its definition. Exact hits now keep the order the lexical ranking gave them (BM25), and RRF orders only the remaining candidates, as agreed in the design. A verified hit also keeps its exact flag when a higher-scoring graph-expanded copy replaces it. Update the benchmark results. --- src/bin/orangu-search-bench/README.md | 18 ++--- src/embeddings/index.rs | 101 ++++++++++++++++++++++---- 2 files changed, 96 insertions(+), 23 deletions(-) diff --git a/src/bin/orangu-search-bench/README.md b/src/bin/orangu-search-bench/README.md index 838eaa9e..6b75b45d 100644 --- a/src/bin/orangu-search-bench/README.md +++ b/src/bin/orangu-search-bench/README.md @@ -30,25 +30,25 @@ The fixture avoids phrases that only exist in test code. ## Results -Run on the `hybrid-search-phases` branch, 2026-11-13: 22,492 indexed symbols, +Run on the `hybrid-search-phases` branch, 2026-11-14: 22,493 indexed symbols, 18 scored queries (10 conceptual, 7 exact, 1 phrase). Embeddings: `ggml-org/embeddinggemma-300M-GGUF:Q8_0` (768 dims, CPU). | metric | semantic-only | lexical-only | hybrid | |---|---|---|---| -| overall MRR | 0.294 | 0.536 | 0.483 | +| overall MRR | 0.294 | 0.541 | 0.567 | | overall Recall@5 | 0.444 | 0.667 | 0.667 | -| conceptual MRR (n=10) | 0.278 | 0.165 | 0.220 | -| conceptual Recall@5 (n=10) | 0.500 | 0.400 | 0.500 | -| exact MRR / Recall@5 (n=7) | 0.357 / 0.429 | 1.000 / 1.000 | 0.857 / 1.000 | -| phrase rank | miss | 1 (verified exact) | 2 | -| mean query latency | 230 ms | 17 ms | 276 ms | +| conceptual MRR (n=10) | 0.278 | 0.173 | 0.220 | +| conceptual Recall@5 (n=10) | 0.500 | 0.400 | 0.400 | +| exact MRR / Recall@5 (n=7) | 0.357 / 0.429 | 1.000 / 1.000 | 1.000 / 1.000 | +| phrase rank | miss | 1 (verified exact) | 1 (verified exact) | +| mean query latency | 85 ms | 6 ms | 101 ms | -Semantic and hybrid latency include one query embedding request (~200 ms). +Semantic and hybrid latency include one query embedding request. | index | cold build | warm build | on disk | |---|---|---|---| -| lexical sidecar | 10.8 s | 2.8 s | 43,921 KiB (42.9 MiB) | +| lexical sidecar | 10.8 s | 2.8 s | 43,927 KiB (42.9 MiB) | | semantic index | ~4.8 h | 26 s | 210 MB | Standalone constants (`MAX_CHUNK_CHARS`, `EMBED_BATCH_TOKEN_BUDGET`, `RRF_K`) diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index 499ed89b..0285a318 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -653,11 +653,17 @@ impl EmbeddingIndex { *fused.entry(hit.id.as_str()).or_default() += 1.0 / (RRF_K + rank as f32 + 1.0); } - let exact: std::collections::HashSet<&str> = lexical_hits - .iter() - .filter(|hit| hit.exact) - .map(|hit| hit.id.as_str()) - .collect(); + // Verified exact hits are prioritised ahead of the fusion and keep the + // order the lexical ranking gave them (BM25 among exact hits, so a + // definition outranks a caller). RRF only orders the remaining + // candidates: its rank-only scores cannot separate lexical #1 from #2 + // once either also appears anywhere in the semantic list. + let mut exact_order: HashMap<&str, usize> = HashMap::new(); + for hit in lexical_hits.iter().filter(|hit| hit.exact) { + let next = exact_order.len(); + exact_order.entry(hit.id.as_str()).or_insert(next); + } + let exact_rank = |id: &str| exact_order.get(id).copied().unwrap_or(usize::MAX); let chunk_by_id: HashMap<&str, &EmbeddedChunk> = self .chunks .iter() @@ -668,7 +674,7 @@ impl EmbeddingIndex { let mut candidates: Vec = fused .into_iter() .map(|(id, score)| { - let is_exact = exact.contains(id); + let is_exact = exact_order.contains_key(id); match chunk_by_id.get(id) { Some(chunk) => hit_from(chunk, score, None, is_exact), None => match lexical_by_id.get(id) { @@ -691,8 +697,8 @@ impl EmbeddingIndex { // Plain RRF can tie a lexical-only top hit with the semantic top hit. // Prefer the semantic candidate on equal scores, then the document id. candidates.sort_by(|a, b| { - b.exact - .cmp(&a.exact) + exact_rank(&a.id) + .cmp(&exact_rank(&b.id)) .then_with(|| b.score.total_cmp(&a.score)) .then_with(|| { from_semantic @@ -746,8 +752,8 @@ impl EmbeddingIndex { let mut hits: Vec = best.into_values().collect(); hits.sort_by(|a, b| { - b.exact - .cmp(&a.exact) + exact_rank(&a.id) + .cmp(&exact_rank(&b.id)) .then_with(|| b.score.total_cmp(&a.score)) .then_with(|| { from_semantic @@ -762,10 +768,18 @@ impl EmbeddingIndex { /// Insert `hit` into `best` unless an equal-or-better score for the same id is /// already recorded. -pub(super) fn insert_hit(best: &mut HashMap, hit: SearchHit) { - match best.get(&hit.id) { - Some(existing) if existing.score >= hit.score => {} - _ => { +pub(super) fn insert_hit(best: &mut HashMap, mut hit: SearchHit) { + match best.get_mut(&hit.id) { + Some(existing) if existing.score >= hit.score => { + existing.exact |= hit.exact; + } + Some(existing) => { + // A verified exact match stays exact even when a higher-scoring + // copy of the same symbol (e.g. a graph expansion) replaces it. + hit.exact |= existing.exact; + *existing = hit; + } + None => { best.insert(hit.id.clone(), hit); } } @@ -1163,6 +1177,65 @@ mod tests { assert!(hits.first().unwrap().exact); } + #[test] + fn hybrid_search_keeps_lexical_order_among_exact_hits() { + let workspace = tempfile::tempdir().unwrap(); + std::fs::write( + workspace.path().join("a.rs"), + "pub fn ping_backend() {}\npub fn ping_backend_retries() {\n ping_backend();\n}\n", + ) + .unwrap(); + let lexical = + crate::embeddings::lexical::LexicalIndex::build_or_update(workspace.path(), 1).unwrap(); + let lexical_order: Vec = lexical + .rank(workspace.path(), "ping_backend") + .into_iter() + .filter(|hit| hit.exact) + .map(|hit| hit.symbol) + .collect(); + assert_eq!(lexical_order, ["ping_backend", "ping_backend_retries"]); + + let index = EmbeddingIndex { + version: INDEX_VERSION, + file_hashes: HashMap::new(), + module_sketches: HashMap::new(), + chunks: vec![ + EmbeddedChunk { + id: "a::ping_backend".into(), + symbol: "ping_backend".into(), + file: "a.rs".into(), + start_line: 1, + end_line: 1, + vector: vec![-1.0, 0.0], + }, + EmbeddedChunk { + id: "a::ping_backend_retries".into(), + symbol: "ping_backend_retries".into(), + file: "a.rs".into(), + start_line: 2, + end_line: 4, + vector: vec![1.0, 0.0], + }, + ], + }; + // Both symbols verify as exact (the caller mentions `ping_backend` in its + // body), and only the caller is a semantic candidate. Plain RRF would + // rank the caller first for appearing in both lists; the definition + // must keep the lead the lexical ranking gave it. + let hits = index.hybrid_search( + workspace.path(), + "ping_backend", + &[1.0, 0.0], + &lexical, + None, + 5, + 16384, + ); + let symbols: Vec<&str> = hits.iter().map(|hit| hit.symbol.as_str()).collect(); + assert_eq!(&symbols[..2], ["ping_backend", "ping_backend_retries"]); + assert!(hits[0].exact && hits[1].exact); + } + #[test] fn hybrid_search_keeps_semantic_lead_for_a_conceptual_query() { let workspace = tempfile::tempdir().unwrap();