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
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:
1 parent
d3bb8c7618
commit
b8400c4fbf
110 files changed
+79395
-250
No files matched your search
+109
-15
@@ -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() {
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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<()> {
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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"] }
|
||||
@@ -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.
|
||||
@@ -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 })
|
||||
}
|
||||
@@ -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"] }
|
||||
|
||||
@@ -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)
|
||||
@@ -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;
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<()> {
|
||||
|
||||
Reference in new issue
Block a user