919 lines
33 KiB
Rust
919 lines
33 KiB
Rust
//! Shared offline hybrid search using real Google embeddings and persisted,
|
|
//! filtered HNSW retrieval. Production-scale validation remains a release gate.
|
|
mod coordination;
|
|
mod media;
|
|
pub use coordination::{
|
|
BudgetedProvider, InferenceBudget, InferencePermit, InferenceRole, PendingForeground,
|
|
};
|
|
mod prepared;
|
|
pub use prepared::PreparedGraphs;
|
|
mod pipeline;
|
|
mod routing;
|
|
pub use routing::SemanticOrigin;
|
|
mod store;
|
|
pub use searchi_media::{FfmpegTools, MediaLimits};
|
|
|
|
use anyhow::{ensure, Result};
|
|
use searchi_core::metrics::{MetricsMode, SpanToken, TraceCollector};
|
|
use searchi_core::SearchIndex;
|
|
pub use searchi_core::{SearchMetrics, SearchOptions, SearchResponse, SearchResult};
|
|
pub use searchi_semantic::WorkerInfo;
|
|
use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
|
|
use serde::Serialize;
|
|
use std::{
|
|
collections::{HashMap, VecDeque},
|
|
path::Path,
|
|
};
|
|
use store::{Store, MAX_SCAN_CHUNKS};
|
|
|
|
const MAX_BATCH: usize = 8;
|
|
const QUERY_CACHE_ENTRIES: usize = 32;
|
|
const BYTE_SLICE: usize = 64 * 1024;
|
|
const BYTE_OVERLAP: usize = 512;
|
|
const MAX_PASSAGE_WINDOWS: usize = 2048;
|
|
const CHUNK_POLICY: &str = "byte64k-overlap512-title512-v1";
|
|
fn generation_fingerprint(model: &str) -> String {
|
|
format!(
|
|
"{model}:{CHUNK_POLICY}:{}:image280-frame70-audio30s-overlap2-distributed240-scene120-v1",
|
|
searchi_semantic::GEMMA_MEDIA_FINGERPRINT
|
|
)
|
|
}
|
|
|
|
fn reusable_policy(expected: &str, active: &str) -> bool {
|
|
if active == expected || active == format!("{expected}:text-only") {
|
|
return true;
|
|
}
|
|
let Some(policy) = active.strip_prefix(&format!("{expected}:custom-")) else {
|
|
return false;
|
|
};
|
|
let values = policy
|
|
.split('-')
|
|
.map(str::parse::<u64>)
|
|
.collect::<std::result::Result<Vec<_>, _>>();
|
|
matches!(values,Ok(v) if v.len()==8 && v[0]>=2 && (5000..=30000).contains(&v[3]) && v[4]<v[3] && v[5]>0 && v[6]<=u32::MAX as u64)
|
|
}
|
|
|
|
/// Resolve compatible publication coverage without broadening model or window
|
|
/// policy compatibility. A retained reader can start under either coverage.
|
|
fn policy_base(policy: &str) -> &str {
|
|
if let Some(base) = policy.strip_suffix(":text-only") {
|
|
return base;
|
|
}
|
|
if let Some((base, _)) = policy.rsplit_once(":custom-") {
|
|
if reusable_policy(base, policy) {
|
|
return base;
|
|
}
|
|
}
|
|
policy
|
|
}
|
|
|
|
fn compatible_query_policy(expected: &str, active: &str) -> bool {
|
|
reusable_policy(policy_base(expected), active)
|
|
}
|
|
|
|
fn policy_coverage(expected: &str, active: &str) -> Option<SemanticCoverage> {
|
|
compatible_query_policy(expected, active).then(|| {
|
|
if active.ends_with(":text-only") {
|
|
SemanticCoverage::TextOnly
|
|
} else {
|
|
SemanticCoverage::FullMedia
|
|
}
|
|
})
|
|
}
|
|
|
|
/// Requested indexing coverage is separate from warmed model readiness.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum SemanticCoverage {
|
|
TextOnly,
|
|
FullMedia,
|
|
}
|
|
|
|
/// Model-specific query/document prompts and actual tokenizer windows belong
|
|
/// to the provider. Storage and ranking never assume a particular dimension.
|
|
pub trait EmbeddingProvider: Send {
|
|
/// Maximum document batch admitted as one inference call. Coordinated
|
|
/// background sessions yield at every document boundary.
|
|
fn document_batch_limit(&self) -> usize {
|
|
MAX_BATCH
|
|
}
|
|
fn fingerprint(&self) -> &str;
|
|
fn dimensions(&self) -> usize;
|
|
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
|
|
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>>;
|
|
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>>;
|
|
fn embed_code_query(&mut self, _query: &str) -> Result<Vec<f32>> {
|
|
anyhow::bail!("Provider does not support Google code retrieval")
|
|
}
|
|
fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
|
|
self.embed_code_query(query)
|
|
}
|
|
fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
|
|
self.embed_query(query)
|
|
}
|
|
fn embed_query_traced(
|
|
&mut self,
|
|
query: &str,
|
|
debug: bool,
|
|
_trace: &mut TraceCollector,
|
|
_parent: Option<SpanToken>,
|
|
) -> Result<Vec<f32>> {
|
|
if debug {
|
|
self.embed_query_debug(query)
|
|
} else {
|
|
self.embed_query(query)
|
|
}
|
|
}
|
|
fn embed_code_query_traced(
|
|
&mut self,
|
|
query: &str,
|
|
debug: bool,
|
|
_trace: &mut TraceCollector,
|
|
_parent: Option<SpanToken>,
|
|
) -> Result<Vec<f32>> {
|
|
if debug {
|
|
self.embed_code_query_debug(query)
|
|
} else {
|
|
self.embed_code_query(query)
|
|
}
|
|
}
|
|
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
|
|
Default::default()
|
|
}
|
|
fn embed_images(&mut self, _paths: &[&Path], _budget: u32) -> Result<Vec<Vec<f32>>> {
|
|
anyhow::bail!("Provider does not support real images")
|
|
}
|
|
fn embed_audio(&mut self, _pcm: &[f32]) -> Result<Vec<f32>> {
|
|
anyhow::bail!("Provider does not support real audio")
|
|
}
|
|
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
|
|
None
|
|
}
|
|
}
|
|
|
|
impl EmbeddingProvider for GemmaEmbedder {
|
|
fn fingerprint(&self) -> &str {
|
|
GEMMA_FINGERPRINT
|
|
}
|
|
fn dimensions(&self) -> usize {
|
|
GEMMA_DIMENSIONS
|
|
}
|
|
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
|
|
self.split_windows(text, title)
|
|
}
|
|
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
|
|
self.embed_documents(texts, title)
|
|
}
|
|
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
|
|
self.embed_query(query)
|
|
}
|
|
fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
|
|
GemmaEmbedder::embed_query_debug(self, query)
|
|
}
|
|
fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
|
|
GemmaEmbedder::embed_code_query(self, query)
|
|
}
|
|
fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
|
|
GemmaEmbedder::embed_code_query_debug(self, query)
|
|
}
|
|
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
|
|
GemmaEmbedder::query_diagnostics(self)
|
|
}
|
|
fn embed_images(&mut self, paths: &[&Path], budget: u32) -> Result<Vec<Vec<f32>>> {
|
|
self.embed_images_with_budget(paths, budget)
|
|
}
|
|
fn embed_audio(&mut self, pcm: &[f32]) -> Result<Vec<f32>> {
|
|
self.embed_audio_pcm(pcm)
|
|
}
|
|
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
|
|
Some(self.info.clone())
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct MediaFailure {
|
|
pub path: String,
|
|
pub message: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct EmbeddingReport {
|
|
pub media_files: usize,
|
|
pub media_failures: usize,
|
|
pub failures: Vec<MediaFailure>,
|
|
pub passages: usize,
|
|
pub chunks: usize,
|
|
pub generation: i64,
|
|
pub model_fingerprint: String,
|
|
/// Completed groups copied or resumed without document/media inference.
|
|
pub reused_sources: usize,
|
|
pub reused_chunks: usize,
|
|
pub encoded_sources: usize,
|
|
pub document_inference_calls: usize,
|
|
pub image_inference_calls: usize,
|
|
pub audio_inference_calls: usize,
|
|
pub resumed: bool,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct EmbeddingStatus {
|
|
/// Coverage of a compatible, prepared active generation, if any.
|
|
pub coverage: Option<SemanticCoverage>,
|
|
pub requested_coverage: SemanticCoverage,
|
|
/// Early setup must preserve an existing compatible full generation.
|
|
pub text_indexing_allowed: bool,
|
|
pub failures: Vec<MediaFailure>,
|
|
pub generation: Option<i64>,
|
|
pub chunks: usize,
|
|
pub model_fingerprint: Option<String>,
|
|
pub stale: bool,
|
|
/// Rebuild is generation-based: all current passages are pending when stale.
|
|
pub pending: usize,
|
|
}
|
|
|
|
#[derive(Clone, Copy, PartialEq, Eq)]
|
|
enum QueryTask {
|
|
Search,
|
|
Code,
|
|
}
|
|
/// Shared local hybrid service. Google EmbeddingGemma 2 is the default and
|
|
/// only shipped provider; no alternative model is silently substituted.
|
|
pub struct HybridSearch {
|
|
lexical: SearchIndex,
|
|
embedder: Box<dyn EmbeddingProvider>,
|
|
store: Store,
|
|
generation_fingerprint: String,
|
|
// Real provider outputs only, scoped to this model session. Queries never
|
|
// persist and results/filters are never cached here.
|
|
query_cache_fingerprint: String,
|
|
query_cache: VecDeque<(QueryTask, String, Vec<f32>)>,
|
|
}
|
|
|
|
impl HybridSearch {
|
|
/// Attach shared admission without replacing the warmed model or query LRU.
|
|
pub fn with_budget(
|
|
mut self,
|
|
budget: &InferenceBudget,
|
|
role: InferenceRole,
|
|
actual_worker_threads: usize,
|
|
) -> Self {
|
|
self.embedder = Box::new(BudgetedProvider::new(
|
|
self.embedder,
|
|
budget,
|
|
role,
|
|
actual_worker_threads,
|
|
));
|
|
self
|
|
}
|
|
/// Connect background publication to an existing warmed reader.
|
|
pub fn share_prepared_graphs(&mut self, graphs: &PreparedGraphs) -> Result<()> {
|
|
self.store.share_prepared_graphs(graphs)
|
|
}
|
|
|
|
/// Actual model/runtime/device reported by the persistent native worker.
|
|
pub fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
|
|
self.embedder.worker_info()
|
|
}
|
|
/// Cheap local status check. No model, runtime worker, or network is loaded.
|
|
/// Corpus revision triggers detect file inserts, updates, and deletions.
|
|
pub fn embedding_status(db: impl AsRef<Path>) -> Result<EmbeddingStatus> {
|
|
Self::embedding_status_for_coverage(db, SemanticCoverage::FullMedia)
|
|
}
|
|
|
|
pub fn embedding_status_for_coverage(
|
|
db: impl AsRef<Path>,
|
|
requested_coverage: SemanticCoverage,
|
|
) -> Result<EmbeddingStatus> {
|
|
let connection =
|
|
rusqlite::Connection::open_with_flags(db, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
|
|
connection.busy_timeout(std::time::Duration::from_secs(5))?;
|
|
let snapshot = connection.unchecked_transaction()?;
|
|
let exists: bool = snapshot.query_row(
|
|
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name='semantic_meta')",
|
|
[],
|
|
|row| row.get(0),
|
|
)?;
|
|
let pending: usize = match requested_coverage {
|
|
SemanticCoverage::TextOnly => snapshot.query_row("SELECT count(*) FROM passages", [], |row| row.get(0))?,
|
|
SemanticCoverage::FullMedia => snapshot.query_row("SELECT (SELECT count(*) FROM passages)+(SELECT count(*) FROM files WHERE extension IN ('png','jpg','jpeg','webp','bmp','mp4','mkv','mov','avi','webm','m4v','mpeg','mpg','mp3','flac','wav','m4a','ogg','opus','aac','aiff','aif'))", [], |row| row.get(0))?,
|
|
};
|
|
if !exists {
|
|
return Ok(EmbeddingStatus {
|
|
coverage: None,
|
|
requested_coverage,
|
|
text_indexing_allowed: true,
|
|
failures: Vec::new(),
|
|
generation: None,
|
|
chunks: 0,
|
|
model_fingerprint: None,
|
|
stale: true,
|
|
pending,
|
|
});
|
|
}
|
|
let version: i64 = snapshot.query_row(
|
|
"SELECT schema_version FROM semantic_meta WHERE singleton=1",
|
|
[],
|
|
|row| row.get(0),
|
|
)?;
|
|
if version < 6 {
|
|
return Ok(EmbeddingStatus {
|
|
coverage: None,
|
|
requested_coverage,
|
|
text_indexing_allowed: true,
|
|
failures: Vec::new(),
|
|
generation: None,
|
|
chunks: 0,
|
|
model_fingerprint: None,
|
|
stale: true,
|
|
pending,
|
|
});
|
|
}
|
|
ensure!(
|
|
(6..=8).contains(&version),
|
|
"Unsupported semantic schema version {version}"
|
|
);
|
|
let (generation,fingerprint,revision,current,dimensions):(Option<i64>,Option<String>,Option<i64>,i64,Option<usize>)=snapshot.query_row(
|
|
"SELECT g.id,g.fingerprint,CASE WHEN g.ann_blob IS NOT NULL AND g.dimensions=768 THEN g.corpus_revision END,m.corpus_revision,g.dimensions FROM semantic_meta m LEFT JOIN semantic_generations g ON g.id=m.active_generation WHERE m.singleton=1",[],
|
|
|row| Ok((row.get(0)?,row.get(1)?,row.get(2)?,row.get(3)?,row.get(4)?)))?;
|
|
let chunks: usize = snapshot.query_row(
|
|
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
|
|
[generation],
|
|
|row| row.get(0),
|
|
)?;
|
|
let mut failures_query=snapshot.prepare("SELECT path,message FROM semantic_media_errors WHERE generation=?1 ORDER BY path LIMIT 1000")?;
|
|
let failures = failures_query
|
|
.query_map([generation], |r| {
|
|
Ok(MediaFailure {
|
|
path: r.get(0)?,
|
|
message: r.get(1)?,
|
|
})
|
|
})?
|
|
.collect::<rusqlite::Result<Vec<_>>>()?;
|
|
let expected = generation_fingerprint(GEMMA_FINGERPRINT);
|
|
let coverage = if dimensions == Some(GEMMA_DIMENSIONS) {
|
|
fingerprint
|
|
.as_deref()
|
|
.and_then(|active| policy_coverage(&expected, active))
|
|
} else {
|
|
None
|
|
};
|
|
let prepared_coverage = if revision.is_some() { coverage } else { None };
|
|
let policy_matches = match requested_coverage {
|
|
SemanticCoverage::TextOnly => prepared_coverage.is_some(),
|
|
// Keep the default full-media status's exact configured policy rule.
|
|
SemanticCoverage::FullMedia => fingerprint.as_deref() == Some(expected.as_str()),
|
|
};
|
|
let stale = revision != Some(current) || !policy_matches;
|
|
Ok(EmbeddingStatus {
|
|
coverage: prepared_coverage,
|
|
requested_coverage,
|
|
text_indexing_allowed: coverage != Some(SemanticCoverage::FullMedia),
|
|
failures,
|
|
generation,
|
|
chunks,
|
|
model_fingerprint: fingerprint,
|
|
stale,
|
|
pending: if stale { pending } else { 0 },
|
|
})
|
|
}
|
|
|
|
pub fn open(
|
|
db: impl AsRef<Path>,
|
|
model_dir: impl AsRef<Path>,
|
|
runtime_library: impl AsRef<Path>,
|
|
threads: usize,
|
|
) -> Result<Self> {
|
|
let embedder = GemmaEmbedder::load(model_dir, runtime_library, threads)?;
|
|
Self::with_provider(db, Box::new(embedder))
|
|
}
|
|
|
|
/// Foreground-only Google session. Indexing must use `open` instead.
|
|
pub fn open_query(
|
|
db: impl AsRef<Path>,
|
|
model_dir: impl AsRef<Path>,
|
|
runtime_library: impl AsRef<Path>,
|
|
threads: usize,
|
|
) -> Result<Self> {
|
|
let embedder = GemmaEmbedder::load_query(model_dir, runtime_library, threads)?;
|
|
Self::with_provider(db, Box::new(embedder))
|
|
}
|
|
|
|
/// Explicit provider injection for future model generations, never fallback.
|
|
pub fn with_provider(
|
|
db: impl AsRef<Path>,
|
|
embedder: Box<dyn EmbeddingProvider>,
|
|
) -> Result<Self> {
|
|
let lexical = SearchIndex::open(&db)?;
|
|
let store = Store::open(db.as_ref(), embedder.dimensions())?;
|
|
let expected = generation_fingerprint(embedder.fingerprint());
|
|
let active = store.active_fingerprint()?;
|
|
let generation_fingerprint = match active {
|
|
Some(active) if reusable_policy(&expected, &active) => active,
|
|
_ => expected,
|
|
};
|
|
Ok(Self {
|
|
query_cache_fingerprint: embedder.fingerprint().to_owned(),
|
|
query_cache: VecDeque::new(),
|
|
lexical,
|
|
embedder,
|
|
store,
|
|
generation_fingerprint,
|
|
})
|
|
}
|
|
|
|
/// Execute real query inference and load the active ANN generation before
|
|
/// foreground use. Call on a background worker after assets are verified.
|
|
pub fn warmup(&mut self) -> Result<()> {
|
|
self.lexical.prewarm()?;
|
|
let vector = self.embedder.embed_query("warm up local file search")?;
|
|
let _ = self.store.search(
|
|
&vector,
|
|
&self.generation_fingerprint,
|
|
1,
|
|
&SearchOptions::default(),
|
|
)?;
|
|
Ok(())
|
|
}
|
|
|
|
/// Build or resume a replacement generation in bounded batches. Completed
|
|
/// groups can be reused; failures retain staging work and the ready reader.
|
|
pub fn index_embeddings(&mut self) -> Result<EmbeddingReport> {
|
|
let fingerprint = format!(
|
|
"{}:text-only",
|
|
generation_fingerprint(self.embedder.fingerprint())
|
|
);
|
|
self.rebuild_generation(fingerprint, None, |service, generation, work| {
|
|
service.build(generation, work)
|
|
})
|
|
}
|
|
|
|
/// Early text setup never replaces a compatible full generation. The
|
|
/// decision is made under the publication build lock, not from UI status.
|
|
/// `None` means full-media indexing must resume after its assets are ready.
|
|
pub fn index_text_embeddings_preserving_media(&mut self) -> Result<Option<EmbeddingReport>> {
|
|
let fingerprint = format!(
|
|
"{}:text-only",
|
|
generation_fingerprint(self.embedder.fingerprint())
|
|
);
|
|
self.rebuild_generation_preserving_media(fingerprint, |service, generation, work| {
|
|
service.build(generation, work)
|
|
})
|
|
}
|
|
|
|
/// Filter each retrieval lane before ranking, preserve filename priority,
|
|
/// then fuse keyword and semantic ranks with reciprocal-rank fusion (k=60).
|
|
pub fn search(
|
|
&mut self,
|
|
query: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
) -> Result<Vec<SearchResult>> {
|
|
let mut trace = TraceCollector::new(MetricsMode::Basic);
|
|
self.search_traced(query, limit, options, &mut trace, None)
|
|
}
|
|
pub fn search_measured(
|
|
&mut self,
|
|
query: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
debug: bool,
|
|
) -> Result<SearchResponse<SearchResult>> {
|
|
let mut trace = TraceCollector::new(if debug {
|
|
MetricsMode::Debug
|
|
} else {
|
|
MetricsMode::Basic
|
|
});
|
|
let results = self.search_traced(query, limit, options, &mut trace, None)?;
|
|
let metrics = trace.finish(results.len());
|
|
Ok(SearchResponse { results, metrics })
|
|
}
|
|
pub fn search_traced(
|
|
&mut self,
|
|
query: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
trace: &mut TraceCollector,
|
|
parent: Option<SpanToken>,
|
|
) -> Result<Vec<SearchResult>> {
|
|
trace.timed("hybrid.search", parent, |trace, parent| {
|
|
self.search_inner(query, limit, options, trace, parent)
|
|
})
|
|
}
|
|
fn search_inner(
|
|
&mut self,
|
|
query: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
trace: &mut TraceCollector,
|
|
parent: Option<SpanToken>,
|
|
) -> Result<Vec<SearchResult>> {
|
|
let candidates = limit.saturating_mul(4).clamp(20, 1000);
|
|
// Core validates filters and query budgets even for empty requests.
|
|
let mut lexical = self
|
|
.lexical
|
|
.search_with_options_traced(query, candidates, options, trace, parent)?;
|
|
if query.trim().is_empty() || limit == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
trace.timed("hybrid.identifier_priority", parent, |_, _| {
|
|
promote_identifiers(query.trim(), &mut lexical)
|
|
});
|
|
let vector = self.query_vector(query, options.code_only, trace, parent)?;
|
|
let semantic =
|
|
self.routed_semantic_traced(query, &vector, candidates, options, trace, parent)?;
|
|
let results = trace.timed("hybrid.rank_fusion", parent, |trace, span| {
|
|
trace.counter(span, "lexical_files", lexical.len() as u64);
|
|
trace.counter(span, "semantic_files", semantic.len() as u64);
|
|
fuse(lexical, semantic, limit.min(1000))
|
|
});
|
|
Ok(results)
|
|
}
|
|
|
|
/// Semantic lane for reproducible relevance evaluation and diagnostics.
|
|
pub fn search_semantic(
|
|
&mut self,
|
|
query: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
) -> Result<Vec<SearchResult>> {
|
|
let _ = self.lexical.search_with_options(query, 0, options)?;
|
|
if query.trim().is_empty() || limit == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut trace = TraceCollector::new(MetricsMode::Basic);
|
|
let vector = self.query_vector(query, options.code_only, &mut trace, None)?;
|
|
self.routed_semantic(query, &vector, limit.min(1000), options)
|
|
}
|
|
|
|
fn query_vector(
|
|
&mut self,
|
|
query: &str,
|
|
code: bool,
|
|
trace: &mut TraceCollector,
|
|
parent: Option<SpanToken>,
|
|
) -> Result<Vec<f32>> {
|
|
let query = query.trim();
|
|
ensure!(
|
|
!query.is_empty() && query.len() <= 4096,
|
|
"Query must be nonempty and at most 4096 UTF-8 bytes"
|
|
);
|
|
let span = trace.begin("semantic.query_embedding", parent);
|
|
let task = if code {
|
|
QueryTask::Code
|
|
} else {
|
|
QueryTask::Search
|
|
};
|
|
trace.counter(span, "code_task", u64::from(code));
|
|
if self.query_cache_fingerprint != self.embedder.fingerprint() {
|
|
self.query_cache.clear();
|
|
self.query_cache_fingerprint = self.embedder.fingerprint().to_owned();
|
|
}
|
|
if let Some(position) = self
|
|
.query_cache
|
|
.iter()
|
|
.position(|(cached_task, cached, _)| *cached_task == task && cached == query)
|
|
{
|
|
let entry = self
|
|
.query_cache
|
|
.remove(position)
|
|
.expect("Located cache entry");
|
|
let vector = entry.2.clone();
|
|
self.query_cache.push_back(entry);
|
|
trace.counter(span, "query_cache_hit", 1);
|
|
trace.counter(span, "query_cache_miss", 0);
|
|
trace.counter(span, "failed", 0);
|
|
// No worker was called: never attach its previous diagnostics.
|
|
trace.end(span, Default::default());
|
|
return Ok(vector);
|
|
}
|
|
trace.counter(span, "query_cache_hit", 0);
|
|
trace.counter(span, "query_cache_miss", 1);
|
|
let vector = (|| {
|
|
let vector = if code {
|
|
self.embedder
|
|
.embed_code_query_traced(query, trace.is_debug(), trace, span)?
|
|
} else {
|
|
self.embedder
|
|
.embed_query_traced(query, trace.is_debug(), trace, span)?
|
|
};
|
|
ensure!(
|
|
vector.len() == self.embedder.dimensions()
|
|
&& vector.iter().all(|value| value.is_finite())
|
|
&& vector.iter().any(|value| *value != 0.0),
|
|
"Provider returned an invalid query vector"
|
|
);
|
|
Ok(vector)
|
|
})();
|
|
trace.counter(span, "failed", u64::from(vector.is_err()));
|
|
trace.end(
|
|
span,
|
|
if trace.is_debug() {
|
|
self.embedder.query_diagnostics()
|
|
} else {
|
|
Default::default()
|
|
},
|
|
);
|
|
let vector = vector?;
|
|
if self.query_cache.len() == QUERY_CACHE_ENTRIES {
|
|
self.query_cache.pop_front();
|
|
}
|
|
self.query_cache
|
|
.push_back((task, query.to_owned(), vector.clone()));
|
|
Ok(vector)
|
|
}
|
|
}
|
|
|
|
fn bounded_title(title: &str) -> &str {
|
|
let mut end = title.len().min(512);
|
|
while !title.is_char_boundary(end) {
|
|
end -= 1;
|
|
}
|
|
&title[..end]
|
|
}
|
|
|
|
fn byte_slices(text: &str) -> Result<Vec<(usize, usize)>> {
|
|
ensure!(
|
|
text.len() <= 2 * 1024 * 1024,
|
|
"Passage exceeds 2 MiB service byte budget"
|
|
);
|
|
let mut slices = Vec::new();
|
|
let mut start = 0;
|
|
while start < text.len() {
|
|
let mut end = (start + BYTE_SLICE).min(text.len());
|
|
while !text.is_char_boundary(end) {
|
|
end -= 1;
|
|
}
|
|
ensure!(end > start, "UTF-8 slice made no progress");
|
|
slices.push((start, end));
|
|
if end == text.len() {
|
|
break;
|
|
}
|
|
start = end.saturating_sub(BYTE_OVERLAP);
|
|
while !text.is_char_boundary(start) {
|
|
start -= 1;
|
|
}
|
|
}
|
|
Ok(slices)
|
|
}
|
|
|
|
fn passage_windows(
|
|
provider: &mut dyn EmbeddingProvider,
|
|
text: &str,
|
|
title: &str,
|
|
) -> Result<Vec<TextWindow>> {
|
|
let mut windows = Vec::new();
|
|
for (start, end) in byte_slices(text)? {
|
|
let fragment = &text[start..end];
|
|
if fragment.trim().is_empty() {
|
|
continue;
|
|
}
|
|
let mut slice = provider.split_windows(fragment, title)?;
|
|
for window in &mut slice {
|
|
window.start_byte += start;
|
|
window.end_byte += start;
|
|
}
|
|
windows.extend(slice);
|
|
ensure!(
|
|
windows.len() <= MAX_PASSAGE_WINDOWS,
|
|
"Passage exceeds {MAX_PASSAGE_WINDOWS} semantic window budget"
|
|
);
|
|
}
|
|
Ok(windows)
|
|
}
|
|
|
|
fn promote_identifiers(query: &str, results: &mut [SearchResult]) {
|
|
if query.chars().any(char::is_whitespace)
|
|
|| !query.chars().any(|c| c.is_ascii_digit() || c == '_')
|
|
{
|
|
return;
|
|
}
|
|
for result in results {
|
|
if result.match_kind == "content" && literal_identifier_match(&result.snippet, query) {
|
|
result.match_kind = "exact_identifier".into();
|
|
}
|
|
}
|
|
}
|
|
|
|
fn literal_identifier_match(text: &str, query: &str) -> bool {
|
|
text.match_indices(query).any(|(start, _)| {
|
|
let end = start + query.len();
|
|
let is_identifier = |c: char| c.is_alphanumeric() || c == '_' || c == '-';
|
|
!text[..start].chars().next_back().is_some_and(is_identifier)
|
|
&& !text[end..].chars().next().is_some_and(is_identifier)
|
|
})
|
|
}
|
|
|
|
fn fuse(
|
|
lexical: Vec<SearchResult>,
|
|
semantic: Vec<SearchResult>,
|
|
limit: usize,
|
|
) -> Vec<SearchResult> {
|
|
let mut results = HashMap::<String, SearchResult>::new();
|
|
for (rank, mut result) in lexical.into_iter().enumerate() {
|
|
let weight = if result.match_kind == "filename" {
|
|
1.15
|
|
} else {
|
|
1.0
|
|
};
|
|
result.score = weight / (61.0 + rank as f64);
|
|
results.insert(result.path.clone(), result);
|
|
}
|
|
for (rank, mut result) in semantic.into_iter().enumerate() {
|
|
let contribution = 1.0 / (61.0 + rank as f64);
|
|
if let Some(existing) = results.get_mut(&result.path) {
|
|
existing.score += contribution;
|
|
if existing.match_kind == "content" {
|
|
existing.match_kind = "hybrid".into();
|
|
}
|
|
} else {
|
|
result.score = contribution;
|
|
results.insert(result.path.clone(), result);
|
|
}
|
|
}
|
|
let tier = |result: &SearchResult| match result.match_kind.as_str() {
|
|
"exact_filename" => 0,
|
|
"exact_identifier" => 1,
|
|
_ => 2,
|
|
};
|
|
let mut results: Vec<_> = results.into_values().collect();
|
|
results.sort_by(|a, b| {
|
|
tier(a)
|
|
.cmp(&tier(b))
|
|
.then_with(|| b.score.total_cmp(&a.score))
|
|
.then_with(|| a.path.cmp(&b.path))
|
|
});
|
|
results.truncate(limit);
|
|
results
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
#[test]
|
|
fn failed_query_inference_retains_diagnostics_in_shared_trace() -> Result<()> {
|
|
// Failure-only provider: no fabricated embeddings are generated.
|
|
struct FailingProvider;
|
|
impl EmbeddingProvider for FailingProvider {
|
|
fn fingerprint(&self) -> &str {
|
|
"failure-injection"
|
|
}
|
|
fn dimensions(&self) -> usize {
|
|
768
|
|
}
|
|
fn split_windows(&mut self, _: &str, _: &str) -> Result<Vec<TextWindow>> {
|
|
anyhow::bail!("unused")
|
|
}
|
|
fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
|
|
anyhow::bail!("unused")
|
|
}
|
|
fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
|
|
anyhow::bail!("injected worker failure")
|
|
}
|
|
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
|
|
std::collections::BTreeMap::from([("request_write_us".into(), 7)])
|
|
}
|
|
}
|
|
let directory = tempfile::tempdir()?;
|
|
let db = directory.path().join("index.db");
|
|
let mut service = HybridSearch::with_provider(&db, Box::new(FailingProvider))?;
|
|
let mut trace = TraceCollector::new(MetricsMode::Debug);
|
|
assert!(service
|
|
.search_traced(
|
|
"ordinary query",
|
|
5,
|
|
&SearchOptions::default(),
|
|
&mut trace,
|
|
None
|
|
)
|
|
.is_err());
|
|
let metrics = trace.finish(0);
|
|
let spans = metrics.trace.unwrap();
|
|
let embedding = spans
|
|
.iter()
|
|
.find(|s| s.name == "semantic.query_embedding")
|
|
.unwrap();
|
|
assert_eq!(embedding.counters["failed"], 1);
|
|
assert_eq!(embedding.counters["request_write_us"], 7);
|
|
let parent = &spans[embedding.parent.unwrap()];
|
|
assert!(
|
|
embedding.start_ms + embedding.duration_ms
|
|
<= parent.start_ms + parent.duration_ms + 1e-6
|
|
);
|
|
Ok(())
|
|
}
|
|
fn result(path: &str, kind: &str) -> SearchResult {
|
|
SearchResult {
|
|
path: path.into(),
|
|
filename: path.into(),
|
|
snippet: "evidence".into(),
|
|
score: 1.0,
|
|
match_kind: kind.into(),
|
|
source: None,
|
|
extraction_error: None,
|
|
}
|
|
}
|
|
#[test]
|
|
fn exact_filename_wins_even_when_another_result_matches_both_lanes() {
|
|
let results = fuse(
|
|
vec![
|
|
result("invoice.pdf", "exact_filename"),
|
|
result("notes.txt", "content"),
|
|
],
|
|
vec![
|
|
result("notes.txt", "semantic"),
|
|
result("another.txt", "semantic"),
|
|
],
|
|
10,
|
|
);
|
|
assert_eq!(results[0].path, "invoice.pdf");
|
|
assert_eq!(results[1].match_kind, "hybrid");
|
|
assert_eq!(results.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn identifiers_use_literal_boundaries_and_survive_semantic_competition() {
|
|
let mut lexical = vec![result("support.txt", "content")];
|
|
lexical[0].snippet = "Error E_SEARCH_071 means upgrade the app".into();
|
|
promote_identifiers("E_SEARCH_071", &mut lexical);
|
|
let results = fuse(lexical, vec![result("wrong.txt", "semantic")], 10);
|
|
assert_eq!(results[0].match_kind, "exact_identifier");
|
|
assert!(!literal_identifier_match(
|
|
"E_SEARCH_071_EXTRA",
|
|
"E_SEARCH_071"
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn keyword_semantic_consensus_beats_weak_partial_filename() {
|
|
let results = fuse(
|
|
vec![
|
|
result("meeting-cancellation.txt", "filename"),
|
|
result("supplier-agreement.txt", "content"),
|
|
],
|
|
vec![result("supplier-agreement.txt", "semantic")],
|
|
10,
|
|
);
|
|
assert_eq!(results[0].path, "supplier-agreement.txt");
|
|
assert_eq!(results[0].match_kind, "hybrid");
|
|
assert_eq!(results[1].match_kind, "filename");
|
|
}
|
|
|
|
#[test]
|
|
fn byte_slices_preserve_utf8_overlap_offsets_and_complete_tail() {
|
|
let text = format!("{}TAIL", "Ä🐈中文 ".repeat(10000));
|
|
let slices = byte_slices(&text).unwrap();
|
|
assert!(slices.len() > 1);
|
|
assert_eq!(slices[0].0, 0);
|
|
assert_eq!(slices.last().unwrap().1, text.len());
|
|
for (start, end) in &slices {
|
|
assert!(text.is_char_boundary(*start) && text.is_char_boundary(*end));
|
|
assert!(end - start <= BYTE_SLICE);
|
|
}
|
|
for pair in slices.windows(2) {
|
|
assert!(pair[1].0 < pair[0].1 && pair[1].1 > pair[0].1);
|
|
}
|
|
assert!(text[slices.last().unwrap().0..].ends_with("TAIL"));
|
|
assert!(byte_slices(&"x".repeat(2 * 1024 * 1024 + 1)).is_err());
|
|
let title = "中文".repeat(120);
|
|
assert!(bounded_title(&title).len() <= 512 && title.starts_with(bounded_title(&title)));
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
|
|
fn google_cpu_preserves_and_retrieves_content_after_64k_byte_boundary() -> Result<()> {
|
|
let model =
|
|
std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
|
|
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
|
|
let mut provider =
|
|
GemmaEmbedder::load_with_device(model, worker, 1, searchi_semantic::Device::Cpu)?;
|
|
let text=format!("{}\nÜbernachtungskosten auf Dienstreisen werden gegen Vorlage der Hotelrechnung erstattet.","Routine system checks were completed successfully. ".repeat(1500));
|
|
assert!(text.len() > BYTE_SLICE);
|
|
let windows = passage_windows(&mut provider, &text, "routine-log.txt")?;
|
|
let last = windows.last().unwrap();
|
|
assert_eq!(last.end_byte, text.len());
|
|
assert!(last.start_byte >= BYTE_SLICE - BYTE_OVERLAP);
|
|
assert!(text.find("Übernachtungskosten").unwrap() > BYTE_SLICE);
|
|
assert!(last.text.contains("Hotelrechnung"));
|
|
assert_eq!(
|
|
text.get(last.start_byte..last.end_byte),
|
|
Some(last.text.as_str())
|
|
);
|
|
let vectors = provider.embed_documents(
|
|
&[windows[0].text.as_str(), last.text.as_str()],
|
|
"routine-log.txt",
|
|
)?;
|
|
let query = provider.embed_query(
|
|
"Where does it say that hotel stays during work travel can be reimbursed?",
|
|
)?;
|
|
let cosine = |vector: &[f32]| vector.iter().zip(&query).map(|(a, b)| a * b).sum::<f32>();
|
|
assert!(
|
|
cosine(&vectors[1]) > cosine(&vectors[0]),
|
|
"Real model did not retrieve the tail window"
|
|
);
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod coverage_tests;
|