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..6b75b45d --- /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-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.541 | 0.567 | +| overall Recall@5 | 0.444 | 0.667 | 0.667 | +| 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. + +| index | cold build | warm build | on disk | +|---|---|---|---| +| 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`) +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 + } + ] +} diff --git a/src/bin/orangu/dispatch.rs b/src/bin/orangu/dispatch.rs index 8def9036..38c8b8d1 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(); @@ -2222,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(), ); @@ -2233,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/chunk.rs b/src/embeddings/chunk.rs index ad512a43..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 @@ -66,13 +70,52 @@ 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. +/// 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(); @@ -82,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); - body.truncate(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, @@ -94,6 +140,7 @@ pub fn chunk_file( start_line: start, end_line: end, text, + body, }) }) .collect() @@ -116,4 +163,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"); + } } diff --git a/src/embeddings/index.rs b/src/embeddings/index.rs index 7e7da26b..0285a318 100644 --- a/src/embeddings/index.rs +++ b/src/embeddings/index.rs @@ -31,11 +31,12 @@ 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::lexical::{LexicalDoc, LexicalIndex}; +use super::source::{SourceFile, scan_workspace}; +use crate::graph::extract::GraphExtractor; use crate::graph::store::GraphStore; const INDEX_VERSION: u32 = 1; @@ -65,7 +66,17 @@ 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; + +/// 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; + +/// 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 ──────────────────────────────────────────────────────────── @@ -106,6 +117,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 @@ -221,6 +234,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 +265,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 +278,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 +289,7 @@ impl EmbeddingIndex { { kept.extend(chunks.iter().cloned()); } else { - stale_files.push((path, rel, hash)); + stale_files.push((file.path.clone(), rel, hash)); } } @@ -449,9 +477,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], @@ -463,6 +491,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 @@ -488,10 +561,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. @@ -504,18 +573,14 @@ 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(); + // 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); - // 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; @@ -542,23 +607,143 @@ 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); + } + + // 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() + .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_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) { + 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| { + exact_rank(&a.id) + .cmp(&exact_rank(&b.id)) + .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, + ), + ); } } } @@ -566,39 +751,46 @@ 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| { + exact_rank(&a.id) + .cmp(&exact_rank(&b.id)) + .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) } } /// 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) { - 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); } } } -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(), @@ -607,9 +799,54 @@ fn hit_from(chunk: &EmbeddedChunk, score: f32, expanded_from: Option) -> end_line: chunk.end_line, score, expanded_from, + exact, } } +/// 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, + exact: bool, +) -> 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, + exact, + } +} + +/// 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 { @@ -708,35 +945,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::*; @@ -749,6 +957,7 @@ mod tests { start_line: 1, end_line: 1, text: "x".repeat(chars), + body: "x".repeat(chars), } } @@ -885,4 +1094,193 @@ mod tests { // "three" should beat "two" due to the diversity penalty 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(); + 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_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(); + 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 new file mode 100644 index 00000000..f4e01516 --- /dev/null +++ b/src/embeddings/lexical.rs @@ -0,0 +1,1071 @@ +// 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, 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; + +/// 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. +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 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, + /// Verified as an exact identifier or quoted-phrase match in the source. + pub exact: bool, +} + +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, + exact: self.exact, + } + } +} + +/// 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() + } + + /// 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; + /// 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 + } + + /// 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 parsed = parse_query(query); + if parsed.terms.is_empty() { + return Vec::new(); + } + 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() { + let is_exact = exact.contains(&index); + if score <= 0.0 && !is_exact { + 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, + exact: is_exact, + }); + } + 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 + } + + /// 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()), + false, + ), + ); + } + } + } + } + } + + let mut hits: Vec = best.into_values().collect(); + 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) + } + + /// 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 + } + + /// 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, + } +} + +/// 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(); + 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) { + 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 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(); + 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")); + } + + #[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 74d76d86..b414f22b 100644 --- a/src/embeddings/mod.rs +++ b/src/embeddings/mod.rs @@ -28,7 +28,9 @@ pub mod chunk; pub mod client; pub mod index; +pub mod lexical; pub mod sketch; +pub mod source; use crate::config::LlmConfiguration; use crate::llm::normalized_openai_endpoint; @@ -40,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, } @@ -49,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 { @@ -63,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, @@ -83,30 +100,109 @@ impl SemanticSearch { !query.trim().is_empty(), "semantic search query must not be empty" ); - let index = EmbeddingIndex::build_or_update( + 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 scan = crate::embeddings::source::scan_workspace(workspace); + let lexical = crate::embeddings::lexical::LexicalIndex::build_or_update_from_scan( workspace, - &self.client, + &scan, + self.compile_workers, + None, + ) + .ok(); + // 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, cancel, self.compile_workers, 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: false, }; - if result.cancelled || index.is_empty() { + if result.cancelled { return Ok(result); } - let query_vector = self.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 = 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) } @@ -157,24 +253,32 @@ 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})"), 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, )); } 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")); + } +} diff --git a/src/tools.rs b/src/tools.rs index 92e4c35b..08841b67 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 } }, @@ -1342,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? } @@ -4006,6 +4012,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(); @@ -4107,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, @@ -4117,13 +4151,70 @@ 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 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}"); + } }