1586 lines
63 KiB
Rust
1586 lines
63 KiB
Rust
mod incremental;
|
|
pub(crate) use incremental::SourceTask;
|
|
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::metrics::{MetricsMode, SpanToken, TraceCollector};
|
|
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>,
|
|
min_duration: Option<u32>,
|
|
max_duration: Option<u32>,
|
|
source_type: Option<String>,
|
|
code_only: bool,
|
|
}
|
|
|
|
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, PathBuf);
|
|
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)}; }
|
|
#[cfg(test)]
|
|
thread_local! { static MEDIA_STAT_CALLS:std::cell::Cell<usize>=const {std::cell::Cell::new(0)}; }
|
|
|
|
impl Store {
|
|
pub(crate) fn active_fingerprint(&self) -> Result<Option<String>> {
|
|
Ok(self.connection.query_row("SELECT g.fingerprint FROM semantic_meta m JOIN semantic_generations g ON g.id=m.active_generation WHERE m.singleton=1",[],|r|r.get(0)).optional()?)
|
|
}
|
|
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..=7).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=6 WHERE singleton=1;",
|
|
)?;
|
|
for (column, declaration) in [
|
|
("source_type", "TEXT NOT NULL DEFAULT 'text'"),
|
|
("file_size", "INTEGER"),
|
|
("file_mtime", "INTEGER"),
|
|
] {
|
|
let exists: bool = initialization.query_row(
|
|
"SELECT EXISTS(SELECT 1 FROM pragma_table_info('semantic_chunks') WHERE name=?1)",
|
|
[column],
|
|
|r| r.get(0),
|
|
)?;
|
|
if !exists {
|
|
initialization.execute_batch(&format!(
|
|
"ALTER TABLE semantic_chunks ADD COLUMN {column} {declaration}"
|
|
))?;
|
|
}
|
|
}
|
|
initialization.execute_batch("CREATE TABLE IF NOT EXISTS semantic_media_errors(generation INTEGER NOT NULL REFERENCES semantic_generations(id) ON DELETE CASCADE,path TEXT NOT NULL,message TEXT NOT NULL,PRIMARY KEY(generation,path))")?;
|
|
incremental::migrate(&initialization)?;
|
|
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, self.lock_path.clone()))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
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 media_batch(&self, after: i64) -> Result<Vec<(i64, String, i64, i64)>> {
|
|
let mut q = self.connection.prepare(
|
|
"SELECT rowid,path,size,mtime FROM files WHERE rowid>?1 ORDER BY rowid LIMIT 64",
|
|
)?;
|
|
let rows = q
|
|
.query_map([after], |r| {
|
|
Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?))
|
|
})?
|
|
.collect::<rusqlite::Result<Vec<_>>>()?;
|
|
Ok(rows)
|
|
}
|
|
pub(crate) fn media_failure(&self, generation: i64, path: &str, message: &str) -> Result<()> {
|
|
self.connection.execute(
|
|
"DELETE FROM semantic_chunks WHERE generation=?1 AND path=?2 AND source_type!='text'",
|
|
params![generation, path],
|
|
)?;
|
|
self.connection.execute("INSERT OR REPLACE INTO semantic_media_errors(generation,path,message) VALUES(?1,?2,?3)",params![generation,path,message.chars().take(2048).collect::<String>()])?;
|
|
Ok(())
|
|
}
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(crate) fn insert_media(
|
|
&self,
|
|
generation: i64,
|
|
file_id: i64,
|
|
ordinal: usize,
|
|
identity: &searchi_media::MediaIdentity,
|
|
source_type: &str,
|
|
timestamp: u32,
|
|
vector: &[f32],
|
|
) -> Result<()> {
|
|
ensure!(
|
|
vector.len() == self.dimensions && vector.iter().all(|v| v.is_finite()),
|
|
"Invalid media embedding"
|
|
);
|
|
let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
|
|
ensure!(
|
|
(norm - 1.).abs() < 1e-3,
|
|
"Media embeddings must be normalized"
|
|
);
|
|
let path = identity.path.to_string_lossy();
|
|
let valid:bool=self.connection.query_row("SELECT EXISTS(SELECT 1 FROM files WHERE rowid=?1 AND path=?2 AND size=?3 AND mtime=?4)",params![file_id,path,identity.bytes as i64,identity.modified_ns],|r|r.get(0))?;
|
|
ensure!(valid, "Media file changed since filesystem indexing");
|
|
let kind = if source_type == "image" {
|
|
"image"
|
|
} else {
|
|
"timestamp_ms"
|
|
};
|
|
let snippet = if source_type == "image" {
|
|
"Image content".to_string()
|
|
} else {
|
|
format!(
|
|
"{} content at {:.3} seconds",
|
|
source_type,
|
|
f64::from(timestamp) / 1000.
|
|
)
|
|
};
|
|
let bytes: Vec<u8> = vector.iter().flat_map(|v| v.to_le_bytes()).collect();
|
|
self.connection.execute("INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector,source_type,file_size,file_mtime) VALUES(?1,?2,?3,?4,?5,?6,?7,0,0,?8,?9,?10,?11,?12)",params![generation,-file_id,ordinal as i64,path,kind,timestamp,identity.sha256,snippet,bytes,source_type,identity.bytes as i64,identity.modified_ns])?;
|
|
Ok(())
|
|
}
|
|
|
|
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(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) fn activate(&mut self, generation: i64) -> Result<()> {
|
|
self.activate_revision(generation, None)
|
|
}
|
|
|
|
fn activate_revision(&mut self, generation: i64, revision: Option<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_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
|
|
if let Some(expected) = revision {
|
|
let current: i64 = transaction.query_row(
|
|
"SELECT corpus_revision FROM semantic_meta WHERE singleton=1",
|
|
[],
|
|
|r| r.get(0),
|
|
)?;
|
|
ensure!(
|
|
current == expected,
|
|
"Indexed corpus changed during semantic rebuild; retry required"
|
|
);
|
|
incremental::validate_completion(&transaction, generation)?;
|
|
}
|
|
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)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
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>> {
|
|
self.search_origin(query, fingerprint, limit, options, None)
|
|
}
|
|
pub(crate) fn search_origin(
|
|
&mut self,
|
|
query: &[f32],
|
|
fingerprint: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
source_type: Option<&str>,
|
|
) -> Result<Vec<SearchResult>> {
|
|
let mut trace = TraceCollector::new(MetricsMode::Basic);
|
|
self.search_origin_traced(
|
|
query,
|
|
fingerprint,
|
|
limit,
|
|
options,
|
|
source_type,
|
|
&mut trace,
|
|
None,
|
|
)
|
|
}
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(crate) fn search_origin_traced(
|
|
&mut self,
|
|
query: &[f32],
|
|
fingerprint: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
source_type: Option<&str>,
|
|
trace: &mut TraceCollector,
|
|
parent: Option<SpanToken>,
|
|
) -> Result<Vec<SearchResult>> {
|
|
trace.timed("semantic.retrieval", parent, |trace, parent| {
|
|
self.search_inner(
|
|
query,
|
|
fingerprint,
|
|
limit,
|
|
options,
|
|
source_type,
|
|
trace,
|
|
parent,
|
|
)
|
|
})
|
|
}
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn search_inner(
|
|
&mut self,
|
|
query: &[f32],
|
|
fingerprint: &str,
|
|
limit: usize,
|
|
options: &SearchOptions,
|
|
source_type: Option<&str>,
|
|
trace: &mut TraceCollector,
|
|
parent: Option<SpanToken>,
|
|
) -> 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 generation_span = trace.begin("semantic.generation_snapshot", parent);
|
|
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());
|
|
};
|
|
trace.end(generation_span, Default::default());
|
|
let ann_load = trace.begin("semantic.ann_load", parent);
|
|
trace.counter(
|
|
ann_load,
|
|
"cache_hit",
|
|
u64::from(
|
|
self.ann_cache
|
|
.as_ref()
|
|
.is_some_and(|(generation, _)| *generation == active),
|
|
),
|
|
);
|
|
if !self
|
|
.ann_cache
|
|
.as_ref()
|
|
.is_some_and(|(generation, _)| *generation == active)
|
|
{
|
|
let blob_span = trace.begin("semantic.ann_blob_read", ann_load);
|
|
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)?)),
|
|
)?;
|
|
trace.counter(blob_span, "bytes", serialized.len() as u64);
|
|
trace.end(blob_span, Default::default());
|
|
let checksum_span = trace.begin("semantic.ann_checksum", ann_load);
|
|
ensure!(
|
|
format!("{:x}", Sha256::digest(&serialized)) == checksum,
|
|
"ANN generation SHA256 mismatch; rebuild embeddings explicitly"
|
|
);
|
|
trace.end(checksum_span, Default::default());
|
|
let deserialize_span = trace.begin("semantic.ann_deserialize", ann_load);
|
|
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)?;
|
|
trace.counter(deserialize_span, "chunks", ann.size() as u64);
|
|
trace.end(deserialize_span, Default::default());
|
|
self.ann_cache = Some((active, ann));
|
|
}
|
|
trace.end(ann_load, Default::default());
|
|
let filter_span = trace.begin("semantic.filter_normalization", parent);
|
|
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,
|
|
min_duration: options.min_duration_ms,
|
|
max_duration: options.max_duration_ms,
|
|
source_type: source_type.map(str::to_owned),
|
|
code_only: options.code_only,
|
|
};
|
|
trace.end(filter_span, Default::default());
|
|
let eligibility_span = trace.begin("semantic.eligibility", parent);
|
|
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();
|
|
trace.counter(
|
|
eligibility_span,
|
|
"cache_hit",
|
|
u64::from(existing.as_ref().is_some_and(|c| {
|
|
c.generation == active && c.filter == filter && c.revision == revision
|
|
})),
|
|
);
|
|
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<_>>>()?;
|
|
trace.counter(eligibility_span, "dirty_paths", paths.len() as u64);
|
|
if paths.len() > 1000 {
|
|
cache = collect_eligibility_traced(
|
|
&snapshot,
|
|
active,
|
|
revision,
|
|
&filter,
|
|
None,
|
|
trace,
|
|
eligibility_span,
|
|
)?;
|
|
} 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_traced(
|
|
&snapshot,
|
|
active,
|
|
revision,
|
|
&filter,
|
|
Some(&path),
|
|
trace,
|
|
eligibility_span,
|
|
)?;
|
|
Arc::make_mut(&mut cache.keys).extend(changed.keys.iter().copied());
|
|
cache.by_path.extend(changed.by_path);
|
|
}
|
|
cache.revision = revision;
|
|
}
|
|
}
|
|
cache
|
|
}
|
|
_ => collect_eligibility_traced(
|
|
&snapshot,
|
|
active,
|
|
revision,
|
|
&filter,
|
|
None,
|
|
trace,
|
|
eligibility_span,
|
|
)?,
|
|
};
|
|
ensure!(
|
|
cache.keys.len() <= MAX_SCAN_CHUNKS,
|
|
"Eligibility cache exceeds chunk resource budget"
|
|
);
|
|
trace.counter(eligibility_span, "eligible_chunks", cache.keys.len() as u64);
|
|
trace.end(eligibility_span, Default::default());
|
|
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;
|
|
let mut filesystem_validity = HashMap::<String, bool>::new();
|
|
let mut passes = 0u64;
|
|
loop {
|
|
passes += 1;
|
|
let pass_span = trace.begin("semantic.ann_pass", parent);
|
|
trace.counter(pass_span, "requested_chunks", candidate_count as u64);
|
|
let ann_span = trace.begin("semantic.ann_search", pass_span);
|
|
let matches =
|
|
ann.filtered_search(query, candidate_count, |key| eligible.contains(&key))?;
|
|
trace.counter(ann_span, "returned_chunks", matches.keys.len() as u64);
|
|
trace.end(ann_span, Default::default());
|
|
let materialization_span = trace.begin("semantic.parent_materialization", pass_span);
|
|
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,c.source_type,c.file_size,c.file_mtime 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, size, mtime) =
|
|
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: match row.get::<_, String>(5)?.as_str() {
|
|
"text" => "semantic".into(),
|
|
other => format!("semantic_{other}"),
|
|
},
|
|
source: Some(SourceLocation {
|
|
kind: row.get(3)?,
|
|
number: row.get(4)?,
|
|
}),
|
|
extraction_error: None,
|
|
},
|
|
row.get::<_, Option<i64>>(6)?,
|
|
row.get::<_, Option<i64>>(7)?,
|
|
))
|
|
})?;
|
|
if let (Some(size), Some(mtime)) = (size, mtime) {
|
|
let valid = *filesystem_validity
|
|
.entry(result.path.clone())
|
|
.or_insert_with(|| {
|
|
#[cfg(test)]
|
|
MEDIA_STAT_CALLS.with(|count| count.set(count.get() + 1));
|
|
let validation_span = trace.begin(
|
|
"semantic.parent_filesystem_validation",
|
|
materialization_span,
|
|
);
|
|
let valid =
|
|
std::fs::metadata(&result.path)
|
|
.ok()
|
|
.is_some_and(|metadata| {
|
|
metadata.is_file()
|
|
&& i64::try_from(metadata.len()).ok() == Some(size)
|
|
&& metadata
|
|
.modified()
|
|
.ok()
|
|
.and_then(|t| {
|
|
t.duration_since(std::time::UNIX_EPOCH).ok()
|
|
})
|
|
.and_then(|d| i64::try_from(d.as_nanos()).ok())
|
|
== Some(mtime)
|
|
});
|
|
trace.counter(validation_span, "valid", u64::from(valid));
|
|
trace.end(validation_span, Default::default());
|
|
valid
|
|
});
|
|
if !valid {
|
|
continue;
|
|
}
|
|
}
|
|
if best.contains_key(&result.path) {
|
|
continue;
|
|
}
|
|
best.insert(result.path.clone(), result);
|
|
}
|
|
results = best.into_values().collect::<Vec<_>>();
|
|
trace.counter(materialization_span, "unique_parents", results.len() as u64);
|
|
trace.counter(
|
|
materialization_span,
|
|
"filesystem_checks",
|
|
filesystem_validity.len() as u64,
|
|
);
|
|
trace.end(materialization_span, Default::default());
|
|
trace.end(pass_span, Default::default());
|
|
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);
|
|
trace.counter(parent, "ann_passes", passes);
|
|
trace.counter(parent, "result_files", results.len() as u64);
|
|
trace.counter(
|
|
parent,
|
|
"filesystem_checks",
|
|
filesystem_validity.len() as u64,
|
|
);
|
|
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)
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::too_many_arguments)]
|
|
fn collect_eligibility_traced(
|
|
connection: &Connection,
|
|
generation: i64,
|
|
revision: i64,
|
|
filter: &FilterKey,
|
|
path: Option<&str>,
|
|
trace: &mut TraceCollector,
|
|
parent: Option<SpanToken>,
|
|
) -> Result<EligibilityCache> {
|
|
trace.timed(
|
|
"semantic.eligibility_filtered_scan_and_hash",
|
|
parent,
|
|
|trace, span| {
|
|
let result = collect_eligibility(connection, generation, revision, filter, path);
|
|
if let Ok(cache) = &result {
|
|
trace.counter(span, "eligible_chunks", cache.keys.len() as u64);
|
|
}
|
|
trace.counter(span, "single_dirty_file", u64::from(path.is_some()));
|
|
result
|
|
},
|
|
)
|
|
}
|
|
|
|
fn collect_eligibility(
|
|
connection: &Connection,
|
|
generation: i64,
|
|
revision: i64,
|
|
filter: &FilterKey,
|
|
path: Option<&str>,
|
|
) -> Result<EligibilityCache> {
|
|
let code = if filter.code_only {
|
|
format!(" AND {}", searchi_core::code_file_predicate())
|
|
} else {
|
|
String::new()
|
|
};
|
|
let mut statement=connection.prepare(&format!("SELECT c.rowid,c.passage_id,c.content_hash,p.text,c.path,c.source_type FROM semantic_chunks c
|
|
LEFT JOIN passages p ON c.source_type='text' AND 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 ((c.source_type='text' AND p.rowid IS NOT NULL) OR (c.source_type!='text' AND c.file_size=f.size AND c.file_mtime=f.mtime))
|
|
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 (?10 IS NULL OR EXISTS(SELECT 1 FROM media_metadata m WHERE m.path=f.path AND m.duration_ms>=?10))
|
|
AND (?11 IS NULL OR EXISTS(SELECT 1 FROM media_metadata m WHERE m.path=f.path AND m.duration_ms<=?11))
|
|
AND (?12 IS NULL OR c.source_type=?12)
|
|
AND (?9 IS NULL OR c.path=?9) {code} 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,
|
|
filter.min_duration,
|
|
filter.max_duration,
|
|
filter.source_type
|
|
])?;
|
|
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 row.get::<_, String>(5)? != "text" {
|
|
row.get::<_, String>(2)?
|
|
} else 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 code_filter_precedes_ann_top_k_and_cache_distinguishes_scope() -> Result<()> {
|
|
// Retrieval geometry only: production vectors come from Google.
|
|
let directory = tempfile::tempdir()?;
|
|
let root = directory.path().join("sources");
|
|
std::fs::create_dir(&root)?;
|
|
for n in 0..40 {
|
|
std::fs::write(
|
|
root.join(format!("document-{n}.md")),
|
|
"parse JSON and retry requests",
|
|
)?;
|
|
}
|
|
std::fs::write(
|
|
root.join("network.rs"),
|
|
"fn retry_request() { parse_json(); }",
|
|
)?;
|
|
std::fs::write(root.join("Makefile"), "build: compile_sources\n")?;
|
|
let db = directory.path().join("index.db");
|
|
let mut lexical = SearchIndex::open(&db)?;
|
|
lexical.index_folder(&root)?;
|
|
let mut store = Store::open(&db, DIMENSIONS)?;
|
|
let generation = store.begin("code-filter-proof")?;
|
|
let mut query = vec![0.; DIMENSIONS];
|
|
query[0] = 1.;
|
|
let mut weaker = query.clone();
|
|
weaker[0] = 0.9;
|
|
weaker[1] = (1.0_f32 - 0.9_f32.powi(2)).sqrt();
|
|
for passage in lexical.passage_batch(0, 100)? {
|
|
let vector = if searchi_core::is_code_path(Path::new(&passage.path)) {
|
|
weaker.clone()
|
|
} else {
|
|
query.clone()
|
|
};
|
|
let window = TextWindow {
|
|
text: passage.text.clone(),
|
|
start_byte: 0,
|
|
end_byte: passage.text.len(),
|
|
};
|
|
store.insert(generation, &passage, 0, &[window], &[vector])?;
|
|
}
|
|
store.activate(generation)?;
|
|
assert!(
|
|
store.search(&query, "code-filter-proof", 1, &SearchOptions::default())?[0]
|
|
.filename
|
|
.ends_with(".md")
|
|
);
|
|
let code = SearchOptions {
|
|
code_only: true,
|
|
..Default::default()
|
|
};
|
|
let hits = store.search(&query, "code-filter-proof", 2, &code)?;
|
|
assert_eq!(hits.len(), 2);
|
|
assert!(hits
|
|
.iter()
|
|
.all(|h| searchi_core::is_code_path(Path::new(&h.path))));
|
|
let incompatible = SearchOptions {
|
|
code_only: true,
|
|
extension: Some("md".into()),
|
|
..Default::default()
|
|
};
|
|
assert!(store
|
|
.search(&query, "code-filter-proof", 2, &incompatible)?
|
|
.is_empty());
|
|
std::fs::remove_file(root.join("network.rs"))?;
|
|
lexical.index_folder(&root)?;
|
|
assert_eq!(
|
|
store.search(&query, "code-filter-proof", 2, &code)?.len(),
|
|
1
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn many_video_windows_do_not_crowd_out_other_parent_files() -> Result<()> {
|
|
// Geometry verifies grouped retrieval, not model inference or codecs.
|
|
let directory = tempfile::tempdir()?;
|
|
let root = directory.path().join("media");
|
|
std::fs::create_dir(&root)?;
|
|
for name in ["long.mp4", "other.mp4", "excluded.png"] {
|
|
std::fs::write(root.join(name), b"parent-grouping identity fixture")?;
|
|
}
|
|
let db = directory.path().join("index.db");
|
|
let mut lexical = SearchIndex::open(&db)?;
|
|
lexical.index_folder(&root)?;
|
|
let mut store = Store::open(&db, DIMENSIONS)?;
|
|
let generation = store.begin("grouped-retrieval-proof")?;
|
|
let mut query = vec![0.; DIMENSIONS];
|
|
query[0] = 1.;
|
|
let mut second = query.clone();
|
|
second[0] = 0.95;
|
|
second[1] = (1.0_f32 - 0.95_f32.powi(2)).sqrt();
|
|
for (id, path, _, _) in store.media_batch(0)? {
|
|
let identity = searchi_media::MediaIdentity::read(Path::new(&path), 1024)?;
|
|
if path.ends_with("long.mp4") {
|
|
for ordinal in 0..96 {
|
|
store.insert_media(
|
|
generation,
|
|
id,
|
|
ordinal,
|
|
&identity,
|
|
"video",
|
|
ordinal as u32 * 2000,
|
|
&query,
|
|
)?;
|
|
}
|
|
} else if path.ends_with("other.mp4") {
|
|
store.insert_media(generation, id, 0, &identity, "video", 4000, &second)?;
|
|
} else {
|
|
store.insert_media(generation, id, 0, &identity, "image", 0, &query)?;
|
|
}
|
|
}
|
|
store.activate(generation)?;
|
|
let options = SearchOptions {
|
|
extension: Some("mp4".into()),
|
|
..Default::default()
|
|
};
|
|
MEDIA_STAT_CALLS.with(|count| count.set(0));
|
|
let mut trace = TraceCollector::new(MetricsMode::Debug);
|
|
let results = store.search_origin_traced(
|
|
&query,
|
|
"grouped-retrieval-proof",
|
|
2,
|
|
&options,
|
|
None,
|
|
&mut trace,
|
|
None,
|
|
)?;
|
|
assert_eq!(results.len(), 2);
|
|
assert!(results[0].path.ends_with("long.mp4"));
|
|
assert!(results[1].path.ends_with("other.mp4"));
|
|
assert_eq!(results[1].source.as_ref().unwrap().number, 4000);
|
|
assert_eq!(
|
|
MEDIA_STAT_CALLS.with(std::cell::Cell::get),
|
|
2,
|
|
"Each parent is stat-checked once despite repeated ANN expansions"
|
|
);
|
|
let metrics = trace.finish(results.len());
|
|
let spans = metrics.trace.unwrap();
|
|
assert!(
|
|
spans
|
|
.iter()
|
|
.filter(|s| s.name == "semantic.ann_pass")
|
|
.count()
|
|
> 1
|
|
);
|
|
assert_eq!(
|
|
spans
|
|
.iter()
|
|
.filter(|s| s.name == "semantic.parent_filesystem_validation")
|
|
.count(),
|
|
2
|
|
);
|
|
assert!(
|
|
spans
|
|
.iter()
|
|
.find(|s| s.name == "semantic.retrieval")
|
|
.unwrap()
|
|
.counters["ann_passes"]
|
|
> 1
|
|
);
|
|
for span in &spans {
|
|
if let Some(parent) = span.parent {
|
|
let parent = &spans[parent];
|
|
assert!(span.start_ms >= parent.start_ms);
|
|
assert!(
|
|
span.start_ms + span.duration_ms <= parent.start_ms + parent.duration_ms + 1e-6
|
|
);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn media_origins_persist_filters_and_changes_remove_stale_candidates() -> Result<()> {
|
|
let directory = tempfile::tempdir()?;
|
|
let root = directory.path().join("media");
|
|
std::fs::create_dir(&root)?;
|
|
let file = root.join("clip.mp4");
|
|
std::fs::write(&file, b"storage identity fixture; not an inference test")?;
|
|
let db = directory.path().join("index.db");
|
|
let mut lexical = SearchIndex::open(&db)?;
|
|
lexical.index_folder(&root)?;
|
|
let mut store = Store::open(&db, DIMENSIONS)?;
|
|
let (id, path, _, _) = store.media_batch(0)?.remove(0);
|
|
let identity = searchi_media::MediaIdentity::read(Path::new(&path), 1024)?;
|
|
lexical.set_media_metadata(
|
|
&path,
|
|
"artist: real extracted metadata fixture",
|
|
65000,
|
|
identity.bytes,
|
|
identity.modified_ns,
|
|
)?;
|
|
let mut vector = vec![0.; DIMENSIONS];
|
|
vector[0] = 1.;
|
|
let generation = store.begin("media-storage-proof")?;
|
|
store.insert_media(generation, id, 0, &identity, "video", 2345, &vector)?;
|
|
store.insert_media(generation, id, 1, &identity, "audio", 28000, &vector)?;
|
|
store.activate(generation)?;
|
|
drop(store);
|
|
let mut store = Store::open(&db, DIMENSIONS)?;
|
|
let mut cold = TraceCollector::new(MetricsMode::Debug);
|
|
let results = store.search_origin_traced(
|
|
&vector,
|
|
"media-storage-proof",
|
|
5,
|
|
&SearchOptions::default(),
|
|
None,
|
|
&mut cold,
|
|
None,
|
|
)?;
|
|
let cold = cold.finish(results.len()).trace.unwrap();
|
|
assert_eq!(
|
|
cold.iter()
|
|
.find(|s| s.name == "semantic.ann_load")
|
|
.unwrap()
|
|
.counters["cache_hit"],
|
|
0
|
|
);
|
|
assert!(cold.iter().any(|s| s.name == "semantic.ann_checksum"));
|
|
let mut warm = TraceCollector::new(MetricsMode::Debug);
|
|
let warm_results = store.search_origin_traced(
|
|
&vector,
|
|
"media-storage-proof",
|
|
5,
|
|
&SearchOptions::default(),
|
|
None,
|
|
&mut warm,
|
|
None,
|
|
)?;
|
|
assert_eq!(
|
|
serde_json::to_value(&results)?,
|
|
serde_json::to_value(&warm_results)?
|
|
);
|
|
let warm = warm.finish(warm_results.len()).trace.unwrap();
|
|
assert_eq!(
|
|
warm.iter()
|
|
.find(|s| s.name == "semantic.ann_load")
|
|
.unwrap()
|
|
.counters["cache_hit"],
|
|
1
|
|
);
|
|
assert_eq!(
|
|
warm.iter()
|
|
.find(|s| s.name == "semantic.eligibility")
|
|
.unwrap()
|
|
.counters["cache_hit"],
|
|
1
|
|
);
|
|
assert!(!warm.iter().any(|s| s.name == "semantic.ann_blob_read"));
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].source.as_ref().unwrap().kind, "timestamp_ms");
|
|
assert!([2345, 28000].contains(&results[0].source.as_ref().unwrap().number));
|
|
let audio = store.search_origin(
|
|
&vector,
|
|
"media-storage-proof",
|
|
5,
|
|
&SearchOptions::default(),
|
|
Some("audio"),
|
|
)?;
|
|
assert_eq!(audio.len(), 1);
|
|
assert_eq!(audio[0].match_kind, "semantic_audio");
|
|
assert_eq!(audio[0].source.as_ref().unwrap().number, 28000);
|
|
let video = store.search_origin(
|
|
&vector,
|
|
"media-storage-proof",
|
|
5,
|
|
&SearchOptions::default(),
|
|
Some("video"),
|
|
)?;
|
|
assert_eq!(video.len(), 1);
|
|
assert_eq!(video[0].match_kind, "semantic_video");
|
|
assert_eq!(video[0].source.as_ref().unwrap().number, 2345);
|
|
assert!(store
|
|
.search_origin(
|
|
&vector,
|
|
"media-storage-proof",
|
|
5,
|
|
&SearchOptions::default(),
|
|
Some("image")
|
|
)?
|
|
.is_empty());
|
|
assert!(
|
|
results[0].match_kind == "semantic_video" || results[0].match_kind == "semantic_audio"
|
|
);
|
|
let filter = SearchOptions {
|
|
extension: Some("png".into()),
|
|
..SearchOptions::default()
|
|
};
|
|
assert!(store
|
|
.search(&vector, "media-storage-proof", 5, &filter)?
|
|
.is_empty());
|
|
let duration_filter = SearchOptions {
|
|
min_duration_ms: Some(66000),
|
|
..Default::default()
|
|
};
|
|
assert!(store
|
|
.search(&vector, "media-storage-proof", 5, &duration_filter)?
|
|
.is_empty());
|
|
let duration_filter = SearchOptions {
|
|
min_duration_ms: Some(64000),
|
|
max_duration_ms: Some(66000),
|
|
..Default::default()
|
|
};
|
|
assert_eq!(
|
|
store
|
|
.search(&vector, "media-storage-proof", 5, &duration_filter)?
|
|
.len(),
|
|
1
|
|
);
|
|
std::fs::write(&file, b"changed")?;
|
|
// Foreground hits check only actual candidate metadata, before watcher
|
|
// reconciliation; they never rehash the complete filesystem corpus.
|
|
assert!(store
|
|
.search(&vector, "media-storage-proof", 5, &SearchOptions::default())?
|
|
.is_empty());
|
|
lexical.index_folder(&root)?;
|
|
assert!(store
|
|
.search(&vector, "media-storage-proof", 5, &SearchOptions::default())?
|
|
.is_empty());
|
|
assert!(store
|
|
.insert_media(generation, id, 2, &identity, "video", 0, &vector)
|
|
.is_err());
|
|
std::fs::remove_file(&file)?;
|
|
lexical.index_folder(&root)?;
|
|
assert!(store
|
|
.search(&vector, "media-storage-proof", 5, &SearchOptions::default())?
|
|
.is_empty());
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn media_failure_removes_partial_vectors_and_persists_reason() -> Result<()> {
|
|
let directory = tempfile::tempdir()?;
|
|
let file = directory.path().join("picture.png");
|
|
std::fs::write(&file, b"identity fixture")?;
|
|
let db = directory.path().join("index.db");
|
|
let mut lexical = SearchIndex::open(&db)?;
|
|
lexical.index_folder(directory.path())?;
|
|
let mut store = Store::open(&db, DIMENSIONS)?;
|
|
let (id, path, _, _) = store
|
|
.media_batch(0)?
|
|
.into_iter()
|
|
.find(|(_, p, _, _)| p.ends_with("picture.png"))
|
|
.unwrap();
|
|
let identity = searchi_media::MediaIdentity::read(Path::new(&path), 1024)?;
|
|
let mut vector = vec![0.; DIMENSIONS];
|
|
vector[0] = 1.;
|
|
let generation = store.begin("failure-proof")?;
|
|
store.insert_media(generation, id, 0, &identity, "image", 0, &vector)?;
|
|
store.media_failure(generation, &path, "Decoder rejected corrupt image")?;
|
|
let chunks: usize = store.connection.query_row(
|
|
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
|
|
[generation],
|
|
|r| r.get(0),
|
|
)?;
|
|
assert_eq!(chunks, 0);
|
|
let reason: String = store.connection.query_row(
|
|
"SELECT message FROM semantic_media_errors WHERE generation=?1",
|
|
[generation],
|
|
|r| r.get(0),
|
|
)?;
|
|
assert!(reason.contains("corrupt"));
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn deleted_media_without_reconciliation_does_not_hide_later_valid_top_k() -> Result<()> {
|
|
let directory = tempfile::tempdir()?;
|
|
let root = directory.path().join("images");
|
|
std::fs::create_dir(&root)?;
|
|
for name in ["a.png", "b.png"] {
|
|
std::fs::write(root.join(name), b"storage fixture")?;
|
|
}
|
|
let db = directory.path().join("index.db");
|
|
let mut lexical = SearchIndex::open(&db)?;
|
|
lexical.index_folder(&root)?;
|
|
let mut store = Store::open(&db, DIMENSIONS)?;
|
|
let generation = store.begin("candidate-stat-proof")?;
|
|
let mut vector = vec![0.; DIMENSIONS];
|
|
vector[0] = 1.;
|
|
for (id, path, _, _) in store.media_batch(0)? {
|
|
let identity = searchi_media::MediaIdentity::read(Path::new(&path), 1024)?;
|
|
store.insert_media(generation, id, 0, &identity, "image", 0, &vector)?;
|
|
}
|
|
store.activate(generation)?;
|
|
std::fs::remove_file(root.join("a.png"))?;
|
|
let results = store.search(
|
|
&vector,
|
|
"candidate-stat-proof",
|
|
1,
|
|
&SearchOptions::default(),
|
|
)?;
|
|
assert_eq!(results.len(), 1);
|
|
assert!(results[0].path.ends_with("b.png"));
|
|
std::fs::remove_file(root.join("b.png"))?;
|
|
assert!(store
|
|
.search(
|
|
&vector,
|
|
"candidate-stat-proof",
|
|
1,
|
|
&SearchOptions::default()
|
|
)?
|
|
.is_empty());
|
|
Ok(())
|
|
}
|
|
|
|
#[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;
|