feat: integrate Google semantic search and branded desktop launcher
Core CI / core (ubuntu-latest) (push) Failing after 3m27s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
Core CI / desktop (ubuntu-latest) (push) Successful in 5m45s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s
Core CI / core (ubuntu-latest) (push) Failing after 3m27s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
Core CI / desktop (ubuntu-latest) (push) Successful in 5m45s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s
This commit is contained in:
1 parent
a90df15503
commit
caf6a2f82e
80 files changed
+14971
-837
No files matched your search
@@ -0,0 +1,20 @@
|
||||
[package]
|
||||
name = "searchi-service"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
description = "Persistent local hybrid retrieval over Searchi lexical and CPU semantic engines"
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
serde.workspace = true
|
||||
searchi-core = { path = "../searchi-core" }
|
||||
searchi-semantic = { path = "../searchi-semantic" }
|
||||
rusqlite = { version = "0.32", features = ["bundled"] }
|
||||
sha2 = "0.10"
|
||||
fs2 = "0.4"
|
||||
usearch = { version = "=2.26.4", default-features = false, features = ["numkong"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
serde_json.workspace = true
|
||||
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"ann_p95_ms": 1.3722,
|
||||
"ann_recall_at10": 0.9850000000000001,
|
||||
"database_ms": 39209.1035,
|
||||
"dimensions": 768,
|
||||
"graph_build_ms": 57950.0959,
|
||||
"graph_memory_bytes": 554453184,
|
||||
"hardware": "skylake",
|
||||
"initial_eligibility_ms": 1612.2069999999999,
|
||||
"kind": "synthetic index-only benchmark, not semantic-model quality",
|
||||
"queries": 20,
|
||||
"threads": 1,
|
||||
"vectors": 100000,
|
||||
"warm_passages_rehashed": 0,
|
||||
"warm_service_p95_ms": 6.4033
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
//! Evaluate the immutable seed corpus with the actual configured Google model.
|
||||
//! Usage: cargo run -p searchi-service --example evaluate -- MODEL_DIR WORKER OUTPUT_JSON
|
||||
use anyhow::{ensure, Context, Result};
|
||||
use searchi_core::SearchIndex;
|
||||
use searchi_service::{HybridSearch, SearchOptions, SearchResult};
|
||||
use serde_json::{json, Value};
|
||||
use std::{fs, path::PathBuf, time::Instant};
|
||||
|
||||
fn expected_rank(results: &[SearchResult], filename: &str) -> Option<usize> {
|
||||
results
|
||||
.iter()
|
||||
.position(|result| result.filename == filename)
|
||||
.map(|rank| rank + 1)
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let args: Vec<_> = std::env::args_os().skip(1).collect();
|
||||
ensure!(
|
||||
args.len() == 3,
|
||||
"Supply MODEL_DIR WORKER OUTPUT_JSON; this command never downloads assets"
|
||||
);
|
||||
ensure!(
|
||||
std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1"),
|
||||
"Windows VM evaluation requires SEARCHI_FORCE_CPU=1"
|
||||
);
|
||||
let fixture: Value =
|
||||
serde_json::from_str(include_str!("../../../fixtures/relevance/fixture.json"))?;
|
||||
let temp = tempfile::tempdir()?;
|
||||
let root = temp.path().join("documents");
|
||||
fs::create_dir(&root)?;
|
||||
let documents = fixture["documents"]
|
||||
.as_array()
|
||||
.context("documents missing")?;
|
||||
for document in documents {
|
||||
let filename = document["filename"].as_str().context("filename missing")?;
|
||||
ensure!(
|
||||
PathBuf::from(filename).file_name() == Some(std::ffi::OsStr::new(filename)),
|
||||
"Fixture filename must be a basename"
|
||||
);
|
||||
fs::write(
|
||||
root.join(filename),
|
||||
document["text"].as_str().context("text missing")?,
|
||||
)?;
|
||||
}
|
||||
let db = temp.path().join("index.db");
|
||||
let mut lexical = SearchIndex::open(&db)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let start = Instant::now();
|
||||
let mut hybrid = HybridSearch::open(&db, &args[0], &args[1], 1)?;
|
||||
let load_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
let start = Instant::now();
|
||||
let indexed = hybrid.index_embeddings()?;
|
||||
let indexing_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
drop(hybrid);
|
||||
let start = Instant::now();
|
||||
let mut hybrid = HybridSearch::open(&db, &args[0], &args[1], 1)?;
|
||||
let startup_init_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
let start = Instant::now();
|
||||
hybrid.warmup()?;
|
||||
let background_prewarm_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
let start = Instant::now();
|
||||
let first = hybrid.search(
|
||||
"How can we end our relationship with the supplier?",
|
||||
20,
|
||||
&SearchOptions::default(),
|
||||
)?;
|
||||
let first_foreground_ms = start.elapsed().as_secs_f64() * 1000.0;
|
||||
let first_foreground_top1 = first.first().map(|hit| hit.filename.as_str());
|
||||
let mut rankings = Vec::new();
|
||||
let mut durations = Vec::new();
|
||||
for query in fixture["queries"].as_array().context("queries missing")? {
|
||||
let text = query["text"].as_str().context("query text missing")?;
|
||||
let expected = query["expected"].as_str().context("expected id missing")?;
|
||||
let filename = documents
|
||||
.iter()
|
||||
.find(|doc| doc["id"] == expected)
|
||||
.context("Expected document missing")?["filename"]
|
||||
.as_str()
|
||||
.unwrap();
|
||||
let lexical_hits = lexical.search(text, 20)?;
|
||||
let semantic_hits = hybrid.search_semantic(text, 20, &SearchOptions::default())?;
|
||||
let start = Instant::now();
|
||||
let hybrid_hits = hybrid.search(text, 20, &SearchOptions::default())?;
|
||||
durations.push(start.elapsed().as_secs_f64() * 1000.0);
|
||||
rankings.push(json!({"query":text,"expected":expected,"language":query["language"],
|
||||
"lexical_rank":expected_rank(&lexical_hits,filename),
|
||||
"semantic_rank":expected_rank(&semantic_hits,filename),
|
||||
"hybrid_rank":expected_rank(&hybrid_hits,filename),
|
||||
"hybrid_top3":hybrid_hits.iter().take(3).map(|hit| json!({"filename":hit.filename,"match_kind":hit.match_kind,"source":hit.source})).collect::<Vec<_>>()}));
|
||||
}
|
||||
durations.sort_by(f64::total_cmp);
|
||||
let metrics = |lane: &str| {
|
||||
json!({
|
||||
"top1_accuracy":rankings.iter().filter(|row| row[lane]==1).count() as f64/rankings.len() as f64,
|
||||
"recall_at3":rankings.iter().filter(|row| row[lane].as_u64().is_some_and(|rank| rank<=3)).count() as f64/rankings.len() as f64,
|
||||
"recall_at10":rankings.iter().filter(|row| row[lane].as_u64().is_some_and(|rank| rank<=10)).count() as f64/rankings.len() as f64,
|
||||
})
|
||||
};
|
||||
let report = json!({"model":"google/embeddinggemma-2","device":"cpu","threads":1,
|
||||
"documents":documents.len(),"queries":rankings.len(),"load_ms":load_ms,"indexing_ms":indexing_ms,
|
||||
"startup_init_ms":startup_init_ms,"background_prewarm_ms":background_prewarm_ms,"first_foreground_ms":first_foreground_ms,"first_foreground_top1":first_foreground_top1,
|
||||
"hybrid_p50_ms":durations[durations.len()/2],"hybrid_p95_ms":durations[(durations.len()*95/100).min(durations.len()-1)],
|
||||
"embedding_report":indexed,"lexical":metrics("lexical_rank"),"semantic":metrics("semantic_rank"),"hybrid":metrics("hybrid_rank"),
|
||||
"rankings":rankings,"limitation":"Synthetic seed corpus only; no production quality or large-corpus latency certification"});
|
||||
fs::write(&args[2], serde_json::to_string_pretty(&report)?)?;
|
||||
println!(
|
||||
"{}",
|
||||
json!({"lexical":report["lexical"],"semantic":report["semantic"],"hybrid":report["hybrid"],"p95_ms":report["hybrid_p95_ms"]})
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
//! Explicit index benchmark: synthetic normalized vectors, not neural retrieval
|
||||
//! certification. Includes the real generation store and warmed eligibility path.
|
||||
use super::*;
|
||||
use anyhow::{ensure, Result};
|
||||
use std::{fs, time::Instant};
|
||||
|
||||
fn random(state: &mut u64) -> f32 {
|
||||
*state ^= *state << 13;
|
||||
*state ^= *state >> 7;
|
||||
*state ^= *state << 17;
|
||||
((*state >> 32) as u32 as f32 / u32::MAX as f32) * 2.0 - 1.0
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Explicit CPU index benchmark; run cargo test --release -p searchi-service ann_100k -- --ignored --nocapture"]
|
||||
fn ann_100k_quality_and_warm_metadata_latency() -> Result<()> {
|
||||
ensure!(
|
||||
!cfg!(debug_assertions),
|
||||
"Run this performance benchmark with --release"
|
||||
);
|
||||
const COUNT: usize = 100_000;
|
||||
const DIM: usize = 768;
|
||||
const CLUSTERS: usize = 100;
|
||||
let mut state = 0x19cfa024b9817731_u64;
|
||||
let centers: Vec<Vec<f32>> = (0..CLUSTERS)
|
||||
.map(|_| (0..DIM).map(|_| random(&mut state)).collect())
|
||||
.collect();
|
||||
let mut points = Vec::<Vec<f32>>::with_capacity(COUNT);
|
||||
for i in 0..COUNT {
|
||||
let mut vector: Vec<f32> = centers[i % CLUSTERS]
|
||||
.iter()
|
||||
.map(|value| value + 0.35 * random(&mut state))
|
||||
.collect();
|
||||
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
|
||||
for value in &mut vector {
|
||||
*value /= norm;
|
||||
}
|
||||
points.push(vector);
|
||||
}
|
||||
let directory = tempfile::tempdir()?;
|
||||
let db = directory.path().join("benchmark.db");
|
||||
let lexical = searchi_core::SearchIndex::open(&db)?;
|
||||
drop(lexical);
|
||||
let mut store = Store::open(&db, DIM)?;
|
||||
let generation = store.begin("synthetic-ann-benchmark-not-model-embeddings")?;
|
||||
let started = Instant::now();
|
||||
{
|
||||
let transaction = store.connection.transaction()?;
|
||||
let mut file=transaction.prepare("INSERT INTO files(path,root,filename,content,size,mtime,extension) VALUES(?1,?2,?3,?4,100,0,?5)")?;
|
||||
let mut passage = transaction
|
||||
.prepare("INSERT INTO passages(text,file_path,kind,number) VALUES(?1,?2,'line',1)")?;
|
||||
let mut chunk=transaction.prepare("INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector)
|
||||
VALUES(?1,?2,0,?3,'line',1,?4,0,?5,?6,?7)")?;
|
||||
for (i, vector) in points.iter().enumerate() {
|
||||
let extension = if i % 10 == 0 { "md" } else { "txt" };
|
||||
let filename = format!("document-{i}.{extension}");
|
||||
let path = directory
|
||||
.path()
|
||||
.join(&filename)
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
let text =
|
||||
format!("Synthetic index benchmark record {i}. No inferred semantic labels.");
|
||||
file.execute(params![
|
||||
path,
|
||||
directory.path().to_string_lossy(),
|
||||
filename,
|
||||
text,
|
||||
extension
|
||||
])?;
|
||||
passage.execute(params![text, path])?;
|
||||
let passage_id = transaction.last_insert_rowid();
|
||||
let bytes: Vec<u8> = vector.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||
chunk.execute(params![
|
||||
generation,
|
||||
passage_id,
|
||||
path,
|
||||
content_hash(&text),
|
||||
text.len() as i64,
|
||||
text,
|
||||
bytes
|
||||
])?;
|
||||
}
|
||||
drop(file);
|
||||
drop(passage);
|
||||
drop(chunk);
|
||||
transaction.commit()?;
|
||||
}
|
||||
let database_ms = started.elapsed().as_secs_f64() * 1000.0;
|
||||
let started = Instant::now();
|
||||
store.activate(generation)?;
|
||||
let graph_build_ms = started.elapsed().as_secs_f64() * 1000.0;
|
||||
let ann = &store.ann_cache.as_ref().unwrap().1;
|
||||
let hardware = ann.hardware_acceleration();
|
||||
let graph_bytes = ann.memory_usage();
|
||||
let query_ids: Vec<_> = (0..20).map(|i| (i * 4999 + 37) % COUNT).collect();
|
||||
let mut recall = 0.0;
|
||||
let mut ann_times = Vec::new();
|
||||
for &query_id in &query_ids {
|
||||
let query = &points[query_id];
|
||||
let mut truth: Vec<_> = points
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, vector)| (i, vector.iter().zip(query).map(|(a, b)| a * b).sum::<f32>()))
|
||||
.collect();
|
||||
truth.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
|
||||
let expected: HashSet<_> = truth.iter().take(10).map(|(i, _)| *i as u64 + 1).collect();
|
||||
let started = Instant::now();
|
||||
let hits = ann.search(query, 10)?;
|
||||
ann_times.push(started.elapsed().as_secs_f64() * 1000.0);
|
||||
recall += hits
|
||||
.keys
|
||||
.iter()
|
||||
.filter(|key| expected.contains(key))
|
||||
.count() as f64
|
||||
/ 10.0;
|
||||
}
|
||||
let started = Instant::now();
|
||||
let _ = store.search(
|
||||
&points[query_ids[0]],
|
||||
"synthetic-ann-benchmark-not-model-embeddings",
|
||||
10,
|
||||
&SearchOptions::default(),
|
||||
)?;
|
||||
let initial_eligibility_ms = started.elapsed().as_secs_f64() * 1000.0;
|
||||
let warmed_hashes = HASH_CALLS.with(std::cell::Cell::get);
|
||||
let mut service_times = Vec::new();
|
||||
for &query_id in &query_ids {
|
||||
let started = Instant::now();
|
||||
let _ = store.search(
|
||||
&points[query_id],
|
||||
"synthetic-ann-benchmark-not-model-embeddings",
|
||||
10,
|
||||
&SearchOptions::default(),
|
||||
)?;
|
||||
service_times.push(started.elapsed().as_secs_f64() * 1000.0);
|
||||
}
|
||||
ensure!(
|
||||
HASH_CALLS.with(std::cell::Cell::get) == warmed_hashes,
|
||||
"Warm benchmark queries rehashed unchanged passages"
|
||||
);
|
||||
ann_times.sort_by(f64::total_cmp);
|
||||
service_times.sort_by(f64::total_cmp);
|
||||
let p95 = |values: &[f64]| values[(values.len() * 95 / 100).min(values.len() - 1)];
|
||||
let report = serde_json::json!({"kind":"synthetic index-only benchmark, not semantic-model quality",
|
||||
"vectors":COUNT,"dimensions":DIM,"queries":query_ids.len(),"threads":1,"hardware":hardware,
|
||||
"database_ms":database_ms,"graph_build_ms":graph_build_ms,"graph_memory_bytes":graph_bytes,
|
||||
"ann_recall_at10":recall/query_ids.len() as f64,"ann_p95_ms":p95(&ann_times),
|
||||
"initial_eligibility_ms":initial_eligibility_ms,"warm_service_p95_ms":p95(&service_times),"warm_passages_rehashed":0});
|
||||
println!("{}", serde_json::to_string_pretty(&report)?);
|
||||
if let Some(path) = std::env::var_os("SEARCHI_ANN_REPORT") {
|
||||
fs::write(path, serde_json::to_string_pretty(&report)?)?;
|
||||
}
|
||||
ensure!(
|
||||
report["ann_recall_at10"].as_f64().unwrap() >= 0.9,
|
||||
"ANN recall quality gate failed; preserve report and tune graph parameters"
|
||||
);
|
||||
ensure!(
|
||||
report["warm_service_p95_ms"].as_f64().unwrap() < 50.0,
|
||||
"Warm retrieval latency gate failed"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,531 @@
|
||||
//! Shared offline hybrid search using real Google embeddings and persisted,
|
||||
//! filtered HNSW retrieval. Production-scale validation remains a release gate.
|
||||
mod store;
|
||||
|
||||
use anyhow::{ensure, Result};
|
||||
use searchi_core::SearchIndex;
|
||||
pub use searchi_core::{SearchOptions, SearchResult};
|
||||
pub use searchi_semantic::WorkerInfo;
|
||||
use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
|
||||
use serde::Serialize;
|
||||
use std::{collections::HashMap, path::Path};
|
||||
use store::{Store, MAX_SCAN_CHUNKS};
|
||||
|
||||
const MAX_BATCH: usize = 8;
|
||||
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}")
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
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 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 worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
|
||||
Some(self.info.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct EmbeddingReport {
|
||||
pub passages: usize,
|
||||
pub chunks: usize,
|
||||
pub generation: i64,
|
||||
pub model_fingerprint: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct EmbeddingStatus {
|
||||
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,
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
impl HybridSearch {
|
||||
/// 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> {
|
||||
let connection =
|
||||
rusqlite::Connection::open_with_flags(db, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
|
||||
connection.busy_timeout(std::time::Duration::from_secs(5))?;
|
||||
let exists: bool = connection.query_row(
|
||||
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name='semantic_meta')",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
let pending: usize =
|
||||
connection.query_row("SELECT count(*) FROM passages", [], |row| row.get(0))?;
|
||||
if !exists {
|
||||
return Ok(EmbeddingStatus {
|
||||
generation: None,
|
||||
chunks: 0,
|
||||
model_fingerprint: None,
|
||||
stale: true,
|
||||
pending,
|
||||
});
|
||||
}
|
||||
let version: i64 = connection.query_row(
|
||||
"SELECT schema_version FROM semantic_meta WHERE singleton=1",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
if version < 5 {
|
||||
return Ok(EmbeddingStatus {
|
||||
generation: None,
|
||||
chunks: 0,
|
||||
model_fingerprint: None,
|
||||
stale: true,
|
||||
pending,
|
||||
});
|
||||
}
|
||||
ensure!(
|
||||
version == 5,
|
||||
"Unsupported semantic schema version {version}"
|
||||
);
|
||||
let (generation,fingerprint,revision,current):(Option<i64>,Option<String>,Option<i64>,i64)=connection.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 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)?)))?;
|
||||
let chunks: usize = connection.query_row(
|
||||
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
|
||||
[generation],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
let stale = revision != Some(current)
|
||||
|| fingerprint.as_deref() != Some(generation_fingerprint(GEMMA_FINGERPRINT).as_str());
|
||||
Ok(EmbeddingStatus {
|
||||
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))
|
||||
}
|
||||
|
||||
/// 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 generation_fingerprint = generation_fingerprint(embedder.fingerprint());
|
||||
Ok(Self {
|
||||
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 a replacement generation in bounded CPU batches. Inference or
|
||||
/// storage errors leave the previous generation active. A crashed build is
|
||||
/// reclaimed on the next explicit rebuild under an OS advisory lock.
|
||||
pub fn index_embeddings(&mut self) -> Result<EmbeddingReport> {
|
||||
let _lock = self.store.build_lock()?;
|
||||
let generation = self.store.begin(&self.generation_fingerprint)?;
|
||||
let built = self.build(generation);
|
||||
match built {
|
||||
Ok(report) => {
|
||||
self.store.activate(generation)?;
|
||||
Ok(report)
|
||||
}
|
||||
Err(error) => {
|
||||
self.store.abandon(generation)?;
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build(&mut self, generation: i64) -> Result<EmbeddingReport> {
|
||||
let mut report = EmbeddingReport {
|
||||
passages: 0,
|
||||
chunks: 0,
|
||||
generation,
|
||||
model_fingerprint: self.generation_fingerprint.clone(),
|
||||
};
|
||||
let mut after_id = 0;
|
||||
loop {
|
||||
let passages = self.lexical.passage_batch(after_id, 64)?;
|
||||
if passages.is_empty() {
|
||||
break;
|
||||
}
|
||||
for passage in passages {
|
||||
after_id = passage.id;
|
||||
if passage.text.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
let title = bounded_title(&passage.filename);
|
||||
let windows = passage_windows(&mut *self.embedder, &passage.text, title)?;
|
||||
ensure!(
|
||||
report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
|
||||
"Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
|
||||
);
|
||||
for (batch_index, batch) in windows.chunks(MAX_BATCH).enumerate() {
|
||||
let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
|
||||
let vectors = self.embedder.embed_documents(&texts, title)?;
|
||||
self.store.insert(
|
||||
generation,
|
||||
&passage,
|
||||
batch_index * MAX_BATCH,
|
||||
batch,
|
||||
&vectors,
|
||||
)?;
|
||||
}
|
||||
report.chunks += windows.len();
|
||||
report.passages += 1;
|
||||
}
|
||||
}
|
||||
Ok(report)
|
||||
}
|
||||
|
||||
/// 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 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(query, candidates, options)?;
|
||||
if query.trim().is_empty() || limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
promote_identifiers(query.trim(), &mut lexical);
|
||||
let vector = self.embedder.embed_query(query.trim())?;
|
||||
let semantic =
|
||||
self.store
|
||||
.search(&vector, &self.generation_fingerprint, candidates, options)?;
|
||||
Ok(fuse(lexical, semantic, limit.min(1000)))
|
||||
}
|
||||
|
||||
/// 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 vector = self.embedder.embed_query(query.trim())?;
|
||||
self.store.search(
|
||||
&vector,
|
||||
&self.generation_fingerprint,
|
||||
limit.min(1000),
|
||||
options,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
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::*;
|
||||
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(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,855 @@
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
fs::{File, OpenOptions},
|
||||
path::{Path, PathBuf},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use anyhow::{ensure, Context, Result};
|
||||
use fs2::FileExt;
|
||||
use rusqlite::{params, Connection, OptionalExtension};
|
||||
use searchi_core::{IndexedPassage, SearchOptions, SearchResult, SourceLocation};
|
||||
use searchi_semantic::TextWindow;
|
||||
use sha2::{Digest, Sha256};
|
||||
use usearch::{Index, IndexOptions, MetricKind, ScalarKind};
|
||||
|
||||
pub(crate) const MAX_SCAN_CHUNKS: usize = 100_000;
|
||||
pub(crate) const MIN_COSINE: f32 = 0.25;
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
struct FilterKey {
|
||||
folder: Option<String>,
|
||||
extension: Option<String>,
|
||||
min_size: Option<i64>,
|
||||
max_size: Option<i64>,
|
||||
modified_after: Option<i64>,
|
||||
modified_before: Option<i64>,
|
||||
}
|
||||
|
||||
struct EligibilityCache {
|
||||
generation: i64,
|
||||
revision: i64,
|
||||
filter: FilterKey,
|
||||
keys: Arc<HashSet<u64>>,
|
||||
by_path: HashMap<String, Vec<u64>>,
|
||||
}
|
||||
|
||||
pub(crate) struct Store {
|
||||
connection: Connection,
|
||||
lock_path: PathBuf,
|
||||
dimensions: usize,
|
||||
ann_cache: Option<(i64, Index)>,
|
||||
eligibility_cache: Option<EligibilityCache>,
|
||||
}
|
||||
|
||||
pub(crate) struct BuildLock(File);
|
||||
impl Drop for BuildLock {
|
||||
fn drop(&mut self) {
|
||||
let _ = FileExt::unlock(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn content_hash(text: &str) -> String {
|
||||
#[cfg(test)]
|
||||
HASH_CALLS.with(|count| count.set(count.get() + 1));
|
||||
format!("{:x}", Sha256::digest(text.as_bytes()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
thread_local! { static HASH_CALLS:std::cell::Cell<usize>=const {std::cell::Cell::new(0)}; }
|
||||
|
||||
impl Store {
|
||||
pub(crate) fn open(db: &Path, dimensions: usize) -> Result<Self> {
|
||||
ensure!(
|
||||
(1..=4096).contains(&dimensions),
|
||||
"Embedding dimensions must be between 1 and 4096"
|
||||
);
|
||||
ensure!(
|
||||
db != Path::new(":memory:"),
|
||||
"Hybrid search requires a persistent database path"
|
||||
);
|
||||
let mut connection = Connection::open(db)?;
|
||||
connection.busy_timeout(std::time::Duration::from_secs(5))?;
|
||||
connection.execute_batch("PRAGMA foreign_keys=ON;")?;
|
||||
let initialization =
|
||||
connection.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
|
||||
initialization.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS semantic_meta (
|
||||
singleton INTEGER PRIMARY KEY CHECK(singleton=1), schema_version INTEGER NOT NULL,
|
||||
active_generation INTEGER, corpus_revision INTEGER NOT NULL DEFAULT 0, dirty_floor INTEGER NOT NULL DEFAULT 0);
|
||||
INSERT OR IGNORE INTO semantic_meta(singleton,schema_version) VALUES(1,5);",
|
||||
)?;
|
||||
let version: i64 = initialization.query_row(
|
||||
"SELECT schema_version FROM semantic_meta WHERE singleton=1",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
ensure!(
|
||||
(1..=5).contains(&version),
|
||||
"Unsupported semantic schema version {version}"
|
||||
);
|
||||
initialization.execute_batch("
|
||||
CREATE TABLE IF NOT EXISTS semantic_generations(id INTEGER PRIMARY KEY, fingerprint TEXT NOT NULL, dimensions INTEGER NOT NULL, corpus_revision INTEGER NOT NULL DEFAULT 0, ann_blob BLOB, ann_sha256 TEXT);
|
||||
CREATE TABLE IF NOT EXISTS semantic_chunks(
|
||||
generation INTEGER NOT NULL REFERENCES semantic_generations(id) ON DELETE CASCADE,
|
||||
passage_id INTEGER NOT NULL, ordinal INTEGER NOT NULL,
|
||||
path TEXT NOT NULL, kind TEXT NOT NULL, number INTEGER NOT NULL,
|
||||
content_hash TEXT NOT NULL, start_byte INTEGER NOT NULL, end_byte INTEGER NOT NULL,
|
||||
snippet TEXT NOT NULL, vector BLOB NOT NULL,
|
||||
PRIMARY KEY(generation,passage_id,ordinal));")?;
|
||||
if version < 5 {
|
||||
for (table,column,alter) in [
|
||||
("semantic_meta","corpus_revision","ALTER TABLE semantic_meta ADD COLUMN corpus_revision INTEGER NOT NULL DEFAULT 0"),
|
||||
("semantic_meta","dirty_floor","ALTER TABLE semantic_meta ADD COLUMN dirty_floor INTEGER NOT NULL DEFAULT 0"),
|
||||
("semantic_generations","dimensions","ALTER TABLE semantic_generations ADD COLUMN dimensions INTEGER NOT NULL DEFAULT 384"),
|
||||
("semantic_generations","corpus_revision","ALTER TABLE semantic_generations ADD COLUMN corpus_revision INTEGER NOT NULL DEFAULT 0"),
|
||||
("semantic_generations","ann_blob","ALTER TABLE semantic_generations ADD COLUMN ann_blob BLOB"),
|
||||
("semantic_generations","ann_sha256","ALTER TABLE semantic_generations ADD COLUMN ann_sha256 TEXT"),
|
||||
] {
|
||||
let exists:bool=initialization.query_row("SELECT EXISTS(SELECT 1 FROM pragma_table_info(?1) WHERE name=?2)",params![table,column],|row| row.get(0))?;
|
||||
if !exists { initialization.execute_batch(alter)?; }
|
||||
}
|
||||
initialization.execute_batch("DROP TRIGGER IF EXISTS semantic_files_insert;
|
||||
DROP TRIGGER IF EXISTS semantic_files_update; DROP TRIGGER IF EXISTS semantic_files_delete;")?;
|
||||
}
|
||||
initialization.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS semantic_dirty_paths(path TEXT PRIMARY KEY, revision INTEGER NOT NULL);
|
||||
CREATE INDEX IF NOT EXISTS semantic_dirty_revision ON semantic_dirty_paths(revision);
|
||||
CREATE INDEX IF NOT EXISTS semantic_chunk_path ON semantic_chunks(generation,path);
|
||||
CREATE TRIGGER IF NOT EXISTS semantic_files_insert AFTER INSERT ON files BEGIN
|
||||
UPDATE semantic_meta SET corpus_revision=corpus_revision+1 WHERE singleton=1;
|
||||
INSERT INTO semantic_dirty_paths(path,revision) VALUES(new.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
|
||||
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision; END;
|
||||
CREATE TRIGGER IF NOT EXISTS semantic_files_update AFTER UPDATE ON files BEGIN
|
||||
UPDATE semantic_meta SET corpus_revision=corpus_revision+1 WHERE singleton=1;
|
||||
INSERT INTO semantic_dirty_paths(path,revision) VALUES(old.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
|
||||
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision;
|
||||
INSERT INTO semantic_dirty_paths(path,revision) VALUES(new.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
|
||||
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision; END;
|
||||
CREATE TRIGGER IF NOT EXISTS semantic_files_delete AFTER DELETE ON files BEGIN
|
||||
UPDATE semantic_meta SET corpus_revision=corpus_revision+1 WHERE singleton=1;
|
||||
INSERT INTO semantic_dirty_paths(path,revision) VALUES(old.path,(SELECT corpus_revision FROM semantic_meta WHERE singleton=1))
|
||||
ON CONFLICT(path) DO UPDATE SET revision=excluded.revision; END;
|
||||
UPDATE semantic_meta SET schema_version=5 WHERE singleton=1;",
|
||||
)?;
|
||||
initialization.commit()?;
|
||||
let canonical = db.canonicalize()?;
|
||||
let lock_path = PathBuf::from(format!("{}.semantic.lock", canonical.to_string_lossy()));
|
||||
Ok(Self {
|
||||
connection,
|
||||
lock_path,
|
||||
dimensions,
|
||||
ann_cache: None,
|
||||
eligibility_cache: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn build_lock(&self) -> Result<BuildLock> {
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(false)
|
||||
.open(&self.lock_path)?;
|
||||
file.try_lock_exclusive()
|
||||
.context("Another semantic rebuild is running")?;
|
||||
Ok(BuildLock(file))
|
||||
}
|
||||
|
||||
pub(crate) fn begin(&mut self, fingerprint: &str) -> Result<i64> {
|
||||
let transaction = self.connection.transaction()?;
|
||||
// Only called while holding the cross-process lock: abandoned staging
|
||||
// generations may be removed safely after a crash.
|
||||
transaction.execute("DELETE FROM semantic_generations WHERE id != COALESCE((SELECT active_generation FROM semantic_meta WHERE singleton=1),-1)", [])?;
|
||||
transaction.execute(
|
||||
"INSERT INTO semantic_generations(fingerprint,dimensions,corpus_revision) SELECT ?1,?2,corpus_revision FROM semantic_meta WHERE singleton=1",
|
||||
params![fingerprint, self.dimensions as i64],
|
||||
)?;
|
||||
let id = transaction.last_insert_rowid();
|
||||
transaction.commit()?;
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub(crate) fn insert(
|
||||
&mut self,
|
||||
generation: i64,
|
||||
passage: &IndexedPassage,
|
||||
first_ordinal: usize,
|
||||
windows: &[TextWindow],
|
||||
vectors: &[Vec<f32>],
|
||||
) -> Result<()> {
|
||||
ensure!(
|
||||
windows.len() == vectors.len(),
|
||||
"Window and embedding counts differ"
|
||||
);
|
||||
let hash = content_hash(&passage.text);
|
||||
let transaction = self.connection.transaction()?;
|
||||
for (i, (window, vector)) in windows.iter().zip(vectors).enumerate() {
|
||||
ensure!(
|
||||
vector.len() == self.dimensions && vector.iter().all(|v| v.is_finite()),
|
||||
"Invalid persisted vector"
|
||||
);
|
||||
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
|
||||
ensure!(
|
||||
norm.is_finite() && (norm - 1.0).abs() < 1e-3,
|
||||
"Persisted embeddings must be L2 normalized"
|
||||
);
|
||||
ensure!(
|
||||
passage.text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
|
||||
"Window does not match passage source range"
|
||||
);
|
||||
let bytes: Vec<u8> = vector
|
||||
.iter()
|
||||
.flat_map(|value| value.to_le_bytes())
|
||||
.collect();
|
||||
transaction.execute("INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector)
|
||||
VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)", params![generation, passage.id,
|
||||
(first_ordinal+i) as i64, passage.path, passage.source.kind, passage.source.number,
|
||||
hash, window.start_byte as i64, window.end_byte as i64, window.text, bytes])?;
|
||||
}
|
||||
transaction.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn activate(&mut self, generation: i64) -> Result<()> {
|
||||
// A failed graph build or serialization leaves the old generation active.
|
||||
let ann = self.build_ann(generation)?;
|
||||
let length = ann.serialized_length();
|
||||
ensure!(
|
||||
length <= 512 * 1024 * 1024,
|
||||
"Serialized ANN exceeds 512 MiB generation budget"
|
||||
);
|
||||
let mut serialized = vec![0_u8; length];
|
||||
ann.save_to_buffer(&mut serialized)?;
|
||||
let checksum = format!("{:x}", Sha256::digest(&serialized));
|
||||
let transaction = self.connection.transaction()?;
|
||||
ensure!(
|
||||
transaction.query_row(
|
||||
"SELECT EXISTS(SELECT 1 FROM semantic_generations WHERE id=?1)",
|
||||
[generation],
|
||||
|row| row.get::<_, bool>(0)
|
||||
)?,
|
||||
"Missing staging generation"
|
||||
);
|
||||
transaction.execute(
|
||||
"UPDATE semantic_generations SET ann_blob=?1,ann_sha256=?2 WHERE id=?3",
|
||||
params![serialized, checksum, generation],
|
||||
)?;
|
||||
transaction.execute(
|
||||
"UPDATE semantic_meta SET active_generation=?1 WHERE singleton=1",
|
||||
[generation],
|
||||
)?;
|
||||
transaction.execute(
|
||||
"DELETE FROM semantic_generations WHERE id!=?1",
|
||||
[generation],
|
||||
)?;
|
||||
transaction.execute_batch("UPDATE semantic_meta SET dirty_floor=max(dirty_floor,corpus_revision-100000) WHERE singleton=1;
|
||||
DELETE FROM semantic_dirty_paths WHERE revision<(SELECT dirty_floor FROM semantic_meta WHERE singleton=1);")?;
|
||||
transaction.commit()?;
|
||||
self.ann_cache = Some((generation, ann));
|
||||
self.eligibility_cache = None;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn new_ann(&self) -> Result<Index> {
|
||||
Ok(Index::new(&IndexOptions {
|
||||
dimensions: self.dimensions,
|
||||
metric: MetricKind::Cos,
|
||||
quantization: ScalarKind::F32,
|
||||
connectivity: 16,
|
||||
expansion_add: 128,
|
||||
expansion_search: 64,
|
||||
multi: false,
|
||||
})?)
|
||||
}
|
||||
|
||||
fn build_ann(&self, generation: i64) -> Result<Index> {
|
||||
let count: usize = self.connection.query_row(
|
||||
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
|
||||
[generation],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
ensure!(
|
||||
count <= MAX_SCAN_CHUNKS,
|
||||
"ANN generation exceeds {MAX_SCAN_CHUNKS} chunks"
|
||||
);
|
||||
let ann = self.new_ann()?;
|
||||
ann.reserve_capacity_and_threads(count.max(1), 1)?;
|
||||
let mut statement = self.connection.prepare(
|
||||
"SELECT rowid,vector FROM semantic_chunks WHERE generation=?1 ORDER BY rowid",
|
||||
)?;
|
||||
let mut rows = statement.query([generation])?;
|
||||
while let Some(row) = rows.next()? {
|
||||
let key: i64 = row.get(0)?;
|
||||
let bytes: Vec<u8> = row.get(1)?;
|
||||
ensure!(
|
||||
bytes.len() == self.dimensions * 4,
|
||||
"Invalid staging embedding length"
|
||||
);
|
||||
let vector: Vec<f32> = bytes
|
||||
.as_chunks::<4>()
|
||||
.0
|
||||
.iter()
|
||||
.map(|bytes| f32::from_le_bytes(*bytes))
|
||||
.collect();
|
||||
ensure!(
|
||||
vector.iter().all(|v| v.is_finite()),
|
||||
"Invalid staging embedding values"
|
||||
);
|
||||
ann.add(u64::try_from(key)?, &vector)?;
|
||||
}
|
||||
Ok(ann)
|
||||
}
|
||||
|
||||
pub(crate) fn abandon(&mut self, generation: i64) -> Result<()> {
|
||||
self.connection.execute("DELETE FROM semantic_generations WHERE id=?1 AND id!=COALESCE((SELECT active_generation FROM semantic_meta WHERE singleton=1),-1)", [generation])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn search(
|
||||
&mut self,
|
||||
query: &[f32],
|
||||
fingerprint: &str,
|
||||
limit: usize,
|
||||
options: &SearchOptions,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
ensure!(
|
||||
query.len() == self.dimensions && query.iter().all(|v| v.is_finite()),
|
||||
"Invalid query vector"
|
||||
);
|
||||
if limit == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
// A consistent read snapshot covers eligibility and materialization.
|
||||
let snapshot = self.connection.unchecked_transaction()?;
|
||||
let active:Option<i64>=snapshot.query_row("SELECT g.id FROM semantic_generations g JOIN semantic_meta m ON m.active_generation=g.id WHERE g.fingerprint=?1 AND g.dimensions=?2",
|
||||
params![fingerprint,self.dimensions as i64],|row|row.get(0)).optional()?;
|
||||
let Some(active) = active else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
if !self
|
||||
.ann_cache
|
||||
.as_ref()
|
||||
.is_some_and(|(generation, _)| *generation == active)
|
||||
{
|
||||
let length: Option<usize> = snapshot.query_row(
|
||||
"SELECT length(ann_blob) FROM semantic_generations WHERE id=?1",
|
||||
[active],
|
||||
|row| row.get(0),
|
||||
)?;
|
||||
ensure!(
|
||||
length.is_some_and(|bytes| bytes <= 512 * 1024 * 1024),
|
||||
"Missing or oversized ANN generation; rebuild embeddings explicitly"
|
||||
);
|
||||
let (serialized, checksum): (Vec<u8>, String) = snapshot.query_row(
|
||||
"SELECT ann_blob,ann_sha256 FROM semantic_generations WHERE id=?1",
|
||||
[active],
|
||||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||||
)?;
|
||||
ensure!(
|
||||
format!("{:x}", Sha256::digest(&serialized)) == checksum,
|
||||
"ANN generation SHA256 mismatch; rebuild embeddings explicitly"
|
||||
);
|
||||
let ann = self.new_ann()?;
|
||||
ann.load_from_buffer(&serialized)?;
|
||||
ensure!(
|
||||
ann.dimensions() == self.dimensions,
|
||||
"ANN dimensions differ from embedding provider"
|
||||
);
|
||||
ann.reserve_capacity_and_threads(ann.size().max(1), 1)?;
|
||||
self.ann_cache = Some((active, ann));
|
||||
}
|
||||
let folder = options
|
||||
.folder
|
||||
.as_ref()
|
||||
.map(std::fs::canonicalize)
|
||||
.transpose()?
|
||||
.map(|p| {
|
||||
format!(
|
||||
"{}{}",
|
||||
p.to_string_lossy().trim_end_matches(['/', '\\']),
|
||||
std::path::MAIN_SEPARATOR
|
||||
)
|
||||
});
|
||||
let extension = options
|
||||
.extension
|
||||
.as_ref()
|
||||
.map(|ext| ext.trim_start_matches('.').to_ascii_lowercase());
|
||||
let min_size = options.min_size.map(i64::try_from).transpose()?;
|
||||
let max_size = options.max_size.map(i64::try_from).transpose()?;
|
||||
let filter = FilterKey {
|
||||
folder,
|
||||
extension,
|
||||
min_size,
|
||||
max_size,
|
||||
modified_after: options.modified_after_unix_ms,
|
||||
modified_before: options.modified_before_unix_ms,
|
||||
};
|
||||
let (revision, floor): (i64, i64) = snapshot.query_row(
|
||||
"SELECT corpus_revision,dirty_floor FROM semantic_meta WHERE singleton=1",
|
||||
[],
|
||||
|row| Ok((row.get(0)?, row.get(1)?)),
|
||||
)?;
|
||||
let existing = self.eligibility_cache.take();
|
||||
let mut cache = match existing {
|
||||
Some(mut cache)
|
||||
if cache.generation == active
|
||||
&& cache.filter == filter
|
||||
&& cache.revision >= floor =>
|
||||
{
|
||||
if cache.revision != revision {
|
||||
let mut statement=snapshot.prepare("SELECT path FROM semantic_dirty_paths WHERE revision>?1 AND revision<=?2 ORDER BY path LIMIT 1001")?;
|
||||
let paths = statement
|
||||
.query_map(params![cache.revision, revision], |row| {
|
||||
row.get::<_, String>(0)
|
||||
})?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
if paths.len() > 1000 {
|
||||
cache = collect_eligibility(&snapshot, active, revision, &filter, None)?;
|
||||
} else {
|
||||
for path in paths {
|
||||
if let Some(old) = cache.by_path.remove(&path) {
|
||||
for key in old {
|
||||
Arc::make_mut(&mut cache.keys).remove(&key);
|
||||
}
|
||||
}
|
||||
let changed = collect_eligibility(
|
||||
&snapshot,
|
||||
active,
|
||||
revision,
|
||||
&filter,
|
||||
Some(&path),
|
||||
)?;
|
||||
Arc::make_mut(&mut cache.keys).extend(changed.keys.iter().copied());
|
||||
cache.by_path.extend(changed.by_path);
|
||||
}
|
||||
cache.revision = revision;
|
||||
}
|
||||
}
|
||||
cache
|
||||
}
|
||||
_ => collect_eligibility(&snapshot, active, revision, &filter, None)?,
|
||||
};
|
||||
ensure!(
|
||||
cache.keys.len() <= MAX_SCAN_CHUNKS,
|
||||
"Eligibility cache exceeds chunk resource budget"
|
||||
);
|
||||
cache.revision = revision;
|
||||
let eligible = Arc::clone(&cache.keys);
|
||||
self.eligibility_cache = Some(cache);
|
||||
if eligible.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let ann = &self.ann_cache.as_ref().expect("loaded ANN cache").1;
|
||||
let mut candidate_count = limit.saturating_mul(4).min(eligible.len()).max(1);
|
||||
let mut results;
|
||||
loop {
|
||||
let matches =
|
||||
ann.filtered_search(query, candidate_count, |key| eligible.contains(&key))?;
|
||||
let mut ranked: Vec<_> = matches.keys.into_iter().zip(matches.distances).collect();
|
||||
ranked.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
|
||||
let mut best = HashMap::<String, SearchResult>::new();
|
||||
let mut materialize=snapshot.prepare("SELECT c.path,f.filename,c.snippet,c.kind,c.number FROM semantic_chunks c JOIN files f ON f.path=c.path WHERE c.rowid=?1 AND c.generation=?2")?;
|
||||
for (key, distance) in ranked {
|
||||
let cosine = 1.0 - distance;
|
||||
ensure!(cosine.is_finite(), "ANN produced nonfinite distance");
|
||||
if cosine < MIN_COSINE {
|
||||
continue;
|
||||
}
|
||||
let result =
|
||||
materialize.query_row(params![i64::try_from(key)?, active], |row| {
|
||||
Ok(SearchResult {
|
||||
path: row.get(0)?,
|
||||
filename: row.get(1)?,
|
||||
snippet: row.get(2)?,
|
||||
score: f64::from(cosine),
|
||||
match_kind: "semantic".into(),
|
||||
source: Some(SourceLocation {
|
||||
kind: row.get(3)?,
|
||||
number: row.get(4)?,
|
||||
}),
|
||||
extraction_error: None,
|
||||
})
|
||||
})?;
|
||||
if best.contains_key(&result.path) {
|
||||
continue;
|
||||
}
|
||||
best.insert(result.path.clone(), result);
|
||||
}
|
||||
results = best.into_values().collect::<Vec<_>>();
|
||||
if results.len() >= limit || candidate_count == eligible.len() {
|
||||
break;
|
||||
}
|
||||
candidate_count = candidate_count.saturating_mul(2).min(eligible.len());
|
||||
}
|
||||
results.sort_by(|a, b| {
|
||||
b.score
|
||||
.total_cmp(&a.score)
|
||||
.then_with(|| a.path.cmp(&b.path))
|
||||
});
|
||||
results.truncate(limit);
|
||||
snapshot.commit()?;
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn exact_search(
|
||||
&self,
|
||||
query: &[f32],
|
||||
fingerprint: &str,
|
||||
limit: usize,
|
||||
options: &SearchOptions,
|
||||
) -> Result<Vec<SearchResult>> {
|
||||
ensure!(
|
||||
query.len() == self.dimensions && query.iter().all(|v| v.is_finite()),
|
||||
"Invalid query vector"
|
||||
);
|
||||
let active: Option<i64> = self.connection.query_row("SELECT g.id FROM semantic_generations g JOIN semantic_meta m ON m.active_generation=g.id WHERE g.fingerprint=?1 AND g.dimensions=?2", params![fingerprint,self.dimensions as i64], |row| row.get(0)).optional()?;
|
||||
let Some(active) = active else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let folder = options
|
||||
.folder
|
||||
.as_ref()
|
||||
.map(std::fs::canonicalize)
|
||||
.transpose()?
|
||||
.map(|p| {
|
||||
format!(
|
||||
"{}{}",
|
||||
p.to_string_lossy().trim_end_matches(['/', '\\']),
|
||||
std::path::MAIN_SEPARATOR
|
||||
)
|
||||
});
|
||||
let extension = options
|
||||
.extension
|
||||
.as_ref()
|
||||
.map(|ext| ext.trim_start_matches('.').to_ascii_lowercase());
|
||||
let min_size = options.min_size.map(i64::try_from).transpose()?;
|
||||
let max_size = options.max_size.map(i64::try_from).transpose()?;
|
||||
let mut statement = self.connection.prepare("SELECT c.passage_id,c.path,f.filename,c.snippet,c.vector,c.content_hash,p.text,c.kind,c.number
|
||||
FROM semantic_chunks c JOIN passages p ON p.rowid=c.passage_id AND p.file_path=c.path AND p.kind=c.kind AND p.number=c.number
|
||||
JOIN files f ON f.path=c.path WHERE c.generation=?1
|
||||
AND (?2 IS NULL OR substr(f.path,1,length(?2))=?2)
|
||||
AND (?3 IS NULL OR f.extension=?3)
|
||||
AND (?4 IS NULL OR f.size>=?4) AND (?5 IS NULL OR f.size<=?5)
|
||||
AND (?6 IS NULL OR f.mtime/1000000>=?6) AND (?7 IS NULL OR f.mtime/1000000<=?7)
|
||||
ORDER BY c.passage_id,c.ordinal LIMIT ?8")?;
|
||||
// Fetch one extra to detect the baseline cap instead of silently hiding
|
||||
// candidates. Metadata filtering precedes this cap and all top-k limits.
|
||||
let mut rows = statement.query(params![
|
||||
active,
|
||||
folder,
|
||||
extension,
|
||||
min_size,
|
||||
max_size,
|
||||
options.modified_after_unix_ms,
|
||||
options.modified_before_unix_ms,
|
||||
MAX_SCAN_CHUNKS as i64 + 1
|
||||
])?;
|
||||
let mut hashes = HashMap::<i64, String>::new();
|
||||
let mut best = HashMap::<String, SearchResult>::new();
|
||||
let mut scanned = 0;
|
||||
while let Some(row) = rows.next()? {
|
||||
scanned += 1;
|
||||
ensure!(
|
||||
scanned <= MAX_SCAN_CHUNKS,
|
||||
"Exact-vector baseline exceeds {MAX_SCAN_CHUNKS} chunks; ANN integration required"
|
||||
);
|
||||
let passage_id: i64 = row.get(0)?;
|
||||
let current_hash = if let Some(hash) = hashes.get(&passage_id) {
|
||||
hash.clone()
|
||||
} else {
|
||||
let hash = content_hash(&row.get::<_, String>(6)?);
|
||||
hashes.insert(passage_id, hash.clone());
|
||||
hash
|
||||
};
|
||||
if current_hash != row.get::<_, String>(5)? {
|
||||
continue;
|
||||
}
|
||||
let bytes: Vec<u8> = row.get(4)?;
|
||||
ensure!(
|
||||
bytes.len() == self.dimensions * 4,
|
||||
"Corrupt stored embedding length"
|
||||
);
|
||||
let mut cosine = 0.0;
|
||||
for (dimension, value) in bytes.as_chunks::<4>().0.iter().enumerate() {
|
||||
let value = f32::from_le_bytes(*value);
|
||||
ensure!(value.is_finite(), "Corrupt stored embedding value");
|
||||
cosine += value * query[dimension];
|
||||
}
|
||||
if cosine < MIN_COSINE {
|
||||
continue;
|
||||
}
|
||||
let path: String = row.get(1)?;
|
||||
if best
|
||||
.get(&path)
|
||||
.is_some_and(|result| result.score >= f64::from(cosine))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
best.insert(
|
||||
path.clone(),
|
||||
SearchResult {
|
||||
path,
|
||||
filename: row.get(2)?,
|
||||
snippet: row.get(3)?,
|
||||
score: f64::from(cosine),
|
||||
match_kind: "semantic".into(),
|
||||
source: Some(SourceLocation {
|
||||
kind: row.get(7)?,
|
||||
number: row.get(8)?,
|
||||
}),
|
||||
extraction_error: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
let mut results: Vec<_> = best.into_values().collect();
|
||||
results.sort_by(|a, b| {
|
||||
b.score
|
||||
.total_cmp(&a.score)
|
||||
.then_with(|| a.path.cmp(&b.path))
|
||||
});
|
||||
results.truncate(limit);
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_eligibility(
|
||||
connection: &Connection,
|
||||
generation: i64,
|
||||
revision: i64,
|
||||
filter: &FilterKey,
|
||||
path: Option<&str>,
|
||||
) -> Result<EligibilityCache> {
|
||||
let mut statement=connection.prepare("SELECT c.rowid,c.passage_id,c.content_hash,p.text,c.path FROM semantic_chunks c
|
||||
JOIN passages p ON p.rowid=c.passage_id AND p.file_path=c.path AND p.kind=c.kind AND p.number=c.number
|
||||
JOIN files f ON f.path=c.path WHERE c.generation=?1
|
||||
AND (?2 IS NULL OR substr(f.path,1,length(?2))=?2) AND (?3 IS NULL OR f.extension=?3)
|
||||
AND (?4 IS NULL OR f.size>=?4) AND (?5 IS NULL OR f.size<=?5)
|
||||
AND (?6 IS NULL OR f.mtime/1000000>=?6) AND (?7 IS NULL OR f.mtime/1000000<=?7)
|
||||
AND (?9 IS NULL OR c.path=?9) ORDER BY c.passage_id,c.ordinal LIMIT ?8")?;
|
||||
let mut rows = statement.query(params![
|
||||
generation,
|
||||
filter.folder,
|
||||
filter.extension,
|
||||
filter.min_size,
|
||||
filter.max_size,
|
||||
filter.modified_after,
|
||||
filter.modified_before,
|
||||
MAX_SCAN_CHUNKS as i64 + 1,
|
||||
path
|
||||
])?;
|
||||
let mut hashes = HashMap::<i64, String>::new();
|
||||
let mut keys = HashSet::new();
|
||||
let mut by_path = HashMap::<String, Vec<u64>>::new();
|
||||
let mut count = 0;
|
||||
while let Some(row) = rows.next()? {
|
||||
count += 1;
|
||||
ensure!(
|
||||
count <= MAX_SCAN_CHUNKS,
|
||||
"Corpus exceeds {MAX_SCAN_CHUNKS} chunk resource budget"
|
||||
);
|
||||
let passage_id: i64 = row.get(1)?;
|
||||
let hash = if let Some(hash) = hashes.get(&passage_id) {
|
||||
hash.clone()
|
||||
} else {
|
||||
let hash = content_hash(&row.get::<_, String>(3)?);
|
||||
hashes.insert(passage_id, hash.clone());
|
||||
hash
|
||||
};
|
||||
if hash == row.get::<_, String>(2)? {
|
||||
let key = u64::try_from(row.get::<_, i64>(0)?)?;
|
||||
keys.insert(key);
|
||||
by_path.entry(row.get(4)?).or_default().push(key);
|
||||
}
|
||||
}
|
||||
Ok(EligibilityCache {
|
||||
generation,
|
||||
revision,
|
||||
filter: filter.clone(),
|
||||
keys: Arc::new(keys),
|
||||
by_path,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use searchi_core::SearchIndex;
|
||||
const DIMENSIONS: usize = 768;
|
||||
|
||||
#[test]
|
||||
fn generation_persists_failed_staging_preserves_active_and_changes_hide_stale_vectors(
|
||||
) -> Result<()> {
|
||||
let directory = tempfile::tempdir()?;
|
||||
let root = directory.path().join("docs");
|
||||
std::fs::create_dir(&root)?;
|
||||
let db = directory.path().join("index.db");
|
||||
std::fs::write(root.join("notes.txt"), "original relevant content")?;
|
||||
let mut index = SearchIndex::open(&db)?;
|
||||
index.index_folder(&root)?;
|
||||
let passage = index.passage_batch(0, 1)?.remove(0);
|
||||
let mut vector = vec![0.0; DIMENSIONS];
|
||||
vector[0] = 1.0;
|
||||
let windows = [TextWindow {
|
||||
text: passage.text.clone(),
|
||||
start_byte: 0,
|
||||
end_byte: passage.text.len(),
|
||||
}];
|
||||
let mut store = Store::open(&db, DIMENSIONS)?;
|
||||
let generation = store.begin("test-storage-vectors")?;
|
||||
store.insert(generation, &passage, 0, &windows, &[vector.clone()])?;
|
||||
store.activate(generation)?;
|
||||
assert_eq!(
|
||||
store.exact_search(
|
||||
&vector,
|
||||
"test-storage-vectors",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?[0]
|
||||
.path,
|
||||
store.search(
|
||||
&vector,
|
||||
"test-storage-vectors",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?[0]
|
||||
.path
|
||||
);
|
||||
drop(store);
|
||||
let mut store = Store::open(&db, DIMENSIONS)?;
|
||||
assert_eq!(
|
||||
store
|
||||
.search(
|
||||
&vector,
|
||||
"test-storage-vectors",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert!(store
|
||||
.search(&vector, "different-model", 10, &SearchOptions::default())?
|
||||
.is_empty());
|
||||
let staging = store.begin("new-model")?;
|
||||
assert_eq!(
|
||||
store
|
||||
.search(
|
||||
&vector,
|
||||
"test-storage-vectors",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
store.abandon(staging)?;
|
||||
std::fs::write(
|
||||
root.join("notes.txt"),
|
||||
"replacement unrelated content longer than before",
|
||||
)?;
|
||||
index.index_folder(&root)?;
|
||||
assert!(store
|
||||
.search(
|
||||
&vector,
|
||||
"test-storage-vectors",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?
|
||||
.is_empty());
|
||||
std::fs::remove_file(root.join("notes.txt"))?;
|
||||
index.index_folder(&root)?;
|
||||
assert!(store
|
||||
.search(
|
||||
&vector,
|
||||
"test-storage-vectors",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?
|
||||
.is_empty());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ann_filters_before_top_k_checks_checksum_and_status_tracks_corpus_changes() -> Result<()> {
|
||||
let directory = tempfile::tempdir()?;
|
||||
let root = directory.path().join("docs");
|
||||
std::fs::create_dir(&root)?;
|
||||
let db = directory.path().join("index.db");
|
||||
std::fs::write(root.join("closest.md"), "closest metadata excluded source")?;
|
||||
std::fs::write(root.join("allowed.txt"), "allowed metadata source")?;
|
||||
let mut lexical = SearchIndex::open(&db)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let passages = lexical.passage_batch(0, 10)?;
|
||||
let mut store = Store::open(&db, DIMENSIONS)?;
|
||||
let fingerprint = crate::generation_fingerprint(searchi_semantic::GEMMA_FINGERPRINT);
|
||||
let generation = store.begin(&fingerprint)?;
|
||||
let mut query = vec![0.0; DIMENSIONS];
|
||||
query[0] = 1.0;
|
||||
for passage in &passages {
|
||||
let mut vector = query.clone();
|
||||
if passage.filename == "allowed.txt" {
|
||||
vector[0] = 0.8;
|
||||
vector[1] = 0.6;
|
||||
}
|
||||
let windows = [TextWindow {
|
||||
text: passage.text.clone(),
|
||||
start_byte: 0,
|
||||
end_byte: passage.text.len(),
|
||||
}];
|
||||
store.insert(generation, passage, 0, &windows, &[vector])?;
|
||||
}
|
||||
store.activate(generation)?;
|
||||
let options = SearchOptions {
|
||||
extension: Some("txt".into()),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
store.search(&query, &fingerprint, 1, &options)?[0].filename,
|
||||
"allowed.txt"
|
||||
);
|
||||
let warmed_hashes = HASH_CALLS.with(std::cell::Cell::get);
|
||||
let _ = store.search(&query, &fingerprint, 1, &options)?;
|
||||
assert_eq!(
|
||||
HASH_CALLS.with(std::cell::Cell::get),
|
||||
warmed_hashes,
|
||||
"Warm query rehashed unchanged corpus"
|
||||
);
|
||||
assert_eq!(
|
||||
store.exact_search(&query, &fingerprint, 1, &options)?[0].filename,
|
||||
"allowed.txt"
|
||||
);
|
||||
assert!(!crate::HybridSearch::embedding_status(&db)?.stale);
|
||||
std::fs::write(
|
||||
root.join("allowed.txt"),
|
||||
"replacement changes the generation revision",
|
||||
)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let status = crate::HybridSearch::embedding_status(&db)?;
|
||||
let before_delta = HASH_CALLS.with(std::cell::Cell::get);
|
||||
assert!(store.search(&query, &fingerprint, 1, &options)?.is_empty());
|
||||
assert!(
|
||||
HASH_CALLS.with(std::cell::Cell::get) - before_delta <= 1,
|
||||
"Single file change rehashed unrelated passages"
|
||||
);
|
||||
assert!(status.stale);
|
||||
assert_eq!(status.pending, 2);
|
||||
store.connection.execute(
|
||||
"UPDATE semantic_generations SET ann_blob=x'010203' WHERE id=?1",
|
||||
[generation],
|
||||
)?;
|
||||
drop(store);
|
||||
let mut store = Store::open(&db, DIMENSIONS)?;
|
||||
assert!(store
|
||||
.search(&query, &fingerprint, 1, &SearchOptions::default())
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("SHA256 mismatch"));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "ann_benchmark.rs"]
|
||||
mod ann_benchmark;
|
||||
@@ -0,0 +1,80 @@
|
||||
use anyhow::Result;
|
||||
use searchi_core::SearchIndex;
|
||||
use searchi_service::{HybridSearch, SearchOptions};
|
||||
use std::fs;
|
||||
|
||||
#[test]
|
||||
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
|
||||
fn real_google_search_persists_and_suppresses_changed_deleted_sources() -> Result<()> {
|
||||
assert_eq!(
|
||||
std::env::var("SEARCHI_FORCE_CPU").unwrap(),
|
||||
"1",
|
||||
"VM tests require CPU only"
|
||||
);
|
||||
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 temp = tempfile::tempdir()?;
|
||||
let root = temp.path().join("documents");
|
||||
fs::create_dir(&root)?;
|
||||
let db = temp.path().join("index.db");
|
||||
fs::write(
|
||||
root.join("agreement.txt"),
|
||||
"The supplier partnership may be ended with sixty days written notice.",
|
||||
)?;
|
||||
fs::write(
|
||||
root.join("garden.txt"),
|
||||
"Tomatoes need regular watering and lots of sunlight.",
|
||||
)?;
|
||||
let mut lexical = SearchIndex::open(&db)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let mut hybrid = HybridSearch::open(&db, &model, &worker, 1)?;
|
||||
let report = hybrid.index_embeddings()?;
|
||||
assert_eq!(report.passages, 2);
|
||||
let hits = hybrid.search(
|
||||
"Wie beenden wir die Zusammenarbeit mit dem Lieferanten?",
|
||||
10,
|
||||
&SearchOptions::default(),
|
||||
)?;
|
||||
assert_eq!(hits[0].filename, "agreement.txt");
|
||||
assert!(hits[0].source.is_some());
|
||||
assert!(hits[0].snippet.contains("sixty days"));
|
||||
let exact = hybrid.search("garden.txt", 10, &SearchOptions::default())?;
|
||||
assert_eq!(exact[0].match_kind, "exact_filename");
|
||||
let restricted = hybrid.search(
|
||||
"Wie beenden wir die Zusammenarbeit mit dem Lieferanten?",
|
||||
10,
|
||||
&SearchOptions {
|
||||
extension: Some("pdf".into()),
|
||||
..Default::default()
|
||||
},
|
||||
)?;
|
||||
assert!(restricted.is_empty());
|
||||
drop(hybrid);
|
||||
let mut hybrid = HybridSearch::open(&db, &model, &worker, 1)?;
|
||||
assert_eq!(
|
||||
hybrid.search(
|
||||
"How can we stop working with the vendor?",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?[0]
|
||||
.filename,
|
||||
"agreement.txt"
|
||||
);
|
||||
fs::write(root.join("agreement.txt"),"Entirely unrelated replacement content about repainting a kitchen and changing cupboard handles.")?;
|
||||
lexical.index_folder(&root)?;
|
||||
assert!(hybrid
|
||||
.search_semantic(
|
||||
"How can we stop working with the vendor?",
|
||||
10,
|
||||
&SearchOptions::default()
|
||||
)?
|
||||
.iter()
|
||||
.all(|hit| hit.filename != "agreement.txt"));
|
||||
fs::remove_file(root.join("garden.txt"))?;
|
||||
lexical.index_folder(&root)?;
|
||||
assert!(hybrid
|
||||
.search_semantic("How do we water vegetables?", 10, &SearchOptions::default())?
|
||||
.iter()
|
||||
.all(|hit| hit.filename != "garden.txt"));
|
||||
Ok(())
|
||||
}
|
||||
Reference in new issue
Block a user