feat: add multimodal search and granular performance diagnostics
Core CI / core (ubuntu-latest) (push) Failing after 3m31s
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 5m39s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s

This commit is contained in:
karl committed 2026-10-12 19:16:23 +02:00
1 parent d3bb8c7618
commit b8400c4fbf
110 files changed
+79395 -250

No files matched your search

+109 -15
View File
@@ -1,9 +1,10 @@
//! Thin command-line adapter around the shared, CPU-only search engine.
//! Thin headless adapter around the shared local search engine.
use std::path::PathBuf;
use anyhow::{bail, Context, Result};
use clap::{Parser, Subcommand};
use searchi_core::metrics::{MetricsMode, SearchResponse, TraceCollector};
use searchi_core::{IndexOptions, SearchIndex, SearchOptions};
use searchi_service::HybridSearch;
@@ -46,6 +47,12 @@ enum Command {
/// Emit a JSON array for scripting.
#[arg(long)]
json: bool,
/// Include basic measured statistics in JSON output.
#[arg(long)]
performance: bool,
/// Include bounded granular stage timings and diagnostic counters.
#[arg(long)]
debug_performance: bool,
/// Restrict results to this folder and descendants.
#[arg(long)]
folder: Option<PathBuf>,
@@ -56,6 +63,11 @@ enum Command {
min_size: Option<u64>,
#[arg(long)]
max_size: Option<u64>,
/// Inclusive media duration limit in milliseconds.
#[arg(long)]
min_duration_ms: Option<u32>,
#[arg(long)]
max_duration_ms: Option<u32>,
/// Inclusive modification time in Unix milliseconds.
#[arg(long)]
modified_after_ms: Option<i64>,
@@ -80,6 +92,14 @@ enum Command {
embedding_worker: PathBuf,
#[arg(long, default_value_t = 1, value_parser = clap::value_parser!(u8).range(1..=8))]
cpu_threads: u8,
/// Standalone decoder; normally discovered beside the installed CLI.
#[arg(long, requires = "ffprobe")]
ffmpeg: Option<PathBuf>,
#[arg(long, requires = "ffmpeg")]
ffprobe: Option<PathBuf>,
/// Explicitly build only document vectors, without media.
#[arg(long)]
text_only: bool,
},
/// Show the indexed file and content counts.
Stats {
@@ -122,10 +142,14 @@ fn main() -> Result<()> {
query,
limit,
json,
performance,
debug_performance,
folder,
extension,
min_size,
max_size,
min_duration_ms,
max_duration_ms,
modified_after_ms,
modified_before_ms,
model_dir,
@@ -141,30 +165,56 @@ fn main() -> Result<()> {
extension,
min_size,
max_size,
min_duration_ms,
max_duration_ms,
modified_after_unix_ms: modified_after_ms,
modified_before_unix_ms: modified_before_ms,
};
let mut trace = TraceCollector::new(if debug_performance {
MetricsMode::Debug
} else {
MetricsMode::Basic
});
let results = if let (Some(model), Some(runtime)) = (model_dir, embedding_worker) {
HybridSearch::open(&db, model, runtime, usize::from(cpu_threads))?
.search(&query, usize::from(limit), &options)
let mut service = trace.timed("cli.session_initialize", None, |_, _| {
HybridSearch::open(&db, model, runtime, usize::from(cpu_threads))
})?;
service
.search_traced(&query, usize::from(limit), &options, &mut trace, None)
.context("Hybrid search failed")?
} else {
SearchIndex::open(&db)
.with_context(|| format!("Cannot open index {}", db.display()))?
.search_with_options(&query, usize::from(limit), &options)
let index = trace
.timed("cli.index_open", None, |_, _| SearchIndex::open(&db))
.with_context(|| format!("Cannot open index {}", db.display()))?;
index
.search_with_options_traced(
&query,
usize::from(limit),
&options,
&mut trace,
None,
)
.context("Search failed")?
};
let response = SearchResponse {
metrics: trace.finish(results.len()),
results,
};
if json {
println!("{}", serde_json::to_string_pretty(&results)?);
} else if results.is_empty() {
if performance || debug_performance {
println!("{}", serde_json::to_string_pretty(&response)?);
} else {
println!("{}", serde_json::to_string_pretty(&response.results)?);
}
} else if response.results.is_empty() {
println!("No matching files.");
} else {
for result in results {
for result in &response.results {
println!("{} [{}]", result.path, result.match_kind);
if let Some(source) = result.source {
if let Some(source) = &result.source {
println!(" {} {}", source.kind, source.number);
}
if let Some(error) = result.extraction_error {
if let Some(error) = &result.extraction_error {
println!(" Text extraction unavailable: {error}");
}
if !result.snippet.is_empty() {
@@ -172,18 +222,36 @@ fn main() -> Result<()> {
}
}
}
if !json {
eprintln!(
"{} results · {:.2} ms (includes CLI initialization)",
response.metrics.result_count, response.metrics.total_ms
);
if debug_performance {
eprintln!("{}", serde_json::to_string_pretty(&response.metrics)?);
}
}
}
Command::Embed {
db,
model_dir,
embedding_worker,
cpu_threads,
ffmpeg,
ffprobe,
text_only,
} => {
require_existing_index(&db)?;
let report =
HybridSearch::open(&db, model_dir, embedding_worker, usize::from(cpu_threads))?
.index_embeddings()
.context("Semantic indexing failed; prior generation remains active")?;
let mut service =
HybridSearch::open(&db, model_dir, embedding_worker, usize::from(cpu_threads))?;
let report = if text_only {
service.index_embeddings()
} else {
let tools = resolve_media_tools(ffmpeg, ffprobe)?;
service
.index_embeddings_with_media(&tools, &searchi_service::MediaLimits::default())
}
.context("Semantic indexing failed; prior generation remains active")?;
println!("{}", serde_json::to_string_pretty(&report)?);
}
Command::Stats { db, json } => {
@@ -204,6 +272,32 @@ fn main() -> Result<()> {
Ok(())
}
fn resolve_media_tools(
ffmpeg: Option<PathBuf>,
ffprobe: Option<PathBuf>,
) -> Result<searchi_service::FfmpegTools> {
if let (Some(ffmpeg), Some(ffprobe)) = (ffmpeg, ffprobe) {
return searchi_service::FfmpegTools::new(ffmpeg, ffprobe);
}
if let Some(tools) = searchi_service::FfmpegTools::from_environment()? {
return Ok(tools);
}
let development =
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../.searchi/media-runtime/bin");
let directory = if cfg!(debug_assertions) && development.is_dir() {
development
} else {
std::env::current_exe()?
.parent()
.context("CLI installation directory is unavailable")?
.join("media-runtime/bin")
};
searchi_service::FfmpegTools::new(
directory.join(if cfg!(windows) { "ffmpeg.exe" } else { "ffmpeg" }),
directory.join(if cfg!(windows) { "ffprobe.exe" } else { "ffprobe" }),
).context("Media tools are missing from this installation; repair the headless package or explicitly use --text-only")
}
/// Read commands must not silently create a new empty database after a typo.
fn require_existing_index(db: &std::path::Path) -> Result<()> {
if !db.is_file() {
+30
View File
@@ -68,6 +68,36 @@ fn indexes_persists_and_returns_search_and_stats_json() {
.unwrap()
.contains("termination"));
for flag in ["--performance", "--debug-performance"] {
let measured = fixture.run(&["search", "--db", "index.db", "termination", "--json", flag]);
assert!(
measured.status.success(),
"{}",
String::from_utf8_lossy(&measured.stderr)
);
let envelope: serde_json::Value = serde_json::from_slice(&measured.stdout).unwrap();
assert_eq!(
envelope["results"],
serde_json::Value::Array(results.clone())
);
assert_eq!(envelope["metrics"]["result_count"], 1);
let total = envelope["metrics"]["total_ms"].as_f64().unwrap();
assert!(total.is_finite() && total >= 0.0);
if flag == "--performance" {
assert!(envelope["metrics"]["trace"].is_null());
} else {
let spans = envelope["metrics"]["trace"].as_array().unwrap();
assert!(spans.iter().any(|span| span["name"] == "cli.index_open"));
assert!(spans.len() > 2);
for span in spans {
let start = span["start_ms"].as_f64().unwrap();
let duration = span["duration_ms"].as_f64().unwrap();
assert!(start >= 0.0 && duration >= 0.0 && start + duration <= total + 0.001);
assert!(!span["name"].as_str().unwrap().contains("termination"));
}
}
}
let stats = fixture.run(&["stats", "--db", "index.db", "--json"]);
assert!(stats.status.success());
let stats: serde_json::Value = serde_json::from_slice(&stats.stdout).unwrap();
@@ -0,0 +1,80 @@
//! Small prepared SQLite instrumentation diagnostic; not a production benchmark.
use anyhow::{ensure, Result};
use searchi_core::SearchIndex;
use std::{fs, path::Path, time::Instant};
fn distribution(samples: &[f64]) -> serde_json::Value {
let mut sorted = samples.to_vec();
sorted.sort_by(f64::total_cmp);
let percentile = |fraction: f64| {
sorted[((sorted.len() as f64 * fraction).ceil() as usize).saturating_sub(1)]
};
serde_json::json!({"samples":sorted.len(),"p50_ms":percentile(0.50),"p95_ms":percentile(0.95),"mean_ms":sorted.iter().sum::<f64>()/sorted.len() as f64})
}
fn main() -> Result<()> {
let fixture = tempfile::tempdir()?;
let root = fixture.path().join("files");
fs::create_dir(&root)?;
for index in 0..1000 {
fs::write(
root.join(format!("record-{index:04}.txt")),
format!(
"Supplier agreement {index}. Cancellation needs written notice. Group {}.",
index % 31
),
)?;
}
let mut index = SearchIndex::open(fixture.path().join("index.sqlite"))?;
let report = index.index_folder(&root)?;
index.prewarm()?;
ensure!(report.indexed == 1000, "Unexpected fixture indexing count");
let queries: Vec<String> = (0..24)
.map(|index| format!("record-{:04}.txt", index * 37))
.collect();
for query in &queries {
index.search_measured(query, 10, false)?;
index.search_measured(query, 10, true)?;
}
let mut basic = Vec::new();
let mut debug = Vec::new();
let mut trace_spans = Vec::new();
// Alternate order to reduce monotonic cache/CPU-order effects.
for repetition in 0..20 {
for query in &queries {
for detailed in if repetition % 2 == 0 {
[false, true]
} else {
[true, false]
} {
let started = Instant::now();
let response = index.search_measured(query, 10, detailed)?;
let elapsed = started.elapsed().as_secs_f64() * 1000.0;
ensure!(
response
.results
.first()
.is_some_and(|result| &result.filename == query),
"Fixture retrieval changed"
);
if detailed {
trace_spans.push(response.metrics.trace.as_ref().unwrap().len());
debug.push(elapsed);
} else {
ensure!(
response.metrics.trace.is_none(),
"Basic unexpectedly stored traces"
);
basic.push(elapsed);
}
}
}
}
let output = serde_json::json!({"kind":"prepared_sqlite_trace_overhead","production_claim":false,"profile":if cfg!(debug_assertions){"debug"}else{"release"},"os":std::env::consts::OS,"arch":std::env::consts::ARCH,"files":1000,"queries":24,"repetitions":20,"measurement":"outer Rust call including metrics assembly; excludes fixture setup/JSON serialization","basic":distribution(&basic),"debug":distribution(&debug),"basic_samples_ms":basic,"debug_samples_ms":debug,"debug_span_count_min":trace_spans.iter().min(),"debug_span_count_max":trace_spans.iter().max(),"limits":"tiny warm lexical corpus; harness serial/no model/media/GPU; other VM work including native photo inference may contend; not isolated production performance"});
let serialized = serde_json::to_string_pretty(&output)?;
if let Some(path) = std::env::args_os().nth(1) {
fs::write(Path::new(&path), serialized)?;
} else {
println!("{serialized}");
}
Ok(())
}
+255 -6
View File
@@ -4,7 +4,10 @@
//! Text and document extraction have byte limits, but still run in this process.
//! Worker isolation remains a production milestone.
mod extraction;
pub mod metrics;
use anyhow::{bail, Context, Result};
pub use metrics::{MetricsMode, SearchMetrics, SearchResponse, SearchSpan};
use metrics::{SpanToken, TraceCollector};
use rusqlite::{params, params_from_iter, types::Value, Connection, OptionalExtension};
use serde::Serialize;
use std::{
@@ -16,7 +19,7 @@ use std::{
};
use walkdir::WalkDir;
const SCHEMA_VERSION: i64 = 9;
const SCHEMA_VERSION: i64 = 10;
const MAX_RESULTS: usize = 1000;
#[derive(Debug, Default, Serialize)]
@@ -53,6 +56,8 @@ pub struct SearchOptions {
pub max_size: Option<u64>,
pub modified_after_unix_ms: Option<i64>,
pub modified_before_unix_ms: Option<i64>,
pub min_duration_ms: Option<u32>,
pub max_duration_ms: Option<u32>,
}
#[derive(Debug, Serialize)]
@@ -66,7 +71,8 @@ pub struct SearchResult {
pub extraction_error: Option<String>,
}
/// A verified matching passage location; numbers are one based.
/// Verified source location. Document numbers are one based; timestamp_ms
/// numbers are milliseconds from media start, and image numbers are zero.
#[derive(Debug, Clone, Serialize)]
pub struct SourceLocation {
pub kind: String,
@@ -227,6 +233,15 @@ impl SearchIndex {
UPDATE name_statistics_state SET dirty=1 WHERE singleton=1; END;
PRAGMA user_version=9;",
)?;
initialization.execute_batch("CREATE INDEX IF NOT EXISTS files_extension ON files(extension,id);
CREATE TABLE IF NOT EXISTS media_metadata(path TEXT PRIMARY KEY, file_id INTEGER NOT NULL UNIQUE, tags TEXT NOT NULL,duration_ms INTEGER NOT NULL);
CREATE VIRTUAL TABLE IF NOT EXISTS media_tags USING fts5(tags,content='media_metadata',content_rowid='file_id');
CREATE TRIGGER IF NOT EXISTS media_metadata_insert AFTER INSERT ON media_metadata BEGIN INSERT INTO media_tags(rowid,tags) VALUES(new.file_id,new.tags); END;
CREATE TRIGGER IF NOT EXISTS media_metadata_delete AFTER DELETE ON media_metadata BEGIN INSERT INTO media_tags(media_tags,rowid,tags) VALUES('delete',old.file_id,old.tags); END;
CREATE TRIGGER IF NOT EXISTS media_metadata_update AFTER UPDATE ON media_metadata BEGIN INSERT INTO media_tags(media_tags,rowid,tags) VALUES('delete',old.file_id,old.tags); INSERT INTO media_tags(rowid,tags) VALUES(new.file_id,new.tags); END;
CREATE TRIGGER IF NOT EXISTS files_media_delete AFTER DELETE ON files BEGIN DELETE FROM media_metadata WHERE path=old.path; END;
CREATE TRIGGER IF NOT EXISTS files_media_update AFTER UPDATE ON files WHEN new.path!=old.path OR new.size!=old.size OR new.mtime!=old.mtime BEGIN DELETE FROM media_metadata WHERE path=old.path; END;
PRAGMA user_version=10;")?;
initialization.commit()?;
}
Ok(Self {
@@ -236,6 +251,32 @@ impl SearchIndex {
})
}
/// Persist genuine extracted media tags, guarded against filesystem changes.
pub fn set_media_metadata(
&mut self,
path: &str,
tags: &str,
duration_ms: u32,
size: u64,
mtime_ns: i64,
) -> Result<()> {
if tags.len() > 32768 {
bail!("Media tags exceed 32 KiB bound");
}
let tx = self.connection.transaction()?;
let id: Option<i64> = tx
.query_row(
"SELECT id FROM files WHERE path=?1 AND size=?2 AND mtime=?3",
params![path, i64::try_from(size)?, mtime_ns],
|r| r.get(0),
)
.optional()?;
let id = id.context("Media metadata file changed since indexing")?;
tx.execute("INSERT INTO media_metadata(path,file_id,tags,duration_ms) VALUES(?1,?2,?3,?4) ON CONFLICT(path) DO UPDATE SET tags=excluded.tags,duration_ms=excluded.duration_ms",params![path,id,tags,duration_ms])?;
tx.commit()?;
Ok(())
}
/// Refresh candidate-frequency statistics off the foreground search path.
/// Stale/missing estimates can affect speed, never literal-match correctness.
/// The app invokes this during background startup and after reconciliation.
@@ -465,6 +506,56 @@ impl SearchIndex {
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let mut trace = TraceCollector::new(MetricsMode::Basic);
self.search_with_options_traced(query, limit, options, &mut trace, None)
}
pub fn search_measured(
&self,
query: &str,
limit: usize,
debug: bool,
) -> Result<SearchResponse<SearchResult>> {
self.search_with_options_measured(query, limit, &SearchOptions::default(), debug)
}
pub fn search_with_options_measured(
&self,
query: &str,
limit: usize,
options: &SearchOptions,
debug: bool,
) -> Result<SearchResponse<SearchResult>> {
let mut trace = TraceCollector::new(if debug {
MetricsMode::Debug
} else {
MetricsMode::Basic
});
let results = self.search_with_options_traced(query, limit, options, &mut trace, None)?;
let metrics = trace.finish(results.len());
Ok(SearchResponse { results, metrics })
}
/// Share one monotonic request collector with downstream retrieval lanes.
pub fn search_with_options_traced(
&self,
query: &str,
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
trace.timed("lexical.search", parent, |trace, parent| {
self.search_inner(query, limit, options, trace, parent)
})
}
fn search_inner(
&self,
query: &str,
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
let normalization = trace.begin("lexical.normalization", parent);
if options
.min_size
.zip(options.max_size)
@@ -479,6 +570,13 @@ impl SearchIndex {
{
bail!("modified-after exceeds modified-before");
}
if options
.min_duration_ms
.zip(options.max_duration_ms)
.is_some_and(|(a, b)| a > b)
{
bail!("Minimum duration exceeds maximum duration");
}
let folder = options
.folder
.as_ref()
@@ -518,11 +616,16 @@ impl SearchIndex {
.map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(" AND ");
trace.counter(normalization, "query_bytes", query.len() as u64);
trace.end(normalization, Default::default());
let planning = trace.begin("lexical.candidate_plan", parent);
let limit = limit.min(MAX_RESULTS);
let filters = " AND (?3 IS NULL OR substr(f.path,1,length(?3))=?3)
AND (?4 IS NULL OR f.extension=?4)
AND (?5 IS NULL OR f.size>=?5) AND (?6 IS NULL OR f.size<=?6)
AND (?7 IS NULL OR f.mtime/1000000>=?7) AND (?8 IS NULL OR f.mtime/1000000<=?8)";
AND (?7 IS NULL OR f.mtime/1000000>=?7) AND (?8 IS NULL OR f.mtime/1000000<=?8)
AND (?10 IS NULL OR EXISTS(SELECT 1 FROM media_metadata m WHERE m.file_id=f.id AND m.duration_ms>=?10))
AND (?11 IS NULL OR EXISTS(SELECT 1 FROM media_metadata m WHERE m.file_id=f.id AND m.duration_ms<=?11))";
let filter_values = [
folder_prefix.map_or(Value::Null, Value::Text),
extension.map_or(Value::Null, Value::Text),
@@ -535,7 +638,9 @@ impl SearchIndex {
.modified_before_unix_ms
.map_or(Value::Null, Value::Integer),
];
let has_filters = filter_values.iter().any(|value| *value != Value::Null);
let has_filters = filter_values.iter().any(|value| *value != Value::Null)
|| options.min_duration_ms.is_some()
|| options.max_duration_ms.is_some();
let mut results = Vec::new();
let mut selected = Vec::<i64>::new();
// Each lane filters its entire candidate set before LIMIT. There is no
@@ -561,13 +666,17 @@ impl SearchIndex {
if cache.len() + grams.len() > 4096 {
cache.clear();
}
let mut frequency_hits = 0u64;
let mut frequency_misses = 0u64;
let mut rare_first = grams
.into_iter()
.map(|gram| {
let term = gram.to_lowercase();
let frequency = if let Some(frequency) = cache.get(&term) {
frequency_hits += 1;
*frequency
} else {
frequency_misses += 1;
let frequency = vocabulary
.query_row([&term], |row| row.get::<_, i64>(0))
.optional()?
@@ -581,6 +690,9 @@ impl SearchIndex {
rare_first.sort();
drop(cache);
let anchor = substring_anchor(&characters, &rare_first);
trace.counter(planning, "frequency_cache_hits", frequency_hits);
trace.counter(planning, "frequency_cache_misses", frequency_misses);
trace.counter(planning, "trigram_index", u64::from(trigram));
let grams = format!("\"{}\"", anchor.replace('"', "\"\""));
let mut lanes = vec![
(
@@ -648,10 +760,30 @@ impl SearchIndex {
"rank,f.id",
));
}
lanes.push((
"media_metadata",
"media_tags CROSS JOIN files f ON f.id=media_tags.rowid".into(),
"media_tags MATCH ?1".into(),
fts.clone(),
"rank,f.id",
));
trace.end(planning, Default::default());
for (kind, from, condition, expression, ordering) in lanes {
if results.len() == limit {
break;
}
let lane = trace.begin(
match kind {
"exact_filename" => "lexical.exact_filename",
"filename" => "lexical.filename_substring",
"path" => "lexical.path_substring",
"name_terms" => "lexical.filename_terms",
"content" => "lexical.content",
"media_metadata" => "lexical.media_metadata",
_ => "lexical.other",
},
parent,
);
let exclusions = if selected.is_empty() {
String::new()
} else {
@@ -664,12 +796,14 @@ impl SearchIndex {
.join(",")
)
};
let rank = if matches!(kind, "content" | "name_terms") {
let rank = if matches!(kind, "content" | "name_terms" | "media_metadata") {
"rank"
} else {
"0.0"
};
let snippet = if kind == "content" {
let snippet = if kind == "media_metadata" {
"snippet(media_tags,0,'','',' … ',24)"
} else if kind == "content" {
"snippet(content_search,0,'','',' … ',24)"
} else if kind == "name_terms" {
"snippet(search,2,'','',' … ',24)"
@@ -698,6 +832,20 @@ impl SearchIndex {
let mut values = vec![Value::Text(expression), Value::Text(query.into())];
values.extend(filter_values.iter().cloned());
values.push(Value::Integer((limit - results.len()) as i64));
if sql.contains("?10") {
values.push(
options
.min_duration_ms
.map_or(Value::Null, |v| Value::Integer(i64::from(v))),
);
values.push(
options
.max_duration_ms
.map_or(Value::Null, |v| Value::Integer(i64::from(v))),
);
}
let sql_span = trace.begin("lexical.filtered_sql", lane);
trace.counter(sql_span, "filters_active", u64::from(has_filters));
let mut statement = self.connection.prepare(&sql)?;
let matches = statement
.query_map(params_from_iter(values), |row| {
@@ -715,12 +863,16 @@ impl SearchIndex {
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
trace.counter(sql_span, "rows", matches.len() as u64);
trace.end(sql_span, Default::default());
for (id, result) in matches {
selected.push(id);
results.push(result);
}
trace.end(lane, Default::default());
}
if !results.is_empty() {
let source_span = trace.begin("lexical.source_locations", parent);
// Ownership has a B-tree index. The rowid IN constraint intersects
// the FTS postings with only the actual finalists' passage IDs.
let placeholders = (0..results.len())
@@ -762,6 +914,9 @@ impl SearchIndex {
result.snippet = snippet.clone();
}
}
trace.counter(source_span, "files", results.len() as u64);
trace.counter(source_span, "matched_sources", sources.len() as u64);
trace.end(source_span, Default::default());
}
Ok(results)
}
@@ -856,6 +1011,100 @@ fn substring_anchor(characters: &[char], frequencies: &[(i64, String)]) -> Strin
#[cfg(test)]
mod tests {
#[test]
fn measured_search_preserves_results_and_contains_private_free_nested_spans(
) -> anyhow::Result<()> {
let directory = tempfile::tempdir()?;
let root = directory.path().join("sensitive-private-location");
std::fs::create_dir(&root)?;
std::fs::write(
root.join("agreement.txt"),
"confidentialneedle cancellation policy",
)?;
let db = directory.path().join("index.db");
let mut index = super::SearchIndex::open(&db)?;
index.index_folder(&root)?;
let basic = index.search_measured("confidentialneedle", 10, false)?;
let debug = index.search_measured("confidentialneedle", 10, true)?;
assert_eq!(
serde_json::to_value(&basic.results)?,
serde_json::to_value(&debug.results)?
);
assert_eq!(basic.metrics.result_count, 1);
assert!(basic.metrics.trace.is_none());
let spans = debug.metrics.trace.as_ref().unwrap();
assert!(spans.iter().any(|s| s.name == "lexical.content"));
assert!(spans.iter().any(|s| s.name == "lexical.source_locations"));
for span in spans {
assert!(span.start_ms >= 0.0 && span.duration_ms >= 0.0);
assert!(span.start_ms + span.duration_ms <= debug.metrics.total_ms + 1e-6);
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
);
}
}
let trace = serde_json::to_string(spans)?;
assert!(!trace.contains("confidentialneedle"));
assert!(!trace.contains("sensitive-private-location"));
Ok(())
}
#[test]
fn genuine_media_tags_search_duration_filter_and_lifecycle() -> anyhow::Result<()> {
let directory = tempfile::tempdir()?;
let root = directory.path().join("music");
std::fs::create_dir(&root)?;
let file = root.join("track.ogg");
std::fs::write(&file, b"storage fixture")?;
let db = directory.path().join("index.db");
let mut index = super::SearchIndex::open(&db)?;
index.index_folder(&root)?;
let path = file.canonicalize()?.to_string_lossy().into_owned();
let (size, mtime): (u64, i64) = index.connection.query_row(
"SELECT size,mtime FROM files WHERE path=?1",
[&path],
|r| Ok((r.get(0)?, r.get(1)?)),
)?;
index.set_media_metadata(
&path,
"artist: Kimiko Ishizaka\ntitle: Variatio 13 a 2 Clav.",
180000,
size,
mtime,
)?;
let results = index.search("Kimiko Ishizaka", 10)?;
assert_eq!(results.len(), 1);
assert_eq!(results[0].match_kind, "media_metadata");
assert!(results[0].snippet.contains("Kimiko"));
let too_short = super::SearchOptions {
max_duration_ms: Some(100000),
..Default::default()
};
assert!(index
.search_with_options("track.ogg", 10, &too_short)?
.is_empty());
let correct = super::SearchOptions {
min_duration_ms: Some(170000),
max_duration_ms: Some(190000),
..Default::default()
};
assert_eq!(index.search_with_options("Kimiko", 10, &correct)?.len(), 1);
drop(index);
let mut index = super::SearchIndex::open(&db)?;
assert_eq!(index.search("Variatio", 10)?.len(), 1);
std::fs::write(&file, b"changed file bytes")?;
index.index_folder(&root)?;
assert!(index.search("Kimiko", 10)?.is_empty());
assert!(index
.set_media_metadata(&path, "stale", 180000, size, mtime)
.is_err());
std::fs::remove_file(&file)?;
index.index_folder(&root)?;
assert!(index.search("Variatio", 10)?.is_empty());
Ok(())
}
use super::*;
#[test]
fn persistence_updates_deletes_and_exact_priority() -> Result<()> {
+292
View File
@@ -0,0 +1,292 @@
//! Monotonic request metrics. Stage names are static labels, never queries/paths.
use serde::{Deserialize, Serialize};
use std::{
collections::BTreeMap,
sync::atomic::{AtomicU64, Ordering},
time::Instant,
};
const MAX_SPANS: usize = 512;
static COLLECTOR_ID: AtomicU64 = AtomicU64::new(1);
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResponse<T> {
pub results: Vec<T>,
pub metrics: SearchMetrics,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchMetrics {
pub total_ms: f64,
pub result_count: usize,
pub trace: Option<Vec<SearchSpan>>,
#[serde(default)]
pub dropped_spans: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchSpan {
pub name: String,
pub start_ms: f64,
pub duration_ms: f64,
pub parent: Option<usize>,
pub counters: BTreeMap<String, u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MetricsMode {
Basic,
Debug,
}
#[derive(Debug, Clone, Copy)]
pub struct SpanToken {
collector: u64,
index: usize,
}
impl SpanToken {
pub fn index(self) -> usize {
self.index
}
}
struct OpenSpan {
value: SearchSpan,
started: Instant,
ended: bool,
}
struct DebugTrace {
id: u64,
spans: Vec<OpenSpan>,
dropped: u64,
}
/// Basic mode stores only one Instant: no span Vec or stage/counter allocation.
pub struct TraceCollector {
started: Instant,
debug: Option<DebugTrace>,
}
impl TraceCollector {
pub fn new(mode: MetricsMode) -> Self {
Self {
started: Instant::now(),
debug: match mode {
MetricsMode::Basic => None,
MetricsMode::Debug => Some(DebugTrace {
id: COLLECTOR_ID.fetch_add(1, Ordering::Relaxed),
spans: Vec::new(),
dropped: 0,
}),
},
}
}
pub fn elapsed_ms(&self) -> f64 {
self.started.elapsed().as_secs_f64() * 1000.0
}
pub fn is_debug(&self) -> bool {
self.debug.is_some()
}
pub fn begin(&mut self, name: &'static str, parent: Option<SpanToken>) -> Option<SpanToken> {
let debug = self.debug.as_mut()?;
let parent_index = match parent {
Some(token)
if token.collector == debug.id
&& debug.spans.get(token.index).is_some_and(|span| !span.ended) =>
{
Some(token.index)
}
Some(_) => {
debug.dropped = debug.dropped.saturating_add(1);
return None;
}
None => None,
};
if debug.spans.len() >= MAX_SPANS {
debug.dropped = debug.dropped.saturating_add(1);
return None;
}
let started = Instant::now();
let index = debug.spans.len();
debug.spans.push(OpenSpan {
value: SearchSpan {
name: name.chars().take(96).collect(),
start_ms: started.duration_since(self.started).as_secs_f64() * 1000.0,
duration_ms: 0.0,
parent: parent_index,
counters: BTreeMap::new(),
},
started,
ended: false,
});
Some(SpanToken {
collector: debug.id,
index,
})
}
/// Invalid/previously ended tokens are harmless. Closing a parent also
/// closes open descendants at the same instant, preserving containment.
pub fn end(&mut self, token: Option<SpanToken>, counters: BTreeMap<String, u64>) {
let Some(token) = token else {
return;
};
let Some(debug) = self.debug.as_mut() else {
return;
};
if token.collector != debug.id
|| token.index >= debug.spans.len()
|| debug.spans[token.index].ended
{
return;
}
let now = Instant::now();
for index in token.index..debug.spans.len() {
let mut ancestor = Some(index);
let mut contained = false;
while let Some(current) = ancestor {
if current == token.index {
contained = true;
break;
}
ancestor = debug.spans[current].value.parent;
}
if contained && !debug.spans[index].ended {
let span = &mut debug.spans[index];
span.value.duration_ms = now.duration_since(span.started).as_secs_f64() * 1000.0;
span.ended = true;
}
}
// At most32 distinct bounded keys; callers should use static diagnostic
// names and counts, never source text or paths.
let stored = &mut debug.spans[token.index].value.counters;
for (key, value) in counters.into_iter().take(32) {
let key: String = key.chars().take(96).collect();
if stored.len() < 32 || stored.contains_key(&key) {
stored.insert(key, value);
}
}
}
pub fn counter(&mut self, token: Option<SpanToken>, name: &'static str, value: u64) {
let Some(token) = token else {
return;
};
let Some(debug) = self.debug.as_mut() else {
return;
};
if token.collector != debug.id {
return;
}
let Some(span) = debug.spans.get_mut(token.index) else {
return;
};
if span.value.counters.len() < 32 || span.value.counters.contains_key(name) {
span.value
.counters
.insert(name.chars().take(96).collect(), value);
}
}
/// Return the closure's value unchanged, including any Result error.
/// The closure receives its parent token for safe nested instrumentation.
pub fn timed<R>(
&mut self,
name: &'static str,
parent: Option<SpanToken>,
run: impl FnOnce(&mut Self, Option<SpanToken>) -> R,
) -> R {
let token = self.begin(name, parent);
let value = run(self, token);
self.end(token, BTreeMap::new());
value
}
pub fn finish(mut self, result_count: usize) -> SearchMetrics {
let now = Instant::now();
if let Some(debug) = self.debug.as_mut() {
for span in &mut debug.spans {
if !span.ended {
span.value.duration_ms =
now.duration_since(span.started).as_secs_f64() * 1000.0;
span.ended = true;
}
}
}
let (trace, dropped_spans) = match self.debug {
Some(debug) => (
Some(debug.spans.into_iter().map(|span| span.value).collect()),
debug.dropped,
),
None => (None, 0),
};
SearchMetrics {
total_ms: now.duration_since(self.started).as_secs_f64() * 1000.0,
result_count,
trace,
dropped_spans,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn nesting_contains_children_and_preserves_error() {
let mut trace = TraceCollector::new(MetricsMode::Debug);
let result: Result<(), &str> = trace.timed("request", None, |trace, parent| {
trace.timed("index", parent, |trace, child| {
trace.counter(child, "rows", 7);
Err("failure")
})
});
assert_eq!(result, Err("failure"));
let metrics = trace.finish(0);
let spans = metrics.trace.unwrap();
assert_eq!(spans[1].parent, Some(0));
assert_eq!(spans[1].counters["rows"], 7);
assert!(spans[1].start_ms >= spans[0].start_ms);
assert!(
spans[1].start_ms + spans[1].duration_ms
<= spans[0].start_ms + spans[0].duration_ms + 1e-6
);
assert!(metrics.total_ms >= spans[0].start_ms + spans[0].duration_ms - 1e-6);
}
#[test]
fn basic_disables_stage_storage_and_preserves_counts() {
let mut trace = TraceCollector::new(MetricsMode::Basic);
assert!(!trace.is_debug());
assert!(trace.begin("unused", None).is_none());
assert_eq!(trace.timed("unused", None, |_, _| 42), 42);
let metrics = trace.finish(3);
assert_eq!(metrics.result_count, 3);
assert!(metrics.trace.is_none());
assert_eq!(metrics.dropped_spans, 0);
}
#[test]
fn trace_is_capped_and_reports_dropped_spans() {
let mut trace = TraceCollector::new(MetricsMode::Debug);
for index in 0..520 {
let token = trace.begin("stage", None);
if index == 0 {
trace.counter(token, "existing", 99);
}
let counters = if index == 0 {
(0..40).map(|key| (format!("counter_{key}"), key)).collect()
} else {
BTreeMap::new()
};
trace.end(token, counters);
}
let metrics = trace.finish(2);
let spans = metrics.trace.unwrap();
assert_eq!(spans.len(), 512);
assert_eq!(spans[0].counters.len(), 32);
assert_eq!(spans[0].counters["existing"], 99);
assert_eq!(metrics.dropped_spans, 8);
}
#[test]
fn foreign_tokens_and_early_parent_close_are_safe() {
let mut one = TraceCollector::new(MetricsMode::Debug);
let mut two = TraceCollector::new(MetricsMode::Debug);
let root = one.begin("root", None);
let child = one.begin("child", root);
assert!(two.begin("foreign", root).is_none());
one.end(root, BTreeMap::new());
one.end(child, BTreeMap::new());
let spans = one.finish(0).trace.unwrap();
assert!(
spans[1].start_ms + spans[1].duration_ms
<= spans[0].start_ms + spans[0].duration_ms + 1e-6
);
}
}
+14
View File
@@ -0,0 +1,14 @@
[package]
name = "searchi-media"
version.workspace = true
edition.workspace = true
license.workspace = true
description = "Bounded local CPU media probing and timestamp-preserving sampling"
[dependencies]
anyhow.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2 = "0.10"
tempfile = "3"
image = { version = "0.25", default-features = false, features = ["png", "jpeg", "webp", "bmp"] }
+70
View File
@@ -0,0 +1,70 @@
# Local media preparation
`searchi-media` prepares actual files for the Google multimodal embedder. It
does not download programs or models and has no UI dependency. Callers supply
verified packaged FFmpeg/ffprobe paths through `FfmpegTools::new`, or explicitly
set both `SEARCHI_FFMPEG_PATH` and `SEARCHI_FFPROBE_PATH`.
FFmpeg handles maintained codecs and container timestamps; serde parses its
bounded JSON, tempfile owns frame cleanup, and sha2 validates file identity.
The embedding crate uses the Rust image library for bounded raster decoding.
USearch supplies the service's persisted HNSW index; no codec or vector index
is reimplemented here.
Every FFmpeg/ffprobe media input restricts protocols to `file,pipe`. A selected
local playlist cannot fetch HTTP/HTTPS/RTSP references. A real loopback-listener
regression verifies probing and both decoding paths never connect. This is a
protocol restriction, not a filesystem sandbox: explicitly selected files on
network-mounted drives still use their operating system's filesystem access.
Windows VM tests use CPU only. Decoder commands specify one thread and
`-hwaccel none`; subprocess windows are hidden. Portable argument arrays avoid
shell interpretation. Production hardware decode and reliable CPU fallback
remain unimplemented; this crate does not silently enable accelerators.
Supported images are PNG, JPEG, WebP and BMP (actual bounded embedding-library
decode, including EXIF orientation). AVIF, HEIC, TIFF, animated GIF and embedded
album artwork are not supported yet. Movie containers include MP4, MKV, MOV,
AVI, WebM, M4V, MPEG/MPG; audio includes MP3, FLAC, WAV, M4A, Ogg, Opus, AAC,
AIFF. Actual codec support comes from the packaged FFmpeg build. Unknown or
corrupt files remain filename-searchable with explicit indexing errors.
Default limits: 16 GiB media input; four-hour known duration; 16 million source
video pixels; 45 seconds per decoder command; 64 KiB ffprobe JSON; 256 KiB
frame diagnostic output; 8 MiB encoded PNG frame. Images are separately limited
to 64 MiB input. Captured pipe output is bounded and timed processes are killed.
Process address-space sandboxing remains a release gate; these limits do not
guarantee a codec's internal allocations fit a fixed RAM ceiling.
Video samples span the whole timeline at two-second intervals, becoming coarser
when the default 240 distributed-frame budget is reached. Up to 120 additional
scene-cut frames are captured within 1.5-second windows around sample anchors.
This is **partial visual coverage**, not exhaustive scene detection. Actual
decoded presentation timestamps are retained relative to container start. The
primary non-cover-art video stream is selected. Full images use model budget
280; video frames use 70. No fabricated captions are stored.
The primary audio stream is decoded to mono 16 kHz PCM, embedded in 30-second
windows with two seconds of overlap. The complete stream is covered or the
600-window budget returns an explicit failure. Standalone music uses the same
real audio embedding; artist/title/album/genre tags are merged from genuine
stream and format metadata and stored separately for lexical search. Audio
embeddings are not transcripts; exact spoken-word retrieval is not provided.
Alternative language tracks and isolated channels remain a release gate.
The service publishes text, image, video and audio vectors as one generation.
It persists origin and page/timestamp locations, policy/model fingerprints,
full file SHA at indexing, and size/mtime guards. Corrupt-file failures discard
that file's partial vectors and persist reasons. Model inference failure aborts
the generation. At query time only actual media candidate hits are stat-checked;
no corpus-wide hash/stat scan occurs. Changes preserving both size and mtime
cannot be detected by this fast guard; periodic content-hash reconciliation is
a release gate. Builds currently rebuild a complete generation, with a
100,000-vector / 10,000-media-file bound; unchanged-vector reuse is future work.
Offline tests cover complete audio windows, frame budgets and identity changes.
The ignored `real_decoder_preserves_frame_timestamps_and_audio` test has been
executed successfully in the CPU Windows VM with provisioned FFmpeg 9.0.2.
It creates genuine FFV1/PCM media, validates decoded two-second PTS and actual
mono sine samples. Linux/macOS execution and installer codec/license delivery
have not been validated.
+789
View File
@@ -0,0 +1,789 @@
//! Local bounded media preparation. No network, model, executable installer, or
//! fabricated caption/transcript lives here. FFmpeg paths come from the caller.
//!
//! Video coverage is distributed sampling, not a guarantee that every scene is
//! represented. Audio windows cover the complete primary stream or return an
//! explicit budget error. Decoder process memory is not sandbox-limited yet.
use anyhow::{bail, ensure, Context, Result};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{
collections::BTreeMap,
fs::{self, File},
io::Read,
path::{Path, PathBuf},
process::{Command, Stdio},
thread,
time::{Duration, Instant, UNIX_EPOCH},
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum MediaKind {
Image,
Video,
Audio,
}
#[cfg(test)]
mod tests {
use super::*;
fn fixture(duration_ms: u32) -> MediaProbe {
MediaProbe {
duration_ms,
start_ms: 0,
video_stream_index: 0,
width: 64,
height: 64,
has_video: true,
has_audio: true,
audio_start_ms: 0,
audio_end_ms: duration_ms,
tags: BTreeMap::new(),
}
}
#[test]
fn audio_covers_tail_with_overlap_and_rejects_budget() -> Result<()> {
let probe = fixture(65_000);
let limits = MediaLimits::default();
let windows = audio_windows(&probe, &limits)?;
assert_eq!(
windows.iter().map(|w| w.start_ms).collect::<Vec<_>>(),
vec![0, 28_000, 56_000]
);
assert_eq!(windows.last().unwrap().duration_ms, 9_000);
let tight = MediaLimits {
max_audio_windows: 2,
..limits
};
assert!(audio_windows(&probe, &tight).is_err());
Ok(())
}
#[test]
fn finer_audio_windows_cover_long_movie_or_explicitly_reject() -> Result<()> {
let probe = fixture(3 * 60 * 60 * 1000);
let limits = MediaLimits {
audio_window_ms: 10_000,
audio_overlap_ms: 2000,
max_audio_windows: 1800,
..Default::default()
};
let windows = audio_windows(&probe, &limits)?;
assert_eq!(windows.first().unwrap().start_ms, 0);
let last = windows.last().unwrap();
assert_eq!(last.start_ms + last.duration_ms, probe.duration_ms);
assert!(windows.len() >= 1300 && windows.len() <= 1440);
let too_small = MediaLimits {
max_audio_windows: 1000,
..limits
};
assert!(audio_windows(&probe, &too_small).is_err());
Ok(())
}
#[test]
fn frame_budget_spans_entire_movie() -> Result<()> {
let probe = fixture(3_600_000);
let times = frame_times(&probe, &MediaLimits::default())?;
assert!(times.len() <= 240);
assert_eq!(times.first(), Some(&0));
assert_eq!(times.last(), Some(&3_599_500));
Ok(())
}
#[test]
fn identity_rejects_changes_and_oversized_input() -> Result<()> {
let temp = tempfile::tempdir()?;
let file = temp.path().join("audio.wav");
fs::write(&file, b"first")?;
assert!(MediaIdentity::read(&file, 4).is_err());
let identity = MediaIdentity::read(&file, 10)?;
fs::write(&file, b"changed size")?;
assert!(identity.ensure_unchanged().is_err());
Ok(())
}
#[test]
#[ignore = "Requires explicitly provisioned CPU FFmpeg/ffprobe executables"]
fn decoder_rejects_network_references_inside_local_media() -> Result<()> {
let tools = FfmpegTools::from_environment()?.context("Set CPU decoder paths")?;
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
listener.set_nonblocking(true)?;
let temp = tempfile::tempdir()?;
let file = temp.path().join("local.m3u8");
fs::write(&file,format!("#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-TARGETDURATION:2\n#EXTINF:2.0,\nhttp://{}/segment.ts\n#EXT-X-ENDLIST\n",listener.local_addr()?))?;
let limits = MediaLimits {
command_timeout: Duration::from_secs(3),
..Default::default()
};
let error = super::probe(&tools, &file, &limits).unwrap_err();
assert!(format!("{error:#}").contains("whitelist"), "{error:#}");
assert!(decode_frames(&tools, &file, &fixture(2000), 0, false, &limits).is_err());
assert!(decode_audio(
&tools,
&file,
AudioWindow {
start_ms: 0,
duration_ms: 2000
},
&limits
)
.is_err());
match listener.accept() {
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
other => panic!("Decoder attempted network connection: {other:?}"),
}
Ok(())
}
#[test]
#[ignore = "Requires explicitly provisioned CPU FFmpeg/ffprobe executables"]
fn real_decoder_preserves_frame_timestamps_and_audio() -> Result<()> {
let tools = FfmpegTools::from_environment()?.context("Set CPU decoder paths")?;
let temp = tempfile::tempdir()?;
let video = temp.path().join("fixture.mkv");
let mut command = Command::new(&tools.ffmpeg);
command
.args([
"-nostdin",
"-v",
"error",
"-threads",
"1",
"-f",
"lavfi",
"-i",
"color=c=red:s=64x64:r=10:d=4",
"-f",
"lavfi",
"-i",
"sine=frequency=440:sample_rate=16000:duration=4",
"-c:v",
"ffv1",
"-c:a",
"pcm_s16le",
"-threads",
"1",
"-y",
])
.arg(&video);
run(command, Duration::from_secs(30), 4096, 65536)?;
let limits = MediaLimits::default();
let probe = probe(&tools, &video, &limits)?;
assert!(probe.has_video && probe.has_audio);
let frames = decode_frames(&tools, &video, &probe, 2000, false, &limits)?;
assert_eq!(frames.frames.len(), 1);
assert!((1900..=2100).contains(&frames.frames[0].timestamp_ms));
let audio = decode_audio(
&tools,
&video,
AudioWindow {
start_ms: 1000,
duration_ms: 2000,
},
&limits,
)?;
assert!((31_000..=32_000).contains(&audio.len()));
assert!(audio.iter().any(|x| x.abs() > 0.01));
assert!(audio.iter().all(|x| x.is_finite() && x.abs() <= 1.0));
let ogg = temp.path().join("overshoot.ogg");
let mut command = Command::new(&tools.ffmpeg);
command
.args([
"-nostdin",
"-v",
"error",
"-threads",
"1",
"-f",
"lavfi",
"-i",
"sine=frequency=440:sample_rate=44100:duration=31.123",
"-c:a",
"libvorbis",
"-metadata",
"artist=CPU decoder fixture",
"-threads",
"1",
"-y",
])
.arg(&ogg);
run(command, Duration::from_secs(30), 4096, 65536)?;
let music_probe = super::probe(&tools, &ogg, &limits)?;
assert_eq!(
music_probe.tags.get("artist").map(String::as_str),
Some("CPU decoder fixture")
);
let pcm = decode_audio(
&tools,
&ogg,
AudioWindow {
start_ms: 0,
duration_ms: 30000,
},
&limits,
)?;
assert_eq!(
pcm.len(),
480000,
"Vorbis packets must not overshoot model limit"
);
Ok(())
}
}
pub fn classify(path: &Path) -> Option<MediaKind> {
match path.extension()?.to_str()?.to_ascii_lowercase().as_str() {
"png" | "jpg" | "jpeg" | "webp" | "bmp" => Some(MediaKind::Image),
"mp4" | "mkv" | "mov" | "avi" | "webm" | "m4v" | "mpeg" | "mpg" => Some(MediaKind::Video),
"mp3" | "flac" | "wav" | "m4a" | "ogg" | "opus" | "aac" | "aiff" | "aif" => {
Some(MediaKind::Audio)
}
_ => None,
}
}
/// Validate bounded raster decoding before treating model errors as fatal.
/// The provider still performs its own bounds and EXIF-aware preprocessing.
pub fn validate_image(path: &Path) -> Result<()> {
ensure!(
fs::metadata(path)?.len() <= 64 * 1024 * 1024,
"Image exceeds 64 MiB"
);
let reader = image::ImageReader::open(path)?.with_guessed_format()?;
let (width, height) = reader.into_dimensions()?;
ensure!(
width <= 16384 && height <= 16384 && u64::from(width) * u64::from(height) <= 16_000_000,
"Image exceeds 16M pixels"
);
let mut reader = image::ImageReader::open(path)?.with_guessed_format()?;
let mut limits = image::Limits::default();
limits.max_alloc = Some(256 * 1024 * 1024);
limits.max_image_width = Some(16384);
limits.max_image_height = Some(16384);
reader.limits(limits);
use image::ImageDecoder;
let mut decoder = reader.into_decoder()?;
decoder
.orientation()
.context("Invalid image EXIF orientation")?;
image::DynamicImage::from_decoder(decoder).context("Image decoder rejected corrupt content")?;
Ok(())
}
#[derive(Debug, Clone)]
pub struct FfmpegTools {
pub ffmpeg: PathBuf,
pub ffprobe: PathBuf,
}
impl FfmpegTools {
pub fn new(ffmpeg: impl AsRef<Path>, ffprobe: impl AsRef<Path>) -> Result<Self> {
ensure!(
ffmpeg.as_ref().is_file() && ffprobe.as_ref().is_file(),
"FFmpeg/ffprobe executable is missing"
);
Ok(Self {
ffmpeg: fs::canonicalize(ffmpeg)?,
ffprobe: fs::canonicalize(ffprobe)?,
})
}
pub fn from_environment() -> Result<Option<Self>> {
match (
std::env::var_os("SEARCHI_FFMPEG_PATH"),
std::env::var_os("SEARCHI_FFPROBE_PATH"),
) {
(None, None) => Ok(None),
(Some(ffmpeg), Some(ffprobe)) => Self::new(ffmpeg, ffprobe).map(Some),
_ => bail!("Set both SEARCHI_FFMPEG_PATH and SEARCHI_FFPROBE_PATH"),
}
}
}
#[derive(Debug, Clone)]
pub struct MediaLimits {
pub max_file_bytes: u64,
pub max_duration_ms: u32,
pub max_frames: usize,
pub frame_interval_ms: u32,
pub max_scene_frames: usize,
pub audio_window_ms: u32,
pub audio_overlap_ms: u32,
pub max_audio_windows: usize,
pub command_timeout: Duration,
}
impl Default for MediaLimits {
fn default() -> Self {
Self {
max_file_bytes: 16 * 1024 * 1024 * 1024,
max_duration_ms: 4 * 60 * 60 * 1000,
max_frames: 240,
frame_interval_ms: 2000,
max_scene_frames: 120,
audio_window_ms: 30_000,
audio_overlap_ms: 2000,
max_audio_windows: 600,
command_timeout: Duration::from_secs(45),
}
}
}
#[derive(Debug, Clone, Serialize)]
pub struct MediaIdentity {
pub path: PathBuf,
pub bytes: u64,
pub modified_ns: i64,
pub sha256: String,
}
impl MediaIdentity {
pub fn read(path: &Path, maximum: u64) -> Result<Self> {
let file = File::open(path)?;
let before = file.metadata()?;
ensure!(
before.len() <= maximum,
"Media file exceeds {maximum} bytes"
);
let modified_ns = modified_ns(&before)?;
let mut reader = file.take(maximum.saturating_add(1));
let mut buffer = [0u8; 65536];
let mut hash = Sha256::new();
let mut count = 0u64;
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
count += read as u64;
ensure!(count <= maximum, "Media grew past file budget");
hash.update(&buffer[..read]);
}
let identity = Self {
path: fs::canonicalize(path)?,
bytes: count,
modified_ns,
sha256: format!("{:x}", hash.finalize()),
};
ensure!(count == before.len(), "Media changed while fingerprinting");
identity.ensure_unchanged()?;
Ok(identity)
}
pub fn ensure_unchanged(&self) -> Result<()> {
let metadata = fs::metadata(&self.path)?;
ensure!(
metadata.len() == self.bytes && modified_ns(&metadata)? == self.modified_ns,
"Media changed during indexing"
);
Ok(())
}
}
fn modified_ns(metadata: &fs::Metadata) -> Result<i64> {
i64::try_from(metadata.modified()?.duration_since(UNIX_EPOCH)?.as_nanos())
.context("Unsupported media timestamp")
}
#[derive(Debug, Clone, Serialize)]
pub struct MediaProbe {
pub duration_ms: u32,
pub start_ms: i64,
pub video_stream_index: u32,
pub width: u32,
pub height: u32,
pub has_video: bool,
pub has_audio: bool,
pub audio_start_ms: u32,
pub audio_end_ms: u32,
pub tags: BTreeMap<String, String>,
}
#[derive(Deserialize)]
struct ProbeJson {
#[serde(default)]
streams: Vec<ProbeStream>,
format: ProbeFormat,
}
#[derive(Deserialize)]
struct ProbeFormat {
duration: String,
#[serde(default)]
start_time: Option<String>,
#[serde(default)]
tags: BTreeMap<String, String>,
}
#[derive(Deserialize)]
struct ProbeStream {
index: u32,
codec_type: String,
#[serde(default)]
width: u32,
#[serde(default)]
height: u32,
#[serde(default)]
duration: Option<String>,
#[serde(default)]
start_time: Option<String>,
#[serde(default)]
tags: BTreeMap<String, String>,
#[serde(default)]
disposition: BTreeMap<String, i64>,
}
fn milliseconds(value: &str) -> Result<i64> {
let value: f64 = value.parse()?;
ensure!(
value.is_finite() && value.abs() < 100_000_000.0,
"Invalid media duration/timestamp"
);
Ok((value * 1000.0).round() as i64)
}
pub fn probe(tools: &FfmpegTools, path: &Path, limits: &MediaLimits) -> Result<MediaProbe> {
let mut command = Command::new(&tools.ffprobe);
command.args(["-protocol_whitelist", "file,pipe"]);
command
.args([
"-v",
"error",
"-threads",
"1",
"-probesize",
"5000000",
"-analyzeduration",
"5000000",
"-show_streams",
"-show_format",
"-of",
"json",
])
.arg(path);
let output = run(command, limits.command_timeout, 65536, 65536)?;
let parsed: ProbeJson =
serde_json::from_slice(&output.stdout).context("Invalid ffprobe JSON")?;
let duration = milliseconds(&parsed.format.duration)?;
ensure!(
duration > 0 && duration <= i64::from(limits.max_duration_ms),
"Media duration exceeds policy or is unknown"
);
let start = parsed
.format
.start_time
.as_deref()
.map(milliseconds)
.transpose()?
.unwrap_or(0);
let video = parsed.streams.iter().find(|stream| {
stream.codec_type == "video" && stream.disposition.get("attached_pic") != Some(&1)
});
if let Some(video) = video {
ensure!(
u64::from(video.width) * u64::from(video.height) <= 16_000_000,
"Video dimensions exceed 16M pixels"
);
}
let audio = parsed
.streams
.iter()
.find(|stream| stream.codec_type == "audio");
let audio_start = audio
.and_then(|stream| stream.start_time.as_deref())
.map(milliseconds)
.transpose()?
.unwrap_or(start)
- start;
let audio_length = audio
.and_then(|stream| stream.duration.as_deref())
.map(milliseconds)
.transpose()?
.unwrap_or(duration);
let mut tags = parsed.format.tags;
for stream in &parsed.streams {
for (key, value) in &stream.tags {
tags.entry(key.clone()).or_insert_with(|| value.clone());
}
}
tags = tags
.into_iter()
.map(|(key, value)| {
(
key.to_ascii_lowercase(),
value.chars().take(2048).collect::<String>(),
)
})
.filter(|(key, _)| {
matches!(
key.as_str(),
"title"
| "artist"
| "album"
| "album_artist"
| "genre"
| "comment"
| "date"
| "composer"
| "description"
| "track"
)
})
.take(16)
.collect();
Ok(MediaProbe {
duration_ms: duration as u32,
start_ms: start,
video_stream_index: video.map_or(0, |stream| stream.index),
width: video.map_or(0, |stream| stream.width),
height: video.map_or(0, |stream| stream.height),
has_video: video.is_some(),
has_audio: audio.is_some(),
audio_start_ms: audio_start.max(0) as u32,
audio_end_ms: (audio_start + audio_length).clamp(0, duration) as u32,
tags,
})
}
/// Uniform coverage is distributed over the entire movie; this is not a prefix
/// truncation. Scene probes may add extra local cuts around these anchors.
pub fn frame_times(probe: &MediaProbe, limits: &MediaLimits) -> Result<Vec<u32>> {
if !probe.has_video {
return Ok(Vec::new());
}
ensure!(
limits.max_frames >= 2 && limits.max_frames <= 2000,
"Frame budget must be 2..2000"
);
let last = probe.duration_ms.saturating_sub(500);
let interval = limits
.frame_interval_ms
.max(last.div_ceil((limits.max_frames - 1) as u32))
.max(1);
let mut times = Vec::new();
let mut time = 0;
while time < last {
times.push(time);
time = time.saturating_add(interval);
}
times.push(last);
ensure!(
times.len() <= limits.max_frames,
"Frame sampling budget exceeded"
);
Ok(times)
}
#[derive(Debug, Clone, Copy, Serialize)]
pub struct AudioWindow {
pub start_ms: u32,
pub duration_ms: u32,
}
pub fn audio_windows(probe: &MediaProbe, limits: &MediaLimits) -> Result<Vec<AudioWindow>> {
if !probe.has_audio {
return Ok(Vec::new());
}
ensure!(
limits.audio_window_ms >= 5_000
&& limits.audio_window_ms <= 30_000
&& limits.audio_overlap_ms < limits.audio_window_ms,
"Invalid audio window policy"
);
let step = limits.audio_window_ms - limits.audio_overlap_ms;
let mut windows = Vec::new();
let mut start = probe.audio_start_ms;
while start < probe.audio_end_ms {
ensure!(
windows.len() < limits.max_audio_windows,
"Whole audio coverage exceeds window budget; no partial index produced"
);
windows.push(AudioWindow {
start_ms: start,
duration_ms: limits.audio_window_ms.min(probe.audio_end_ms - start),
});
if start + limits.audio_window_ms >= probe.audio_end_ms {
break;
}
start = start.saturating_add(step);
}
Ok(windows)
}
pub struct DecodedFrame {
pub path: PathBuf,
pub timestamp_ms: u32,
}
pub struct FrameBatch {
directory: tempfile::TempDir,
pub frames: Vec<DecodedFrame>,
}
impl FrameBatch {
pub fn directory(&self) -> &Path {
self.directory.path()
}
}
pub fn decode_frames(
tools: &FfmpegTools,
path: &Path,
probe: &MediaProbe,
start_ms: u32,
include_scene: bool,
limits: &MediaLimits,
) -> Result<FrameBatch> {
let directory = tempfile::tempdir()?;
let mut command = Command::new(&tools.ffmpeg);
command.args(["-protocol_whitelist", "file,pipe"]);
command.args(["-nostdin","-hide_banner","-loglevel","info","-threads","1","-hwaccel","none","-copyts","-ss"])
.arg(format!("{:.3}",f64::from(start_ms)/1000.0)).args(["-t",if include_scene {"1.5"} else {"0.5"},"-i"]).arg(path)
.arg("-map").arg(format!("0:{}",probe.video_stream_index)).args(["-an","-filter_threads","1","-vf"])
.arg(if include_scene {"select=eq(n\\,0)+gt(scene\\,0.32),showinfo,scale=960:960:force_original_aspect_ratio=decrease"}
else {"showinfo,scale=960:960:force_original_aspect_ratio=decrease"})
.args(["-frames:v",if include_scene {"2"} else {"1"},"-fps_mode","vfr","-threads","1","-y"])
.arg(directory.path().join("frame-%02d.png"));
let output = run(command, limits.command_timeout, 4096, 256 * 1024)?;
let mut timestamps = Vec::new();
for line in String::from_utf8_lossy(&output.stderr).lines() {
if line.contains("showinfo") {
if let Some((_, tail)) = line.split_once("pts_time:") {
let value = tail
.split_whitespace()
.next()
.context("Missing frame PTS")?;
let actual = (milliseconds(value)? - probe.start_ms).max(0);
ensure!(
actual <= i64::from(probe.duration_ms) + 1000,
"Frame PTS lies beyond media timeline"
);
timestamps.push(actual as u32);
}
}
}
let mut frames = Vec::new();
for ordinal in 1..=if include_scene { 2 } else { 1 } {
let file = directory.path().join(format!("frame-{ordinal:02}.png"));
if file.is_file() {
ensure!(
fs::metadata(&file)?.len() <= 8 * 1024 * 1024,
"Decoded frame exceeds 8MiB"
);
let timestamp_ms = *timestamps
.get(ordinal - 1)
.context("Decoder did not report genuine frame PTS")?;
frames.push(DecodedFrame {
path: file,
timestamp_ms,
});
}
}
ensure!(
!frames.is_empty(),
"Decoder produced no frame at requested timeline position"
);
Ok(FrameBatch { directory, frames })
}
pub fn decode_audio(
tools: &FfmpegTools,
path: &Path,
window: AudioWindow,
limits: &MediaLimits,
) -> Result<Vec<f32>> {
ensure!(
window.duration_ms > 0 && window.duration_ms <= 30_000,
"Audio decode window exceeds 30 seconds"
);
let mut command = Command::new(&tools.ffmpeg);
command.args(["-protocol_whitelist", "file,pipe"]);
command
.args([
"-nostdin",
"-hide_banner",
"-v",
"error",
"-threads",
"1",
"-ss",
])
.arg(format!("{:.3}", f64::from(window.start_ms) / 1000.0))
.args(["-t"])
.arg(format!("{:.3}", f64::from(window.duration_ms) / 1000.0))
.arg("-i")
.arg(path)
.arg("-af")
.arg(format!(
"aresample=16000,atrim=end_sample={},asetpts=PTS-STARTPTS",
u64::from(window.duration_ms) * 16
))
.args([
"-map", "0:a:0", "-vn", "-ac", "1", "-ar", "16000", "-threads", "1", "-f", "s16le",
"pipe:1",
]);
let output = run(command, limits.command_timeout, 960_000, 65536)?;
ensure!(
!output.stdout.is_empty() && output.stdout.len() % 2 == 0,
"Empty or invalid decoded PCM audio"
);
Ok(output
.stdout
.as_chunks::<2>()
.0
.iter()
.map(|bytes| f32::from(i16::from_le_bytes([bytes[0], bytes[1]])) / 32768.0)
.collect())
}
struct Output {
stdout: Vec<u8>,
stderr: Vec<u8>,
}
fn run(
mut command: Command,
timeout: Duration,
stdout_limit: usize,
stderr_limit: usize,
) -> Result<Output> {
#[cfg(windows)]
{
use std::os::windows::process::CommandExt;
command.creation_flags(0x08000000);
}
let mut child = command
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.context("Launch configured FFmpeg tool")?;
let stdout = child.stdout.take().context("Missing decoder stdout")?;
let stderr = child.stderr.take().context("Missing decoder stderr")?;
let read = |reader: Box<dyn Read + Send>, limit: usize| {
thread::spawn(move || -> Result<Vec<u8>> {
let mut bytes = Vec::new();
reader.take(limit as u64 + 1).read_to_end(&mut bytes)?;
ensure!(bytes.len() <= limit, "Decoder output exceeded byte budget");
Ok(bytes)
})
};
let out = read(Box::new(stdout), stdout_limit);
let err = read(Box::new(stderr), stderr_limit);
let start = Instant::now();
let status = loop {
if let Some(status) = child.try_wait()? {
break Some(status);
}
if start.elapsed() > timeout {
child.kill()?;
child.wait()?;
break None;
}
thread::sleep(Duration::from_millis(10));
};
let stdout = out
.join()
.map_err(|_| anyhow::anyhow!("Decoder stdout reader panicked"))??;
let stderr = err
.join()
.map_err(|_| anyhow::anyhow!("Decoder stderr reader panicked"))??;
let status = status.context("Decoder exceeded time budget")?;
ensure!(
status.success(),
"Decoder failed: {}",
String::from_utf8_lossy(&stderr)
.chars()
.take(2048)
.collect::<String>()
);
Ok(Output { stdout, stderr })
}
+3
View File
@@ -11,6 +11,9 @@ minilm-proof = ["dep:ort", "dep:tokenizers"]
[dependencies]
anyhow.workspace = true
base64 = "0.22"
image = { version = "0.25", default-features = false, features = ["png", "jpeg", "webp", "bmp"] }
jpeg-decoder = { version = "0.3.2", default-features = false }
ort = { version = "=2.0.0-rc.13", optional = true, default-features = false, features = ["std", "load-dynamic", "api-27"] }
sha2 = "0.10"
tokenizers = { version = "=0.22.2", optional = true, default-features = false, features = ["fancy-regex"] }
+44
View File
@@ -0,0 +1,44 @@
# Google worker image decoding
The production adapter verifies model/projector SHA256 and performs genuine
Google MTMD image/audio inference. Images are decoded with established Rust
libraries, not filename captions. `WorkerInfo.media_supported` comes from the
native handshake; absent fields on older text workers mean false. Media APIs
return an explicit runtime repair error for such workers while text remains
compatible.
JPEG sources up to 64 Mi pixels (including 48 MP camera photos), 16,384 pixels
per side, and 64 MiB compressed size are supported. For JPEG above 16 Mi pixels,
`image-rs/jpeg-decoder` 0.3.2 applies scaled IDCT before final model resizing.
Its default features are disabled, so this path adds no Rayon worker pool.
The source dimensions are checked before pixels are decoded. The secondary
header must match the first to detect file changes. RGB24 and grayscale camera
JPEGs are supported on this scaled path; unusual large CMYK/16-bit JPEGs fail
explicitly rather than silently taking an unbounded full-resolution fallback.
The published decoder offers 1/8, 1/4, 1/2, and full-scale decoding. It is in
maintenance mode; the current image 0.25 zune JPEG adapter exposes no scaled
IDCT interface. See the [decoder API](https://docs.rs/jpeg-decoder/0.3.2/jpeg_decoder/struct.Decoder.html)
and [upstream maintenance note](https://github.com/image-rs/jpeg-decoder).
Small JPEGs retain the existing image/zune path and pinned parity behavior.
PNG, WebP, and BMP remain capped at 16 Mi pixels with explicit diagnostics.
Both paths enforce a 256 MiB advertised decoding-buffer ceiling and apply
EXIF orientation. That ceiling is a buffer/output bound, not an absolute
process RSS guarantee: compressed data, codec working state, and resizing
have additional allocations. Extreme/progressive files need broader resource
profiling. The media fingerprint includes `largejpeg-scaled-v2`; generations
must additionally include visual token budget. Text fingerprints are unchanged.
An executed Windows CPU release test encodes and decodes an actual 8000x6000
JPEG with EXIF rotation, verifies orientation and region colors, and resizes it
to a 280-token grid. Decode plus resize took 176–220 ms and produced 672x912
RGB (1,838,592 bytes). A sampled whole-test peak was 146.86 MiB, including
144 MB of source fixture construction; it is not a decoder-only measurement
or a semantic quality benchmark. An 81 MP JPEG header is rejected before
pixel decode in the default tests.
```powershell
cargo test --release -p searchi-semantic large_phone_jpeg -- --ignored --nocapture
cargo test -p searchi-semantic --lib
```
@@ -41,11 +41,12 @@ set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "No runtime tools" FORCE)
set(LLAMA_BUILD_SERVER OFF CACHE BOOL "No HTTP server" FORCE)
set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "No examples" FORCE)
set(LLAMA_BUILD_TESTS OFF CACHE BOOL "No upstream tests in worker build" FORCE)
set(LLAMA_BUILD_MTMD OFF CACHE BOOL "Text worker; media encoder follows separately" FORCE)
set(LLAMA_BUILD_MTMD ON CACHE BOOL "Google image/audio projector" FORCE)
set(MTMD_VIDEO OFF CACHE BOOL "Frames decoded by bounded media adapter" FORCE)
set(LLAMA_OPENSSL OFF CACHE BOOL "No network runtime dependencies" FORCE)
add_subdirectory("${LLAMA_CPP_DIR}" llama-build)
add_executable(searchi-gemma-worker worker.cpp)
target_include_directories(searchi-gemma-worker PRIVATE "${LLAMA_CPP_DIR}/vendor")
target_link_libraries(searchi-gemma-worker PRIVATE llama)
target_link_libraries(searchi-gemma-worker PRIVATE llama mtmd)
target_compile_definitions(searchi-gemma-worker PRIVATE SEARCHI_LLAMA_TAG="b11461-no-logits-v1")
install(TARGETS searchi-gemma-worker RUNTIME DESTINATION bin)
+108 -6
View File
@@ -1,11 +1,14 @@
// Local NDJSON helper around the pinned llama.cpp encoder. stdout is exclusively
// protocol output; llama diagnostics go to stderr. No model download or HTTP.
#include "llama.h"
#include "mtmd.h"
#include "nlohmann/json.hpp"
#include <algorithm>
#include <cmath>
#include <chrono>
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <memory>
#include <stdexcept>
@@ -13,7 +16,7 @@
#include <vector>
using json = nlohmann::json;
constexpr size_t max_request_bytes = 2 * 1024 * 1024;
constexpr size_t max_request_bytes = 16 * 1024 * 1024;
constexpr size_t max_passage_bytes = 1024 * 1024;
constexpr size_t max_text_bytes = 65536;
constexpr size_t max_windows = 256;
@@ -34,12 +37,14 @@ class Encoder {
llama_context * context = nullptr;
llama_batch_ext * batch = nullptr;
const llama_vocab * vocab = nullptr;
mtmd_context * media = nullptr;
std::string model_path;
void release() {
if (media) mtmd_free(media);
if (batch) llama_batch_ext_free(batch);
if (context) llama_free(context);
if (model) llama_model_free(model);
batch = nullptr; context = nullptr; model = nullptr; vocab = nullptr;
batch = nullptr; context = nullptr; model = nullptr; vocab = nullptr; media = nullptr;
}
bool open(bool gpu) {
auto mp = llama_model_default_params();
@@ -68,6 +73,9 @@ public:
llama_log_set([](ggml_log_level level, const char * text, void *) {
if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
},nullptr);
mtmd_log_set([](ggml_log_level level, const char * text, void *) {
if (level >= GGML_LOG_LEVEL_WARN) std::fputs(text,stderr);
},nullptr);
llama_backend_init();
bool use_gpu = want_gpu && llama_supports_gpu_offload();
require(open(use_gpu) || (use_gpu && open(false)),"Cannot initialize local EmbeddingGemma 2 model/context");
@@ -77,6 +85,52 @@ public:
llama_backend_free();
}
bool healthy() const {return context != nullptr && batch != nullptr;}
void ensure_media(const std::string & projector) {
if (media) return;
auto params = mtmd_context_params_default();
params.use_gpu = device == "gpu"; params.n_threads = threads;
params.print_timings = false; params.warmup = false;
params.image_min_tokens = 1; params.image_max_tokens = 1120;
media = mtmd_init_from_file(projector.c_str(),model,params);
require_inference(media != nullptr,"Cannot initialize Google media projector");
}
std::vector<float> embed_bitmap(mtmd_bitmap * bitmap) {
std::unique_ptr<mtmd_input_chunks,decltype(&mtmd_input_chunks_free)> chunks(mtmd_input_chunks_init(),mtmd_input_chunks_free);
require_inference(chunks != nullptr,"Cannot allocate media chunks");
const char * marker = mtmd_default_marker();
mtmd_input_text text{marker,std::strlen(marker),true,true};
const mtmd_bitmap * bitmaps[] = {bitmap};
require_inference(mtmd_tokenize(media,chunks.get(),&text,bitmaps,1) == 0,"Media tokenization failed");
llama_batch_ext_clear(batch); size_t position = 0;
auto place = [&](int index) {
require_inference(index >= 0 && position < context_tokens,"Media batch exceeds context");
llama_pos pos[4] = {llama_pos(position++),0,0,0};
require_inference(llama_batch_ext_set_pos(batch,index,pos) && llama_batch_ext_set_output_embd(batch,index,true),"Cannot place media row");
};
for (size_t i=0;i<mtmd_input_chunks_size(chunks.get());++i) {
const auto * chunk = mtmd_input_chunks_get(chunks.get(),i);
if (mtmd_input_chunk_get_type(chunk) == MTMD_INPUT_CHUNK_TYPE_TEXT) {
size_t count=0; const auto * tokens=mtmd_input_chunk_get_tokens_text(chunk,&count);
for(size_t j=0;j<count;++j) place(llama_batch_ext_add_token(batch,0,tokens[j]));
} else {
require_inference(mtmd_encode_chunk(media,chunk)==0,"Google media encoder failed");
const float * rows=mtmd_get_output_embd(media);
require_inference(rows != nullptr,"No projected media rows");
size_t width=llama_model_n_embd_inp(model);
for(size_t j=0;j<mtmd_input_chunk_get_n_tokens(chunk);++j) {
llama_embd embd{rows+j*width,1,width};
place(llama_batch_ext_add_embd(batch,0,embd));
}
}
}
require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch)==0,"Media text encoder failed");
const float * pooled=llama_get_embeddings_seq(context,0);
require_inference(pooled != nullptr,"No pooled media vector");
double norm=0; for(int i=0;i<dimensions;++i){require_inference(std::isfinite(pooled[i]),"Nonfinite media vector");norm+=double(pooled[i])*pooled[i];}
require_inference(norm>0 && std::isfinite(norm),"Invalid media vector norm");
std::vector<float> vector(dimensions);for(int i=0;i<dimensions;++i)vector[i]=float(pooled[i]/std::sqrt(norm));
return vector;
}
void retry_cpu() {
require(device == "gpu","Only a GPU inference failure can request a CPU retry");
release();
@@ -103,10 +157,21 @@ public:
require(title.size() <= 512, "Document title exceeds 512 UTF-8 bytes");
return "title: " + (title.empty() ? "none" : title) + " | text: ";
}
std::vector<float> embed(const std::string & text, int cap) {
std::vector<float> embed(const std::string & text, int cap, json * timings = nullptr) {
using clock = std::chrono::steady_clock;
auto started = timings ? clock::now() : clock::time_point{};
auto mark = [&](const char * name) {
if (timings) {
auto now = clock::now();
(*timings)[name] = std::chrono::duration_cast<std::chrono::microseconds>(now-started).count();
started = now;
}
};
require(!text.empty() && text.size() <= max_text_bytes, "Embedding text must be nonempty and at most 64 KiB");
auto tokens = tokenize(text,true);
require(!tokens.empty() && tokens.size() <= size_t(cap), "Prompt exceeds model token budget; split document windows first");
if (timings) (*timings)["token_count"] = tokens.size();
mark("tokenization_us");
llama_batch_ext_clear(batch);
for (size_t i = 0; i < tokens.size(); ++i) {
int index = llama_batch_ext_add_token(batch,0,tokens[i]);
@@ -115,7 +180,9 @@ public:
require(llama_batch_ext_set_pos(batch,index,position), "Cannot set encoder position");
require(llama_batch_ext_set_output_embd(batch,index,true), "Cannot select embedding output");
}
mark("batch_prepare_us");
require_inference(llama_process(context,LLAMA_PROCESS_TYPE_ENCODE,batch) == 0, "llama_process encoder failed");
mark("inference_us");
const float * pooled = llama_get_embeddings_seq(context,0);
require_inference(pooled != nullptr, "Encoder returned no pooled vector");
double norm = 0;
@@ -123,6 +190,7 @@ public:
require_inference(norm > 0 && std::isfinite(norm), "Zero or nonfinite embedding norm");
std::vector<float> vector(dimensions);
for (int i = 0; i < dimensions; ++i) vector[i] = float(pooled[i]/std::sqrt(norm));
mark("pool_normalize_us");
return vector;
}
json windows(const std::string & text, const std::string & title) const {
@@ -157,8 +225,42 @@ public:
}
json handle(const json & request) {
auto op = request.at("op").get<std::string>();
if(op=="image" || op=="audio") {
auto encoded=request.at("data").get<std::string>();
require(encoded.size()<=12000000 && encoded.size()%4==0,"Invalid media payload length");
static const std::string alphabet="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
std::vector<unsigned char> bytes; bytes.reserve(encoded.size()/4*3);
unsigned accumulator=0;int bits=0;bool padding=false;
for(char c:encoded){if(c=='='){padding=true;continue;} require(!padding,"Invalid base64 padding");auto n=alphabet.find(c);require(n!=std::string::npos,"Invalid base64");accumulator=(accumulator<<6)|unsigned(n);bits+=6;if(bits>=8){bits-=8;bytes.push_back(static_cast<unsigned char>(accumulator>>bits));}}
auto projector=request.at("projector").get<std::string>();
require(!projector.empty() && projector.size()<32768,"Invalid projector path");
ensure_media(projector);
mtmd_bitmap * raw=nullptr;
if(op=="image") {
int w=request.at("width"),h=request.at("height");
require(w>=48 && h>=48 && w%48==0 && h%48==0 && uint64_t(w)*h<=1120*48*48,"Invalid bounded image dimensions");
require(bytes.size()==size_t(w)*h*3,"RGB image length mismatch");
raw=mtmd_bitmap_init(w,h,bytes.data());
} else {
require(bytes.size()%4==0 && !bytes.empty() && bytes.size()/4<=30*16000,"Audio must be at most30seconds mono16kHz PCM");
std::vector<float> pcm(bytes.size()/4);
std::memcpy(pcm.data(),bytes.data(),bytes.size());
for(float f:pcm)require(std::isfinite(f) && std::abs(f)<=1,"Audio PCM outside normalized range");
raw=mtmd_bitmap_init_from_audio(pcm.size(),pcm.data());
}
std::unique_ptr<mtmd_bitmap,decltype(&mtmd_bitmap_free)> bitmap(raw,mtmd_bitmap_free);
require_inference(bitmap != nullptr,"Cannot allocate media bitmap");
return {{"vectors",json::array({embed_bitmap(bitmap.get())})}};
}
if (op == "windows") return {{"windows",windows(request.at("text").get<std::string>(),request.value("title",std::string()))}};
if (op == "query") return {{"vectors",json::array({embed("task: search result | query: "+request.at("text").get<std::string>(),512)})}};
if (op == "query") {
bool debug = request.value("debug_performance",false);
json timings;
auto vector = embed("task: search result | query: "+request.at("text").get<std::string>(),512,debug ? &timings : nullptr);
json response = {{"vectors",json::array({vector})}};
if (debug) response["performance"] = timings;
return response;
}
if (op == "documents") {
const auto & texts = request.at("texts");
require(texts.is_array() && !texts.empty() && texts.size() <= 8,"Batch size must be between 1 and 8");
@@ -187,11 +289,11 @@ int main(int argc, char ** argv) {
bool force_cpu = std::getenv("SEARCHI_FORCE_CPU") && std::string(std::getenv("SEARCHI_FORCE_CPU")) == "1";
Encoder encoder(path,threads,device == "auto" && !force_cpu);
std::cout << json({{"ready",true},{"model","google/embeddinggemma-2"},{"dimensions",dimensions},
{"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG}}).dump() << '\n' << std::flush;
{"device",encoder.device},{"threads",threads},{"runtime",SEARCHI_LLAMA_TAG},{"media_supported",true}}).dump() << '\n' << std::flush;
for (;;) {
std::string line; char c;
while (std::cin.get(c) && c != '\n') {
require(line.size() < max_request_bytes,"Request exceeds 2 MiB"); line.push_back(c);
require(line.size() < max_request_bytes,"Request exceeds 16 MiB"); line.push_back(c);
}
if (!std::cin && line.empty()) break;
json response;
+390 -6
View File
@@ -1,24 +1,35 @@
use crate::TextWindow;
use anyhow::{ensure, Context, Result};
use base64::Engine;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
collections::BTreeMap,
fs::File,
io::{BufRead, BufReader, Read, Write},
path::{Path, PathBuf},
process::{Child, ChildStdin, Command, Stdio},
sync::mpsc::{self, Receiver},
thread,
time::Duration,
time::{Duration, Instant},
};
pub const GEMMA_DIMENSIONS: usize = 768;
pub const GEMMA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mean:l2:query512:doc2048:windows-overlap128:prompts-v1";
pub const GEMMA_MEDIA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mtmd:mean:l2:rgb-catmullrom48-exif-largejpeg-scaled-v2:pcm16k30s-v1";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageEmbedding {
pub vector: Vec<f32>,
pub encoded_width: u32,
pub encoded_height: u32,
pub token_budget: u32,
}
const MODEL_FILE: &str = "embeddinggemma-2-Q8_0.gguf";
const MODEL_BYTES: u64 = 309855456;
const MODEL_SHA256: &str = "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135";
const MAX_REQUEST: usize = 2 * 1024 * 1024;
const MAX_REQUEST: usize = 16 * 1024 * 1024;
const MAX_RESPONSE: usize = 4 * 1024 * 1024;
#[derive(Debug, Clone, Copy)]
@@ -35,6 +46,8 @@ pub struct WorkerInfo {
pub device: String,
pub threads: usize,
pub runtime: String,
#[serde(default)]
pub media_supported: bool,
}
/// Persistent isolated native worker. No HTTP or remote model fallback.
@@ -47,6 +60,8 @@ pub struct GemmaEmbedder {
force_cpu: bool,
model_dir: PathBuf,
worker_path: PathBuf,
projector_verified: bool,
query_profile: Option<BTreeMap<String, u64>>,
}
impl Drop for GemmaEmbedder {
@@ -161,6 +176,8 @@ impl GemmaEmbedder {
force_cpu: cpu,
model_dir,
worker_path: worker,
projector_verified: false,
query_profile: None,
})
}
@@ -169,10 +186,12 @@ impl GemmaEmbedder {
self.healthy,
"Gemma worker session failed; reopen it explicitly"
);
let encode_started = self.query_profile.as_ref().map(|_| Instant::now());
let encoded = serde_json::to_vec(&request)?;
self.record_query_duration("request_encode_us", encode_started);
ensure!(
encoded.len() < MAX_REQUEST,
"Gemma request exceeds 2 MiB protocol budget"
"Gemma request exceeds 16 MiB protocol budget"
);
let response = match self.exchange(&encoded) {
Ok(response) => response,
@@ -180,6 +199,7 @@ impl GemmaEmbedder {
// GPU runtimes can abort rather than return an encoder error.
// Restart once as CPU and retry this same bounded request.
let _ = self.child.wait();
let restart_started = self.query_profile.as_ref().map(|_| Instant::now());
let replacement = Self::load_with_device(
&self.model_dir,
&self.worker_path,
@@ -189,7 +209,13 @@ impl GemmaEmbedder {
.with_context(|| {
format!("CPU worker restart failed after GPU transport error: {error}")
})?;
let diagnostics = self.query_profile.take();
*self = replacement;
self.query_profile = diagnostics;
self.record_query_duration("cpu_restart_us", restart_started);
if let Some(profile) = &mut self.query_profile {
profile.insert("cpu_restart_count".into(), 1);
}
self.exchange(&encoded)?
}
Err(error) => return Err(error),
@@ -201,6 +227,11 @@ impl GemmaEmbedder {
);
self.info.device = device.into();
}
if response["fallback"].as_str() == Some("gpu_inference_failure") {
if let Some(profile) = &mut self.query_profile {
profile.insert("worker_cpu_retry_count".into(), 1);
}
}
if response["fatal"] == true {
self.healthy = false;
self.info.ready = false;
@@ -215,18 +246,21 @@ impl GemmaEmbedder {
}
fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
let write_started = self.query_profile.as_ref().map(|_| Instant::now());
let sent = (|| {
self.stdin.write_all(encoded)?;
self.stdin.write_all(b"\n")?;
self.stdin.flush()
})()
.context("Cannot write local Gemma worker request");
self.record_query_duration("request_write_us", write_started);
if let Err(error) = sent {
self.healthy = false;
self.info.ready = false;
let _ = self.child.kill();
return Err(error);
}
let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
Ok(Ok(value)) => value,
other => {
@@ -236,6 +270,7 @@ impl GemmaEmbedder {
anyhow::bail!("Gemma worker request failed: {other:?}");
}
};
self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
Ok(response)
}
@@ -281,15 +316,265 @@ impl GemmaEmbedder {
}
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
self.query_profile = None;
self.embed_query_inner(query, false)
}
/// Opt-in diagnostics contain only durations and counts, never query text.
pub fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
self.query_profile = Some(BTreeMap::new());
self.embed_query_inner(query, true)
}
pub fn query_diagnostics(&self) -> BTreeMap<String, u64> {
self.query_profile.clone().unwrap_or_default()
}
fn record_query_duration(&mut self, name: &str, started: Option<Instant>) {
if let (Some(profile), Some(started)) = (&mut self.query_profile, started) {
let micros = started.elapsed().as_micros().min(u64::MAX as u128) as u64;
*profile.entry(name.into()).or_default() += micros;
}
}
fn embed_query_inner(&mut self, query: &str, debug: bool) -> Result<Vec<f32>> {
ensure!(
!query.trim().is_empty() && query.len() <= 4096,
"Query must be nonempty and at most 4096 UTF-8 bytes"
);
let response = self.request(json!({"op":"query","text":query}))?;
let response =
self.request(json!({"op":"query","text":query,"debug_performance":debug}))?;
if debug {
if let Some(profile) = &mut self.query_profile {
for name in [
"tokenization_us",
"batch_prepare_us",
"inference_us",
"pool_normalize_us",
"token_count",
] {
if let Some(value) = response["performance"][name].as_u64() {
profile.insert(format!("worker_{name}"), value);
}
}
}
}
let validation_started = self.query_profile.as_ref().map(|_| Instant::now());
let result = extract_vectors(response, 1)?.remove(0);
self.record_query_duration("response_validate_us", validation_started);
Ok(result)
}
fn projector(&mut self) -> Result<PathBuf> {
ensure!(
self.info.media_supported,
"Installed embedding worker does not support media; repair/rebuild the native runtime"
);
let path = self.model_dir.join("mmproj-embeddinggemma-2-Q8_0.gguf");
if !self.projector_verified {
verify_artifact(
&path,
554821024,
"c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f",
)?;
self.projector_verified = true;
}
Ok(path)
}
/// Genuine Google pixel embeddings, serial bounded image operations.
pub fn embed_images(&mut self, paths: &[&Path]) -> Result<Vec<Vec<f32>>> {
self.embed_images_with_budget(paths, 280)
}
/// 70 tokens for video frames, 140 for photos, 280 or 1120 for screenshots.
pub fn embed_images_with_budget(
&mut self,
paths: &[&Path],
budget: u32,
) -> Result<Vec<Vec<f32>>> {
Ok(self
.embed_images_with_metadata(paths, budget)?
.into_iter()
.map(|image| image.vector)
.collect())
}
pub fn embed_images_with_metadata(
&mut self,
paths: &[&Path],
budget: u32,
) -> Result<Vec<ImageEmbedding>> {
ensure!(
!paths.is_empty() && paths.len() <= 8,
"Image batch must contain 1–8 paths"
);
let projector = self.projector()?;
let mut vectors = Vec::with_capacity(paths.len());
for path in paths {
let pixels = decode_image(path, budget)?;
let response=self.request(json!({"op":"image","projector":projector,"width":pixels.width(),"height":pixels.height(),"data":base64::engine::general_purpose::STANDARD.encode(pixels.as_raw())}))?;
vectors.push(ImageEmbedding {
vector: extract_vectors(response, 1)?.remove(0),
encoded_width: pixels.width(),
encoded_height: pixels.height(),
token_budget: budget,
});
}
Ok(vectors)
}
/// One mono 16 kHz normalized PCM window, at most 30 seconds. Callers
/// segment entire recordings and preserve real timestamps; no truncation.
pub fn embed_audio_pcm(&mut self, samples: &[f32]) -> Result<Vec<f32>> {
ensure!(
!samples.is_empty() && samples.len() <= 480000,
"Audio window must contain 1–480000 mono 16 kHz samples"
);
ensure!(
samples.iter().all(|v| v.is_finite() && v.abs() <= 1.0),
"Invalid normalized audio PCM"
);
ensure!(
cfg!(target_endian = "little"),
"Native PCM worker requires little-endian host"
);
let projector = self.projector()?;
let bytes: Vec<u8> = samples.iter().flat_map(|v| v.to_le_bytes()).collect();
let response=self.request(json!({"op":"audio","projector":projector,"data":base64::engine::general_purpose::STANDARD.encode(bytes)}))?;
Ok(extract_vectors(response, 1)?.remove(0))
}
}
fn decode_image(path: &Path, budget: u32) -> Result<image::RgbImage> {
ensure!(
(70..=1120).contains(&budget),
"Image token budget must be 70–1120"
);
ensure!(
std::fs::metadata(path)?.len() <= 64 * 1024 * 1024,
"Image file exceeds 64 MiB"
);
let mut reader = image::ImageReader::open(path)?.with_guessed_format()?;
let is_jpeg = reader.format() == Some(image::ImageFormat::Jpeg);
let mut limits = image::Limits::default();
limits.max_alloc = Some(256 * 1024 * 1024);
limits.max_image_width = Some(16384);
limits.max_image_height = Some(16384);
reader.limits(limits);
use image::ImageDecoder;
let mut decoder = reader
.into_decoder()
.context("Cannot initialize bounded image decoder")?;
let (source_width, source_height) = decoder.dimensions();
ensure!(
u64::from(source_width) * u64::from(source_height) <= 64 * 1024 * 1024
&& decoder.total_bytes() <= 256 * 1024 * 1024,
"Decoded image exceeds memory/pixel budget"
);
let orientation = decoder
.orientation()
.context("Cannot read image EXIF orientation")?;
if is_jpeg && u64::from(source_width) * u64::from(source_height) > 16 * 1024 * 1024 {
drop(decoder);
return decode_large_jpeg(path, budget, source_width, source_height, orientation);
}
ensure!(u64::from(source_width)*u64::from(source_height)<=16*1024*1024,"Non-JPEG image exceeds16megapixels; bounded scaled decoding is currently available for JPEG");
let mut decoded = image::DynamicImage::from_decoder(decoder)
.context("Cannot decode supported image pixels")?;
decoded.apply_orientation(orientation);
ensure!(
u64::from(decoded.width()) * u64::from(decoded.height()) <= 16 * 1024 * 1024,
"Image exceeds 16 megapixels"
);
let (w, h) = image_target(decoded.width(), decoded.height(), budget)?;
Ok(image::imageops::resize(
&decoded.to_rgb8(),
w,
h,
image::imageops::FilterType::CatmullRom,
))
}
fn decode_large_jpeg(
path: &Path,
budget: u32,
width: u32,
height: u32,
orientation: image::metadata::Orientation,
) -> Result<image::RgbImage> {
let (target_width, target_height) = image_target(width, height, budget)?;
let mut decoder = jpeg_decoder::Decoder::new(BufReader::new(File::open(path)?));
decoder.set_max_decoding_buffer_size(256 * 1024 * 1024);
decoder
.read_info()
.context("Cannot read scaled JPEG header")?;
let source = decoder.info().context("Missing scaled JPEG header")?;
ensure!(
u32::from(source.width) == width && u32::from(source.height) == height,
"JPEG changed during decoding"
);
ensure!(
matches!(
source.pixel_format,
jpeg_decoder::PixelFormat::RGB24 | jpeg_decoder::PixelFormat::L8
),
"Large JPEG requires supported RGB24 or grayscale scaled decoding"
);
let (w, h) = decoder.scale(u16::try_from(target_width)?, u16::try_from(target_height)?)?;
ensure!(
u64::from(w) * u64::from(h) * 4 <= 256 * 1024 * 1024,
"Scaled JPEG output exceeds memory budget"
);
let pixels = decoder
.decode()
.context("Cannot decode large JPEG with scaled IDCT")?;
let mut decoded = match source.pixel_format {
jpeg_decoder::PixelFormat::RGB24 => image::DynamicImage::ImageRgb8(
image::RgbImage::from_raw(u32::from(w), u32::from(h), pixels)
.context("Invalid scaled JPEG RGB length")?,
),
jpeg_decoder::PixelFormat::L8 => image::DynamicImage::ImageLuma8(
image::GrayImage::from_raw(u32::from(w), u32::from(h), pixels)
.context("Invalid scaled JPEG grayscale length")?,
),
other => anyhow::bail!(
"Large JPEG pixel format {other:?} is not supported by bounded scaled decoder"
),
};
decoded.apply_orientation(orientation);
let (w, h) = image_target(decoded.width(), decoded.height(), budget)?;
Ok(image::imageops::resize(
&decoded.into_rgb8(),
w,
h,
image::imageops::FilterType::CatmullRom,
))
}
fn image_target(width: u32, height: u32, budget: u32) -> Result<(u32, u32)> {
ensure!(
width > 0 && height > 0 && (70..=1120).contains(&budget),
"Invalid image dimensions/budget"
);
let scale = (f64::from(budget) * 2304.0 / (f64::from(width) * f64::from(height))).sqrt();
let mut w = ((f64::from(width) * scale / 48.0).floor() as u32) * 48;
let mut h = ((f64::from(height) * scale / 48.0).floor() as u32) * 48;
if w == 0 {
w = 48;
h = (height / width * 48).min(budget * 48);
}
if h == 0 {
h = 48;
w = (width / height * 48).min(budget * 48);
}
ensure!(
w >= 48 && h >= 48 && u64::from(w) * u64::from(h) <= u64::from(budget) * 2304,
"Image target outside token budget"
);
Ok((w, h))
}
fn can_retry_cpu(actual_device: &str, force_cpu: bool) -> bool {
actual_device == "gpu" && !force_cpu
}
@@ -331,6 +616,10 @@ fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
}
fn verify_model(path: &Path) -> Result<()> {
verify_artifact(path, MODEL_BYTES, MODEL_SHA256)
}
fn verify_artifact(path: &Path, bytes: u64, sha: &str) -> Result<()> {
let file = File::open(path).with_context(|| {
format!(
"Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
@@ -338,7 +627,7 @@ fn verify_model(path: &Path) -> Result<()> {
)
})?;
ensure!(
file.metadata()?.len() == MODEL_BYTES,
file.metadata()?.len() == bytes,
"EmbeddingGemma 2 model size mismatch"
);
let mut reader = BufReader::new(file);
@@ -352,7 +641,7 @@ fn verify_model(path: &Path) -> Result<()> {
hash.update(&buffer[..count]);
}
ensure!(
format!("{:x}", hash.finalize()) == MODEL_SHA256,
format!("{:x}", hash.finalize()) == sha,
"EmbeddingGemma 2 model SHA256 mismatch"
);
Ok(())
@@ -362,6 +651,101 @@ fn verify_model(path: &Path) -> Result<()> {
mod tests {
use super::*;
#[test]
fn old_text_worker_handshake_does_not_claim_media_support() {
let mut handshake = json!({"ready":true,"model":"google/embeddinggemma-2","dimensions":768,"device":"cpu","threads":1,"runtime":"b11461-no-logits-v1"});
let old: WorkerInfo = serde_json::from_value(handshake.clone()).unwrap();
assert!(!old.media_supported);
handshake["media_supported"] = json!(true);
assert!(
serde_json::from_value::<WorkerInfo>(handshake)
.unwrap()
.media_supported
);
}
#[test]
fn media_image_grid_is_bounded_for_extreme_aspects_and_full_budget() {
for budget in [70, 140, 280, 1120] {
for (width, height) in [(1, 16384), (16384, 1), (1920, 1080), (500, 500)] {
let (w, h) = image_target(width, height, budget).unwrap();
assert_eq!(w % 48, 0);
assert_eq!(h % 48, 0);
assert!(u64::from(w) * u64::from(h) <= u64::from(budget) * 2304);
}
}
assert!(image_target(0, 100, 280).is_err());
assert!(image_target(100, 100, 1121).is_err());
}
#[test]
fn artifact_verification_rejects_projector_size_and_checksum_mismatch() {
let path = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/LICENSE-DigUp.txt");
let bytes = std::fs::read(&path).unwrap();
let sha = format!("{:x}", Sha256::digest(&bytes));
assert!(verify_artifact(&path, bytes.len() as u64, &sha).is_ok());
assert!(verify_artifact(&path, bytes.len() as u64 + 1, &sha).is_err());
assert!(verify_artifact(&path, bytes.len() as u64, &"0".repeat(64)).is_err());
}
#[test]
#[ignore = "Explicit48MP JPEG resource test; no model needed, run release CPU"]
fn large_phone_jpeg_scaled_decode_preserves_exif_and_bounds() {
use image::ImageEncoder;
let image = image::RgbImage::from_fn(8000, 6000, |x, _| {
if x < 4000 {
image::Rgb([240, 10, 10])
} else {
image::Rgb([10, 10, 240])
}
});
let mut encoded = Vec::new();
let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut encoded, 80);
// Little-endian TIFF orientation tag6: rotate90clockwise.
encoder
.set_exif_metadata(vec![
b'I', b'I', 42, 0, 8, 0, 0, 0, 1, 0, 0x12, 1, 3, 0, 1, 0, 0, 0, 6, 0, 0, 0, 0, 0,
0, 0,
])
.unwrap();
encoder.encode_image(&image).unwrap();
drop(image);
let path =
std::env::temp_dir().join(format!("searchi-phone48mp-{}.jpg", std::process::id()));
std::fs::write(&path, &encoded).unwrap();
drop(encoded);
let started = std::time::Instant::now();
let result = decode_image(&path, 280);
let _ = std::fs::remove_file(&path);
let pixels = result.unwrap();
println!(
"48MPactualJPEG scaleddecode+EXIF+resize {:?}; output{}x{} RGB{}bytes",
started.elapsed(),
pixels.width(),
pixels.height(),
pixels.as_raw().len()
);
assert!(pixels.height() > pixels.width());
assert!(pixels.get_pixel(pixels.width() / 2, pixels.height() / 4)[0] > 200);
assert!(pixels.get_pixel(pixels.width() / 2, pixels.height() * 3 / 4)[2] > 200);
assert!(u64::from(pixels.width()) * u64::from(pixels.height()) <= 280 * 2304);
}
#[test]
fn oversized_jpeg_header_rejected_before_pixel_decode() {
let mut encoded = Vec::new();
image::codecs::jpeg::JpegEncoder::new(&mut encoded)
.encode_image(&image::RgbImage::new(16, 16))
.unwrap();
let sof = encoded.windows(2).position(|p| p == [0xff, 0xc0]).unwrap();
encoded[sof + 5..sof + 7].copy_from_slice(&9000u16.to_be_bytes());
encoded[sof + 7..sof + 9].copy_from_slice(&9000u16.to_be_bytes());
let path =
std::env::temp_dir().join(format!("searchi-too-large-{}.jpg", std::process::id()));
std::fs::write(&path, encoded).unwrap();
let result = decode_image(&path, 280);
let _ = std::fs::remove_file(&path);
assert!(result
.unwrap_err()
.to_string()
.contains("memory/pixel budget"));
}
#[test]
fn missing_google_model_fails_without_starting_worker() {
assert!(GemmaEmbedder::load("missing-google-model", "missing-worker", 1).is_err());
}
+4 -1
View File
@@ -1,7 +1,10 @@
//! Local Google EmbeddingGemma 2 inference through an isolated native worker.
//! The rejected MiniLM experiment is available only under `minilm-proof`.
mod gemma;
pub use gemma::{Device, GemmaEmbedder, WorkerInfo, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
pub use gemma::{
Device, GemmaEmbedder, ImageEmbedding, WorkerInfo, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT,
GEMMA_MEDIA_FINGERPRINT,
};
/// A verified source byte range covered by a model token window.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+14
View File
@@ -0,0 +1,14 @@
# Pinned Google media reference fixtures
The `digup-photo.jpg`, `digup-screen.png`, `digup-speech.wav`, and
`digup-media-reference.json` files are copied unchanged from
`ARahim3/DigUp`, commit `6faa29e4c315a3a591de1b41bd78751bdf39588d`,
`Tests/LlamaRuntimeTests/Reference/`. They and the existing text reference
fixture are retained under the adjacent `LICENSE-DigUp.txt` MIT notice.
The source test uses image budgets 280 (screen), 140 (photo), 70 (frame),
and a mono 16 kHz speech recording. The native CPU test compares actual
768-dimensional Google vectors with the pinned llama.cpp b11461 reference
at cosine >=0.999. This is runtime parity evidence, not a production media
relevance corpus. Independent images, audio, and video require separate
quality evaluation.
@@ -0,0 +1 @@
{"made_with":"llama.cpp b11461","vector_version":1,"vectors":{"document":[0.0023304052,-0.021797102,-0.025986522,-0.025894934,0.0046086097,-0.029744277,-0.04186634,0.03234369,0.015701862,-0.13444579,-0.049679857,-0.067124754,-0.008888296,0.049783457,0.04075608,-0.0027419934,0.0494676,0.0092225745,-0.003042357,0.021539804,-0.066572726,-0.04184304,-0.008146685,-0.015800655,0.0028940293,0.026850378,0.023585942,0.012522315,-0.018817414,0.19809185,-0.0541878,-0.033899978,0.06363426,-0.012184333,0.026218394,-0.030783277,-0.016514625,0.0010143976,0.016249185,0.00673781,0.07569905,-0.032087695,0.01618515,0.008316216,0.009405525,0.035788074,0.011760172,-0.0018986556,0.03811712,0.04686964,-0.0019338331,-0.06858261,-0.020950682,0.0061196685,-0.012488069,0.055386145,0.00828158,-0.028973866,0.029866058,0.019298207,0.0064102146,0.03210796,0.031232368,-0.024290783,-0.009193791,-0.00788917,0.07007842,-0.036232397,0.048318822,0.03329114,-0.021271994,-0.01810213,0.056723867,0.026635557,-0.0066668643,-0.03339305,-0.09275154,-0.005457068,-0.037602138,0.06852603,-0.086277,0.0108406255,-0.022301262,0.054799795,0.048695795,-0.0094250385,0.04326083,-0.05584021,0.009427591,0.018164966,0.0077456376,0.012890579,0.011140378,0.009403773,-0.013322473,0.014014791,0.0045265807,0.0026267597,-0.02777846,0.03337617,-0.014311906,-0.0027436074,0.0016559366,0.04267254,-0.0074306116,0.010555651,0.00801582,0.019190526,0.026427642,-0.02305468,0.017346293,-0.254231,-0.037654083,0.004254816,-0.029559055,0.30139065,-0.0017701226,0.012945749,0.045082387,-0.005745192,0.023981802,0.11611545,0.02849787,0.0050043385,0.019032571,0.01543062,0.013486391,-0.0073420904,-0.0120351305,0.027507395,0.017718071,0.0048503615,0.0070377868,-0.00946547,0.015142535,0.013139711,0.019906517,-0.038481727,0.060138095,0.0027306078,-0.010124241,0.020947324,-0.005082965,0.009129657,-0.06960446,0.04814667,0.012469531,-0.017168682,-0.007835558,-0.018339474,0.04499317,-0.048769984,-0.021295255,0.0008460758,0.042394932,0.02629723,0.015862137,-0.027827857,0.02877754,0.044078708,0.02864664,0.015502866,-0.011931367,0.006051371,-0.0676744,-0.032051444,-0.0051915697,0.022663914,0.010766299,-0.045957707,-0.071434505,-0.0071970616,-0.00080499344,-0.019176258,0.039329063,0.07615527,-0.02934876,0.03178097,0.030727895,-0.0014563488,0.033275653,-0.016663045,0.016650142,0.03149862,-0.042355653,-0.01447139,-0.031222051,-0.0371221,-0.071084455,0.02895739,0.012906462,-0.034588896,0.005292062,0.04567936,0.057948757,0.016161263,0.026154758,0.037441123,0.020264363,0.0021852585,0.010503882,0.02564951,-0.0034963347,-0.002975977,-0.05896342,0.090469755,-0.02141068,-0.018266847,0.03979448,0.006404804,0.048994806,-0.0013431136,0.003767527,0.009192836,0.03172842,0.0018301627,-0.032883424,0.070755646,-0.008929331,0.007731602,0.039488316,-0.0026821592,0.047416102,0.017939921,-0.046493225,-0.003571527,0.017453436,0.027679116,-0.045083456,-0.0018627502,-0.0058596916,-0.016741525,-0.041559335,-0.07398639,0.035852198,-0.01250213,-0.027620463,-0.022970477,-0.007722595,0.006369175,-0.027909646,0.05087572,0.04319605,0.0249135,-0.045503743,-0.050388955,-0.07399011,-0.031614304,-0.009977491,-0.011919013,0.017587284,0.046735547,0.051120915,-0.01913088,-0.0040965537,-0.030126736,0.022246089,-0.024694057,-0.0011967963,-0.0076620826,-0.0001045683,0.025630128,0.0023922755,-0.016251149,0.029261485,0.026775412,-0.0011348778,-0.023079926,0.0046075485,0.013615946,-0.022271913,0.0063756416,-0.011097959,-0.057059642,0.017879428,0.039326653,-0.010385263,-0.024296982,-0.0056849266,-0.021441974,0.023893535,-0.004312103,-0.0029023795,0.011346968,-0.05022057,-0.00962676,-0.030292084,0.0022458949,-0.0012382208,0.014016518,0.0130033605,-0.0075721424,0.022480633,0.0378293,-0.0060748337,0.020296352,-0.038013346,0.02737668,0.051704004,0.029559156,-0.008555367,-0.040540576,0.0028047708,-0.01888827,-0.017475538,0.022049233,-0.018655647,0.00035712068,0.0007242559,-0.031585597,0.002809768,-0.010847336,0.019464653,-0.038227003,0.004094894,-0.0045381063,-0.032818098,-0.025449943,-0.0048597767,-0.052715924,0.024247607,-0.025617195,-0.028776655,-0.01491966,0.027141426,0.022653097,0.01347185,0.0010195217,0.0077305734,0.006165836,-0.009121481,-0.011312435,0.0010023145,0.033321623,0.0020216978,-0.04601626,-0.020557044,-0.023531808,0.013099387,0.049464565,-0.0083905,-0.12218178,-0.0090951715,-0.037435923,-0.043180246,0.008593723,0.003133877,0.03239143,-0.02740052,0.028175933,0.00713765,0.010822759,-0.011402951,-0.042760924,-0.03611006,0.011457286,0.060568936,0.0027335465,-0.024305783,-0.004420147,0.0068441983,0.002837347,-0.02218109,-0.022486687,-0.01941801,0.01667096,-0.03958557,-0.0077348636,-0.0012779543,-0.009598079,-0.008962751,0.033793483,0.018625712,-0.006633453,0.026242023,0.021685809,0.013799533,-0.0060389466,0.0020937775,0.02297482,-0.02592566,0.0011351658,0.00793861,-0.01809629,0.050504446,-0.026299631,-0.028017156,0.006029957,-0.02648004,0.03030571,0.016809817,0.29351303,-0.028621448,0.01694766,0.021401705,0.012320341,0.0015229117,-0.Line truncated
Binary file not shown.

After

Width:  |  Height:  |  Size: 165 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.3 KiB

Binary file not shown.
@@ -1,6 +1,49 @@
use anyhow::Result;
use searchi_semantic::{Device, GemmaEmbedder, GEMMA_DIMENSIONS};
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_debug_query_reports_actual_worker_stages_without_changing_vectors() -> 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 embedder = GemmaEmbedder::load_with_device(model, worker, 1, Device::Cpu)?;
let query = "supplier contract cancellation policy";
let debug = embedder.embed_query_debug(query)?;
let diagnostics = embedder.query_diagnostics();
for name in [
"request_encode_us",
"request_write_us",
"worker_wait_and_response_decode_us",
"response_validate_us",
"worker_tokenization_us",
"worker_batch_prepare_us",
"worker_inference_us",
"worker_pool_normalize_us",
"worker_token_count",
] {
assert!(
diagnostics.contains_key(name),
"Missing actual stage {name}: {diagnostics:?}"
);
}
assert!(diagnostics["worker_inference_us"] > 0);
assert!(diagnostics["worker_token_count"] > 0);
assert!(!serde_json::to_string(&diagnostics)?.contains(query));
let normal = embedder.embed_query(query)?;
assert!(embedder.query_diagnostics().is_empty());
let cosine: f64 = debug
.iter()
.zip(&normal)
.map(|(a, b)| f64::from(*a) * f64::from(*b))
.sum();
assert!(
cosine >= 0.99999,
"Debug tracing changed embeddings: {cosine}"
);
println!("CPU query diagnostics: {diagnostics:?}; debug/normal cosine {cosine:.8}");
Ok(())
}
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_real_embeddings_and_verified_tail_windows() -> Result<()> {
@@ -0,0 +1,57 @@
use anyhow::Result;
use searchi_semantic::{Device, GemmaEmbedder};
use std::path::Path;
#[test]
#[ignore = "Requires explicitly provisioned Google media model and CPU worker"]
fn google_cpu_pixels_match_digup_reference() -> Result<()> {
let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set model path");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set worker path");
let mut embedder = GemmaEmbedder::load_with_device(model, worker, 1, Device::Cpu)?;
let fixtures = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
let reference: serde_json::Value =
serde_json::from_str(include_str!("fixtures/digup-media-reference.json"))?;
for (name, file, budget) in [
("screen-280", "digup-screen.png", 280),
("photo-140", "digup-photo.jpg", 140),
("frame-70", "digup-photo.jpg", 70),
] {
let path = fixtures.join(file);
let vector = embedder
.embed_images_with_budget(&[path.as_path()], budget)?
.remove(0);
let cosine: f64 = vector
.iter()
.zip(reference["vectors"][name].as_array().unwrap())
.map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
.sum();
println!("Google image {name} reference cosine={cosine:.8}");
assert!(cosine >= 0.999, "Image parity below0.999: {name} {cosine}");
}
assert_eq!(embedder.info.device, "cpu");
let wav = include_bytes!("fixtures/digup-speech.wav");
let mut offset = 12;
let mut samples = Vec::new();
while offset + 8 <= wav.len() {
let length = u32::from_le_bytes(wav[offset + 4..offset + 8].try_into()?) as usize;
if &wav[offset..offset + 4] == b"data" {
for pair in wav[offset + 8..offset + 8 + length].as_chunks::<2>().0 {
samples.push(f32::from(i16::from_le_bytes(*pair)) / 32768.0);
}
break;
}
offset += 8 + length + (length % 2);
}
let vector = embedder.embed_audio_pcm(&samples)?;
let cosine: f64 = vector
.iter()
.zip(reference["vectors"]["speech"].as_array().unwrap())
.map(|(a, b)| f64::from(*a) * b.as_f64().unwrap())
.sum();
println!("Google audio reference cosine={cosine:.8}");
assert!(
cosine >= 0.999,
"Audio reference parity below0.999: {cosine}"
);
Ok(())
}
+1
View File
@@ -10,6 +10,7 @@ anyhow.workspace = true
serde.workspace = true
searchi-core = { path = "../searchi-core" }
searchi-semantic = { path = "../searchi-semantic" }
searchi-media = { path = "../searchi-media" }
rusqlite = { version = "0.32", features = ["bundled"] }
sha2 = "0.10"
fs2 = "0.4"
+43
View File
@@ -0,0 +1,43 @@
# Shared hybrid retrieval
`HybridSearch` combines genuine lexical and Google multimodal retrieval. The
CLI and desktop use the same service; it has no desktop-shell dependency.
`search`, `search_semantic` and `search_semantic_with_origin` preserve grouped
parent files. ANN candidates expand until enough unique eligible parents exist
or the complete bounded eligible chunk set is requested. Filename/type/folder,
size/date/duration and explicit source filters apply before ANN top-k. Media
candidate filesystem metadata is checked once per parent per query. Ninety-six
near-identical video vectors cannot hide a second eligible parent (regression
tested). HNSW remains approximate retrieval, not an exact recall guarantee.
Explicit spoken/audio wording prefers actual audio-source vectors, preserving
exact lexical filename priority. Conservative English/German token cues detect
requests such as “someone mentions”, “spoken”, “audio”, “jemand erwähnt” and
“Gespräch”. Visual/document cues veto automatic routing. This finite parser is
not language understanding; paraphrases, negation and mixed intents can remain
ambiguous. If the preferred lane has no eligible results, ordinary retrieval
is used. `SemanticOrigin::All` bypasses the hint; `Audio`, `Image`, `Video` and
`Text` select an explicit channel without fallback or reembedding.
Persisted custom sampling limits can be reopened without generating embeddings
again only when the complete model and preprocessing fingerprint matches.
The custom suffix must contain the eight valid numeric sampling limits. Model
or preprocessing changes never silently reuse incompatible vectors.
`search_measured(query,limit,options,debug)` returns `SearchResponse` with results
and monotonic total milliseconds/result count. Basic mode stores no trace.
`search_traced(...,&mut TraceCollector,parent)` shares the native request clock
so queue/lock time and actual engine stages remain comparable. Debug mode adds
bounded nested spans for lexical SQL, query embedding, worker diagnostics,
filtered eligibility and its cache, ANN read/checksum/deserialization/search
passes, parent materialization and validation, and rank fusion. Static stage
names and numerical counters exclude query/source text and paths. Worker-local
microsecond counters have a separate clock and are not aligned request spans.
Measured/source APIs retain existing result behavior except the explicit
spoken/audio preference. Debug instrumentation does not change vector policy.
Audio timestamps represent segment starts, not exact word locations. Smaller
windows alone have not demonstrated accurate spoken navigation. No-answer
calibration, music discrimination, exact spoken positions, process isolation,
unchanged-vector reuse and production platform checks remain release gates.
+207 -18
View File
@@ -1,10 +1,15 @@
//! Shared offline hybrid search using real Google embeddings and persisted,
//! filtered HNSW retrieval. Production-scale validation remains a release gate.
mod media;
mod routing;
pub use routing::SemanticOrigin;
mod store;
pub use searchi_media::{FfmpegTools, MediaLimits};
use anyhow::{ensure, Result};
use searchi_core::metrics::{MetricsMode, SpanToken, TraceCollector};
use searchi_core::SearchIndex;
pub use searchi_core::{SearchOptions, SearchResult};
pub use searchi_core::{SearchMetrics, SearchOptions, SearchResponse, SearchResult};
pub use searchi_semantic::WorkerInfo;
use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
use serde::Serialize;
@@ -17,7 +22,24 @@ 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}")
format!(
"{model}:{CHUNK_POLICY}:{}:image280-frame70-audio30s-overlap2-distributed240-scene120-v1",
searchi_semantic::GEMMA_MEDIA_FINGERPRINT
)
}
fn reusable_policy(expected: &str, active: &str) -> bool {
if active == expected || active == format!("{expected}:text-only") {
return true;
}
let Some(policy) = active.strip_prefix(&format!("{expected}:custom-")) else {
return false;
};
let values = policy
.split('-')
.map(str::parse::<u64>)
.collect::<std::result::Result<Vec<_>, _>>();
matches!(values,Ok(v) if v.len()==8 && v[0]>=2 && (5000..=30000).contains(&v[3]) && v[4]<v[3] && v[5]>0 && v[6]<=u32::MAX as u64)
}
/// Model-specific query/document prompts and actual tokenizer windows belong
@@ -28,6 +50,18 @@ pub trait EmbeddingProvider: Send {
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>>;
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>>;
fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
self.embed_query(query)
}
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
Default::default()
}
fn embed_images(&mut self, _paths: &[&Path], _budget: u32) -> Result<Vec<Vec<f32>>> {
anyhow::bail!("Provider does not support real images")
}
fn embed_audio(&mut self, _pcm: &[f32]) -> Result<Vec<f32>> {
anyhow::bail!("Provider does not support real audio")
}
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
None
}
@@ -49,13 +83,34 @@ impl EmbeddingProvider for GemmaEmbedder {
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
self.embed_query(query)
}
fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
GemmaEmbedder::embed_query_debug(self, query)
}
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
GemmaEmbedder::query_diagnostics(self)
}
fn embed_images(&mut self, paths: &[&Path], budget: u32) -> Result<Vec<Vec<f32>>> {
self.embed_images_with_budget(paths, budget)
}
fn embed_audio(&mut self, pcm: &[f32]) -> Result<Vec<f32>> {
self.embed_audio_pcm(pcm)
}
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
Some(self.info.clone())
}
}
#[derive(Debug, Serialize)]
pub struct MediaFailure {
pub path: String,
pub message: String,
}
#[derive(Debug, Serialize)]
pub struct EmbeddingReport {
pub media_files: usize,
pub media_failures: usize,
pub failures: Vec<MediaFailure>,
pub passages: usize,
pub chunks: usize,
pub generation: i64,
@@ -64,6 +119,7 @@ pub struct EmbeddingReport {
#[derive(Debug, Serialize)]
pub struct EmbeddingStatus {
pub failures: Vec<MediaFailure>,
pub generation: Option<i64>,
pub chunks: usize,
pub model_fingerprint: Option<String>,
@@ -98,9 +154,10 @@ impl HybridSearch {
|row| row.get(0),
)?;
let pending: usize =
connection.query_row("SELECT count(*) FROM passages", [], |row| row.get(0))?;
connection.query_row("SELECT (SELECT count(*) FROM passages)+(SELECT count(*) FROM files WHERE extension IN ('png','jpg','jpeg','webp','bmp','mp4','mkv','mov','avi','webm','m4v','mpeg','mpg','mp3','flac','wav','m4a','ogg','opus','aac','aiff','aif'))", [], |row| row.get(0))?;
if !exists {
return Ok(EmbeddingStatus {
failures: Vec::new(),
generation: None,
chunks: 0,
model_fingerprint: None,
@@ -113,8 +170,9 @@ impl HybridSearch {
[],
|row| row.get(0),
)?;
if version < 5 {
if version < 6 {
return Ok(EmbeddingStatus {
failures: Vec::new(),
generation: None,
chunks: 0,
model_fingerprint: None,
@@ -123,7 +181,7 @@ impl HybridSearch {
});
}
ensure!(
version == 5,
version == 6,
"Unsupported semantic schema version {version}"
);
let (generation,fingerprint,revision,current):(Option<i64>,Option<String>,Option<i64>,i64)=connection.query_row(
@@ -134,9 +192,19 @@ impl HybridSearch {
[generation],
|row| row.get(0),
)?;
let mut failures_query=connection.prepare("SELECT path,message FROM semantic_media_errors WHERE generation=?1 ORDER BY path LIMIT 1000")?;
let failures = failures_query
.query_map([generation], |r| {
Ok(MediaFailure {
path: r.get(0)?,
message: r.get(1)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let stale = revision != Some(current)
|| fingerprint.as_deref() != Some(generation_fingerprint(GEMMA_FINGERPRINT).as_str());
Ok(EmbeddingStatus {
failures,
generation,
chunks,
model_fingerprint: fingerprint,
@@ -162,7 +230,12 @@ impl HybridSearch {
) -> Result<Self> {
let lexical = SearchIndex::open(&db)?;
let store = Store::open(db.as_ref(), embedder.dimensions())?;
let generation_fingerprint = generation_fingerprint(embedder.fingerprint());
let expected = generation_fingerprint(embedder.fingerprint());
let active = store.active_fingerprint()?;
let generation_fingerprint = match active {
Some(active) if reusable_policy(&expected, &active) => active,
_ => expected,
};
Ok(Self {
lexical,
embedder,
@@ -189,6 +262,11 @@ impl HybridSearch {
/// 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 previous_fingerprint = self.generation_fingerprint.clone();
self.generation_fingerprint = format!(
"{}:text-only",
generation_fingerprint(self.embedder.fingerprint())
);
let _lock = self.store.build_lock()?;
let generation = self.store.begin(&self.generation_fingerprint)?;
let built = self.build(generation);
@@ -199,6 +277,7 @@ impl HybridSearch {
}
Err(error) => {
self.store.abandon(generation)?;
self.generation_fingerprint = previous_fingerprint;
Err(error)
}
}
@@ -206,6 +285,9 @@ impl HybridSearch {
fn build(&mut self, generation: i64) -> Result<EmbeddingReport> {
let mut report = EmbeddingReport {
media_files: 0,
media_failures: 0,
failures: Vec::new(),
passages: 0,
chunks: 0,
generation,
@@ -253,21 +335,81 @@ impl HybridSearch {
query: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let mut trace = TraceCollector::new(MetricsMode::Basic);
self.search_traced(query, limit, options, &mut trace, None)
}
pub fn search_measured(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
debug: bool,
) -> Result<SearchResponse<SearchResult>> {
let mut trace = TraceCollector::new(if debug {
MetricsMode::Debug
} else {
MetricsMode::Basic
});
let results = self.search_traced(query, limit, options, &mut trace, None)?;
let metrics = trace.finish(results.len());
Ok(SearchResponse { results, metrics })
}
pub fn search_traced(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
trace.timed("hybrid.search", parent, |trace, parent| {
self.search_inner(query, limit, options, trace, parent)
})
}
fn search_inner(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
let candidates = limit.saturating_mul(4).clamp(20, 1000);
// Core validates filters and query budgets even for empty requests.
let mut lexical = self
.lexical
.search_with_options(query, candidates, options)?;
.search_with_options_traced(query, candidates, options, trace, parent)?;
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())?;
trace.timed("hybrid.identifier_priority", parent, |_, _| {
promote_identifiers(query.trim(), &mut lexical)
});
let embedding_span = trace.begin("semantic.query_embedding", parent);
let vector = if trace.is_debug() {
self.embedder.embed_query_debug(query.trim())
} else {
self.embedder.embed_query(query.trim())
};
trace.counter(embedding_span, "failed", u64::from(vector.is_err()));
trace.end(
embedding_span,
if trace.is_debug() {
self.embedder.query_diagnostics()
} else {
Default::default()
},
);
let vector = vector?;
let semantic =
self.store
.search(&vector, &self.generation_fingerprint, candidates, options)?;
Ok(fuse(lexical, semantic, limit.min(1000)))
self.routed_semantic_traced(query, &vector, candidates, options, trace, parent)?;
let results = trace.timed("hybrid.rank_fusion", parent, |trace, span| {
trace.counter(span, "lexical_files", lexical.len() as u64);
trace.counter(span, "semantic_files", semantic.len() as u64);
fuse(lexical, semantic, limit.min(1000))
});
Ok(results)
}
/// Semantic lane for reproducible relevance evaluation and diagnostics.
@@ -282,12 +424,7 @@ impl HybridSearch {
return Ok(Vec::new());
}
let vector = self.embedder.embed_query(query.trim())?;
self.store.search(
&vector,
&self.generation_fingerprint,
limit.min(1000),
options,
)
self.routed_semantic(query, &vector, limit.min(1000), options)
}
}
@@ -417,6 +554,58 @@ fn fuse(
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn failed_query_inference_retains_diagnostics_in_shared_trace() -> Result<()> {
// Failure-only provider: no fabricated embeddings are generated.
struct FailingProvider;
impl EmbeddingProvider for FailingProvider {
fn fingerprint(&self) -> &str {
"failure-injection"
}
fn dimensions(&self) -> usize {
768
}
fn split_windows(&mut self, _: &str, _: &str) -> Result<Vec<TextWindow>> {
anyhow::bail!("unused")
}
fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
anyhow::bail!("unused")
}
fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
anyhow::bail!("injected worker failure")
}
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
std::collections::BTreeMap::from([("request_write_us".into(), 7)])
}
}
let directory = tempfile::tempdir()?;
let db = directory.path().join("index.db");
let mut service = HybridSearch::with_provider(&db, Box::new(FailingProvider))?;
let mut trace = TraceCollector::new(MetricsMode::Debug);
assert!(service
.search_traced(
"ordinary query",
5,
&SearchOptions::default(),
&mut trace,
None
)
.is_err());
let metrics = trace.finish(0);
let spans = metrics.trace.unwrap();
let embedding = spans
.iter()
.find(|s| s.name == "semantic.query_embedding")
.unwrap();
assert_eq!(embedding.counters["failed"], 1);
assert_eq!(embedding.counters["request_write_us"], 7);
let parent = &spans[embedding.parent.unwrap()];
assert!(
embedding.start_ms + embedding.duration_ms
<= parent.start_ms + parent.duration_ms + 1e-6
);
Ok(())
}
fn result(path: &str, kind: &str) -> SearchResult {
SearchResult {
path: path.into(),
+232
View File
@@ -0,0 +1,232 @@
use super::*;
use searchi_media::*;
#[derive(Debug)]
struct InferenceFailure(anyhow::Error);
impl std::fmt::Display for InferenceFailure {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Media inference failed: {}", self.0)
}
}
impl std::error::Error for InferenceFailure {}
impl HybridSearch {
/// Publish text and actual image/frame/audio vectors in one atomic generation.
pub fn index_embeddings_with_media(
&mut self,
tools: &FfmpegTools,
limits: &MediaLimits,
) -> Result<EmbeddingReport> {
let previous_fingerprint = self.generation_fingerprint.clone();
self.generation_fingerprint = generation_fingerprint(self.embedder.fingerprint());
let defaults = MediaLimits::default();
if limits.max_frames != defaults.max_frames
|| limits.frame_interval_ms != defaults.frame_interval_ms
|| limits.max_scene_frames != defaults.max_scene_frames
|| limits.audio_window_ms != defaults.audio_window_ms
|| limits.audio_overlap_ms != defaults.audio_overlap_ms
|| limits.max_audio_windows != defaults.max_audio_windows
|| limits.max_duration_ms != defaults.max_duration_ms
|| limits.max_file_bytes != defaults.max_file_bytes
{
self.generation_fingerprint.push_str(&format!(
":custom-{}-{}-{}-{}-{}-{}-{}-{}",
limits.max_frames,
limits.frame_interval_ms,
limits.max_scene_frames,
limits.audio_window_ms,
limits.audio_overlap_ms,
limits.max_audio_windows,
limits.max_duration_ms,
limits.max_file_bytes
));
}
let _lock = self.store.build_lock()?;
let generation = self.store.begin(&self.generation_fingerprint)?;
let built = (|| {
let mut report = self.build(generation)?;
let mut after = 0;
let mut media_count = 0usize;
loop {
let files = self.store.media_batch(after)?;
if files.is_empty() {
break;
}
for (id, path, _, _) in files {
after = id;
let path = Path::new(&path);
let Some(kind) = classify(path) else {
continue;
};
media_count += 1;
ensure!(
media_count <= 10000,
"Media generation exceeds 10000-file resource budget"
);
let before_chunks = report.chunks;
match self.build_media_file(
generation,
id,
path,
kind,
tools,
limits,
&mut report,
) {
Ok(()) => report.media_files += 1,
Err(error) => {
if error.downcast_ref::<InferenceFailure>().is_some() {
return Err(error);
}
self.store.media_failure(
generation,
&path.to_string_lossy(),
&format!("{error:#}"),
)?;
report.chunks = before_chunks;
if report.failures.len() < 1000 {
report.failures.push(MediaFailure {
path: path.to_string_lossy().into_owned(),
message: format!("{error:#}").chars().take(2048).collect(),
});
}
report.media_failures += 1;
}
}
}
}
Ok(report)
})();
match built {
Ok(report) => {
self.store.activate(generation)?;
Ok(report)
}
Err(error) => {
self.store.abandon(generation)?;
self.generation_fingerprint = previous_fingerprint;
Err(error)
}
}
}
#[allow(clippy::too_many_arguments)]
fn build_media_file(
&mut self,
generation: i64,
id: i64,
path: &Path,
kind: MediaKind,
tools: &FfmpegTools,
limits: &MediaLimits,
report: &mut EmbeddingReport,
) -> Result<()> {
let identity = MediaIdentity::read(
path,
if kind == MediaKind::Image {
64 * 1024 * 1024
} else {
limits.max_file_bytes
},
)?;
let mut ordinal = 0;
if kind == MediaKind::Image {
validate_image(path)?;
ensure!(
report.chunks < MAX_SCAN_CHUNKS,
"Media exceeds generation chunk budget"
);
let vectors = self
.embedder
.embed_images(&[path], 280)
.map_err(|e| anyhow::Error::new(InferenceFailure(e)))?;
ensure!(vectors.len() == 1, "Image embedding count mismatch");
self.store
.insert_media(generation, id, ordinal, &identity, "image", 0, &vectors[0])?;
report.chunks += 1;
} else {
let metadata = probe(tools, path, limits)?;
let tags = metadata
.tags
.iter()
.map(|(k, v)| format!("{k}: {v}"))
.collect::<Vec<_>>()
.join("\n");
self.lexical.set_media_metadata(
&path.to_string_lossy(),
&tags,
metadata.duration_ms,
identity.bytes,
identity.modified_ns,
)?;
if metadata.has_video && kind == MediaKind::Video {
let mut seen = std::collections::HashSet::new();
let mut scenes = 0;
for time in frame_times(&metadata, limits)? {
let frames = decode_frames(
tools,
path,
&metadata,
time,
scenes < limits.max_scene_frames,
limits,
)?;
for (i, frame) in frames.frames.iter().enumerate() {
if !seen.insert(frame.timestamp_ms) {
continue;
}
if i > 0 {
scenes += 1;
}
ensure!(
report.chunks < MAX_SCAN_CHUNKS,
"Media exceeds generation chunk budget"
);
let vectors = self
.embedder
.embed_images(&[frame.path.as_path()], 70)
.map_err(|e| anyhow::Error::new(InferenceFailure(e)))?;
ensure!(vectors.len() == 1, "Frame embedding count mismatch");
self.store.insert_media(
generation,
id,
ordinal,
&identity,
"video",
frame.timestamp_ms,
&vectors[0],
)?;
ordinal += 1;
report.chunks += 1;
}
}
}
if metadata.has_audio {
for window in audio_windows(&metadata, limits)? {
ensure!(
report.chunks < MAX_SCAN_CHUNKS,
"Audio exceeds generation chunk budget"
);
let pcm = decode_audio(tools, path, window, limits)?;
let vector = self
.embedder
.embed_audio(&pcm)
.map_err(|e| anyhow::Error::new(InferenceFailure(e)))?;
self.store.insert_media(
generation,
id,
ordinal,
&identity,
"audio",
window.start_ms,
&vector,
)?;
ordinal += 1;
report.chunks += 1;
}
}
}
identity.ensure_unchanged()?;
Ok(())
}
}
+213
View File
@@ -0,0 +1,213 @@
use super::*;
/// Explicit semantic source channel, independent of filename/type filters.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SemanticOrigin {
All,
Text,
Image,
Video,
Audio,
}
impl SemanticOrigin {
fn source_type(self) -> Option<&'static str> {
match self {
Self::All => None,
Self::Text => Some("text"),
Self::Image => Some("image"),
Self::Video => Some("video"),
Self::Audio => Some("audio"),
}
}
}
/// A conservative intent hint, not a language understanding model. Explicit
/// document/visual cues veto it; ordinary topic and instrument queries do not
/// silently restrict modality. Word boundaries avoid filename substrings.
fn preferred_origin(query: &str) -> Option<SemanticOrigin> {
if query.split_whitespace().count() == 1 && Path::new(query).extension().is_some() {
return None;
}
let lower = query.to_lowercase();
let words: Vec<_> = lower
.split(|c: char| !c.is_alphanumeric())
.filter(|s| !s.is_empty())
.collect();
let has = |values: &[&str]| words.iter().any(|word| values.contains(word));
if has(&[
"image",
"photo",
"picture",
"screenshot",
"diagram",
"document",
"pdf",
"text",
"paper",
"bild",
"foto",
"screenshot",
"dokument",
"textdatei",
]) {
return None;
}
if has(&[
"audio",
"spoken",
"speech",
"recording",
"conversation",
"interview",
"podcast",
"sound",
"music",
"gesprochen",
"gespräch",
"gespraech",
"geräusch",
"geraeusch",
"musik",
"hörbar",
"hörbarer",
"hörbare",
"hoerbar",
"hoerbarer",
"hoerbare",
"tonaufnahme",
]) {
return Some(SemanticOrigin::Audio);
}
if (has(&["someone", "speaker"])
&& has(&[
"mentions",
"mentioned",
"says",
"said",
"speaks",
"speaking",
]))
|| (has(&["jemand", "sprecher", "sprecherin"])
&& has(&["erwähnt", "erwaehnt", "sagt", "gesagt", "spricht"]))
{
return Some(SemanticOrigin::Audio);
}
None
}
impl HybridSearch {
pub(super) fn routed_semantic(
&mut self,
query: &str,
vector: &[f32],
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let mut trace = TraceCollector::new(MetricsMode::Basic);
self.routed_semantic_traced(query, vector, limit, options, &mut trace, None)
}
pub(super) fn routed_semantic_traced(
&mut self,
query: &str,
vector: &[f32],
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
trace.timed("semantic.routing", parent, |trace, parent| {
if let Some(origin) = preferred_origin(query) {
trace.counter(parent, "audio_preference", 1);
let preferred = self.store.search_origin_traced(
vector,
&self.generation_fingerprint,
limit,
options,
origin.source_type(),
trace,
parent,
)?;
if !preferred.is_empty() {
return Ok(preferred);
}
}
trace.counter(parent, "general_lane", 1);
self.store.search_origin_traced(
vector,
&self.generation_fingerprint,
limit,
options,
None,
trace,
parent,
)
})
}
/// Query one actual source channel, or All to bypass automatic spoken/audio
/// preference. This never embeds files again or changes an active policy.
pub fn search_semantic_with_origin(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
origin: SemanticOrigin,
) -> 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_origin(
&vector,
&self.generation_fingerprint,
limit.min(1000),
options,
origin.source_type(),
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn persisted_custom_limits_require_identical_model_and_preprocessing() {
let expected = "pinned-model:pinned-media:preprocessing-v1";
assert!(reusable_policy(
expected,
&format!("{expected}:custom-240-2000-120-10000-2000-1800-14400000-17179869184")
));
for active in ["other-model:pinned-media:preprocessing-v1:custom-240-2000-120-10000-2000-1800-14400000-17179869184","pinned-model:pinned-media:preprocessing-v2:custom-240-2000-120-10000-2000-1800-14400000-17179869184","pinned-model:pinned-media:preprocessing-v1:custom-unknown"] {assert!(!reusable_policy(expected,active));}
}
#[test]
fn explicit_spoken_intent_is_conservative_and_bilingual() {
for query in [
"where someone mentions a contract",
"a speaker says the date",
"spoken discussion of ecology",
"audio of a waterfall",
"jemand erwähnt das Budget",
"Wo spricht jemand über Kündigung?",
"Gespräch über die neue Wohnung",
"hörbarer Regen",
] {
assert_eq!(
preferred_origin(query),
Some(SemanticOrigin::Audio),
"{query}"
);
}
for query in [
"contract deadline",
"violin solo",
"photo of someone speaking",
"Bild von einem Sprecher",
"a document that someone mentioned",
"speech.png",
"a screenshot of an audio waveform",
] {
assert_eq!(preferred_origin(query), None, "{query}");
}
}
}
+651 -24
View File
@@ -8,6 +8,7 @@ use std::{
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};
@@ -24,6 +25,9 @@ struct FilterKey {
max_size: Option<i64>,
modified_after: Option<i64>,
modified_before: Option<i64>,
min_duration: Option<u32>,
max_duration: Option<u32>,
source_type: Option<String>,
}
struct EligibilityCache {
@@ -57,8 +61,13 @@ pub(crate) fn content_hash(text: &str) -> String {
#[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),
@@ -85,7 +94,7 @@ impl Store {
|row| row.get(0),
)?;
ensure!(
(1..=5).contains(&version),
(1..=6).contains(&version),
"Unsupported semantic schema version {version}"
);
initialization.execute_batch("
@@ -130,8 +139,25 @@ impl Store {
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;",
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))")?;
initialization.commit()?;
let canonical = db.canonicalize()?;
let lock_path = PathBuf::from(format!("{}.semantic.lock", canonical.to_string_lossy()));
@@ -170,6 +196,67 @@ impl Store {
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,
@@ -312,6 +399,61 @@ impl Store {
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()),
@@ -321,17 +463,30 @@ impl Store {
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],
@@ -346,10 +501,15 @@ impl Store {
[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!(
@@ -357,8 +517,12 @@ impl Store {
"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()
@@ -384,13 +548,25 @@ impl Store {
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),
};
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
@@ -404,8 +580,17 @@ impl Store {
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(&snapshot, active, revision, &filter, None)?;
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) {
@@ -413,12 +598,14 @@ impl Store {
Arc::make_mut(&mut cache.keys).remove(&key);
}
}
let changed = collect_eligibility(
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);
@@ -428,12 +615,22 @@ impl Store {
}
cache
}
_ => collect_eligibility(&snapshot, active, revision, &filter, None)?,
_ => 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);
@@ -443,40 +640,97 @@ impl Store {
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 FROM semantic_chunks c JOIN files f ON f.path=c.path WHERE c.rowid=?1 AND c.generation=?2")?;
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 =
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: "semantic".into(),
source: Some(SourceLocation {
kind: row.get(3)?,
number: row.get(4)?,
}),
extraction_error: None,
})
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;
}
@@ -488,6 +742,13 @@ impl Store {
.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)
}
@@ -614,6 +875,30 @@ impl Store {
}
}
#[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,
@@ -621,12 +906,15 @@ fn collect_eligibility(
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
let mut statement=connection.prepare("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) ORDER BY c.passage_id,c.ordinal LIMIT ?8")?;
let mut rows = statement.query(params![
generation,
@@ -637,7 +925,10 @@ fn collect_eligibility(
filter.modified_after,
filter.modified_before,
MAX_SCAN_CHUNKS as i64 + 1,
path
path,
filter.min_duration,
filter.max_duration,
filter.source_type
])?;
let mut hashes = HashMap::<i64, String>::new();
let mut keys = HashSet::new();
@@ -650,7 +941,9 @@ fn collect_eligibility(
"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) {
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)?);
@@ -678,6 +971,340 @@ mod tests {
use searchi_core::SearchIndex;
const DIMENSIONS: usize = 768;
#[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<()> {