Reuse completed Google embeddings and resume semantic staging
This commit is contained in:
karl committed
1 parent
0b56ab553c
commit
5c18e7b49c
11 files changed
+1677
-148
No files matched your search
@@ -0,0 +1,417 @@
|
||||
//! Incremental acceptance uses genuine Google vectors, never stand-in embeddings.
|
||||
use anyhow::Result;
|
||||
use rusqlite::Connection;
|
||||
use searchi_core::SearchIndex;
|
||||
use searchi_semantic::{GemmaEmbedder, TextWindow};
|
||||
use searchi_service::{EmbeddingProvider, HybridSearch, SearchOptions};
|
||||
use std::{
|
||||
fs,
|
||||
path::Path,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
#[derive(Default, Debug, Clone, PartialEq, Eq)]
|
||||
struct Calls {
|
||||
documents: usize,
|
||||
images: usize,
|
||||
audio: usize,
|
||||
}
|
||||
struct CountedGoogle {
|
||||
inner: GemmaEmbedder,
|
||||
calls: Arc<Mutex<Calls>>,
|
||||
}
|
||||
impl EmbeddingProvider for CountedGoogle {
|
||||
fn fingerprint(&self) -> &str {
|
||||
EmbeddingProvider::fingerprint(&self.inner)
|
||||
}
|
||||
fn dimensions(&self) -> usize {
|
||||
EmbeddingProvider::dimensions(&self.inner)
|
||||
}
|
||||
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
|
||||
self.inner.split_windows(text, title)
|
||||
}
|
||||
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
|
||||
self.calls.lock().unwrap().documents += texts.len();
|
||||
self.inner.embed_documents(texts, title)
|
||||
}
|
||||
fn embed_query(&mut self, text: &str) -> Result<Vec<f32>> {
|
||||
self.inner.embed_query(text)
|
||||
}
|
||||
fn embed_code_query(&mut self, text: &str) -> Result<Vec<f32>> {
|
||||
self.inner.embed_code_query(text)
|
||||
}
|
||||
fn embed_images(&mut self, paths: &[&Path], budget: u32) -> Result<Vec<Vec<f32>>> {
|
||||
self.calls.lock().unwrap().images += paths.len();
|
||||
self.inner.embed_images_with_budget(paths, budget)
|
||||
}
|
||||
fn embed_audio(&mut self, pcm: &[f32]) -> Result<Vec<f32>> {
|
||||
self.calls.lock().unwrap().audio += 1;
|
||||
self.inner.embed_audio_pcm(pcm)
|
||||
}
|
||||
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
|
||||
Some(self.inner.info.clone())
|
||||
}
|
||||
}
|
||||
fn google(db: &Path) -> Result<(HybridSearch, Arc<Mutex<Calls>>)> {
|
||||
assert_eq!(std::env::var("SEARCHI_FORCE_CPU")?, "1");
|
||||
let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR")
|
||||
.expect("Verified Google model directory required");
|
||||
let worker =
|
||||
std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Normal auto runtime worker required");
|
||||
let inner = GemmaEmbedder::load(model, worker, 1)?;
|
||||
assert_eq!(inner.info.device, "cpu");
|
||||
let calls = Arc::new(Mutex::new(Calls::default()));
|
||||
let service = HybridSearch::with_provider(
|
||||
db,
|
||||
Box::new(CountedGoogle {
|
||||
inner,
|
||||
calls: calls.clone(),
|
||||
}),
|
||||
)?;
|
||||
Ok((service, calls))
|
||||
}
|
||||
type Chunk = (String, String, i64, i64, i64, String, Vec<u8>);
|
||||
fn chunks(db: &Path) -> Result<Vec<Chunk>> {
|
||||
let connection = Connection::open(db)?;
|
||||
let mut query = connection.prepare("SELECT path,kind,number,start_byte,end_byte,snippet,vector FROM semantic_chunks WHERE generation=(SELECT active_generation FROM semantic_meta WHERE singleton=1) ORDER BY path,ordinal")?;
|
||||
let rows = query.query_map([], |r| {
|
||||
Ok((
|
||||
r.get(0)?,
|
||||
r.get(1)?,
|
||||
r.get(2)?,
|
||||
r.get(3)?,
|
||||
r.get(4)?,
|
||||
r.get(5)?,
|
||||
r.get(6)?,
|
||||
))
|
||||
})?;
|
||||
Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Requires verified Google assets and the normal native worker; CPU only"]
|
||||
fn genuine_google_reuses_text_code_and_remaps_moves_but_reencodes_titles() -> Result<()> {
|
||||
let temp = tempfile::tempdir()?;
|
||||
let root = temp.path().join("sources");
|
||||
fs::create_dir(&root)?;
|
||||
let document = root.join("agreement.txt");
|
||||
fs::write(
|
||||
&document,
|
||||
"The supplier partnership may be ended with sixty days written notice.",
|
||||
)?;
|
||||
fs::write(
|
||||
root.join("retry.rs"),
|
||||
"pub fn exponential_backoff(attempt: u32) -> u64 { 100 * 2_u64.pow(attempt.min(8)) }\n",
|
||||
)?;
|
||||
let db = temp.path().join("index.db");
|
||||
let mut lexical = SearchIndex::open(&db)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let (mut service, calls) = google(&db)?;
|
||||
service.index_embeddings()?;
|
||||
let first = chunks(&db)?;
|
||||
assert_eq!(first.len(), 2);
|
||||
let initial = calls.lock().unwrap().clone();
|
||||
assert_eq!(initial.documents, 2);
|
||||
service.index_embeddings()?;
|
||||
assert_eq!(
|
||||
*calls.lock().unwrap(),
|
||||
initial,
|
||||
"Unchanged groups must use zero inference"
|
||||
);
|
||||
assert_eq!(
|
||||
chunks(&db)?,
|
||||
first,
|
||||
"Real vector bytes and locations must survive reuse"
|
||||
);
|
||||
|
||||
let moved = root.join("nested");
|
||||
fs::create_dir(&moved)?;
|
||||
fs::rename(&document, moved.join("agreement.txt"))?;
|
||||
lexical.index_folder(&root)?;
|
||||
service.index_embeddings()?;
|
||||
assert_eq!(
|
||||
*calls.lock().unwrap(),
|
||||
initial,
|
||||
"Folder-only moves retain the Google title"
|
||||
);
|
||||
let after_move = chunks(&db)?;
|
||||
let original = first
|
||||
.iter()
|
||||
.find(|c| c.0.ends_with("agreement.txt"))
|
||||
.unwrap();
|
||||
let relocated = after_move
|
||||
.iter()
|
||||
.find(|c| c.0.ends_with("agreement.txt"))
|
||||
.unwrap();
|
||||
assert_eq!(relocated.6, original.6);
|
||||
assert!(Path::new(&relocated.0).starts_with(moved.canonicalize()?));
|
||||
let results = service.search_semantic(
|
||||
"How can we stop working with the vendor?",
|
||||
10,
|
||||
&SearchOptions::default(),
|
||||
)?;
|
||||
assert_eq!(results[0].filename, "agreement.txt");
|
||||
assert!(results[0].source.is_some());
|
||||
|
||||
fs::rename(
|
||||
moved.join("agreement.txt"),
|
||||
moved.join("supplier-contract.txt"),
|
||||
)?;
|
||||
lexical.index_folder(&root)?;
|
||||
service.index_embeddings()?;
|
||||
assert_eq!(
|
||||
calls.lock().unwrap().documents,
|
||||
initial.documents + 1,
|
||||
"Basename changes alter Google's document title prompt"
|
||||
);
|
||||
fs::write(
|
||||
moved.join("supplier-contract.txt"),
|
||||
"The supplier partnership may be ended with ninety days written notice.",
|
||||
)?;
|
||||
lexical.index_folder(&root)?;
|
||||
service.index_embeddings()?;
|
||||
assert_eq!(calls.lock().unwrap().documents, initial.documents + 2);
|
||||
let code_before = first.iter().find(|c| c.0.ends_with("retry.rs")).unwrap();
|
||||
let latest = chunks(&db)?;
|
||||
assert_eq!(
|
||||
latest.iter().find(|c| c.0.ends_with("retry.rs")).unwrap(),
|
||||
code_before
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ffmpeg(tools: &searchi_service::FfmpegTools, output: &Path, arguments: &[&str]) -> Result<()> {
|
||||
let result = std::process::Command::new(&tools.ffmpeg)
|
||||
.args([
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"error",
|
||||
"-nostdin",
|
||||
"-y",
|
||||
"-threads",
|
||||
"1",
|
||||
])
|
||||
.args(arguments)
|
||||
.arg(output)
|
||||
.output()?;
|
||||
anyhow::ensure!(
|
||||
result.status.success(),
|
||||
"CPU fixture encoding failed: {}",
|
||||
String::from_utf8_lossy(&result.stderr)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Requires verified Google assets, normal worker and CPU FFmpeg helpers"]
|
||||
fn genuine_google_mixed_media_reuse_checks_bytes_despite_identical_file_metadata() -> Result<()> {
|
||||
let tools = searchi_service::FfmpegTools {
|
||||
ffmpeg: std::env::var_os("SEARCHI_TEST_FFMPEG")
|
||||
.expect("Set SEARCHI_TEST_FFMPEG")
|
||||
.into(),
|
||||
ffprobe: std::env::var_os("SEARCHI_TEST_FFPROBE")
|
||||
.expect("Set SEARCHI_TEST_FFPROBE")
|
||||
.into(),
|
||||
};
|
||||
let temp = tempfile::tempdir()?;
|
||||
let root = temp.path().join("mixed");
|
||||
fs::create_dir(&root)?;
|
||||
fs::write(
|
||||
root.join("agreement.txt"),
|
||||
"The supplier may terminate the partnership with sixty days written notice.",
|
||||
)?;
|
||||
fs::write(
|
||||
root.join("retry.rs"),
|
||||
"pub fn retry_delay(attempt: u32) -> u64 { 100 * 2_u64.pow(attempt.min(8)) }\n",
|
||||
)?;
|
||||
let image = root.join("color.bmp");
|
||||
ffmpeg(
|
||||
&tools,
|
||||
&image,
|
||||
&[
|
||||
"-f",
|
||||
"lavfi",
|
||||
"-i",
|
||||
"color=c=red:s=64x64",
|
||||
"-frames:v",
|
||||
"1",
|
||||
"-c:v",
|
||||
"bmp",
|
||||
],
|
||||
)?;
|
||||
ffmpeg(
|
||||
&tools,
|
||||
&root.join("tone.wav"),
|
||||
&[
|
||||
"-f",
|
||||
"lavfi",
|
||||
"-i",
|
||||
"sine=frequency=440:duration=1",
|
||||
"-c:a",
|
||||
"pcm_s16le",
|
||||
"-ar",
|
||||
"16000",
|
||||
"-ac",
|
||||
"1",
|
||||
],
|
||||
)?;
|
||||
ffmpeg(
|
||||
&tools,
|
||||
&root.join("scene.mp4"),
|
||||
&[
|
||||
"-f",
|
||||
"lavfi",
|
||||
"-i",
|
||||
"color=c=blue:s=64x64:r=2:d=1",
|
||||
"-an",
|
||||
"-c:v",
|
||||
"mpeg4",
|
||||
],
|
||||
)?;
|
||||
let db = temp.path().join("index.db");
|
||||
let mut lexical = SearchIndex::open(&db)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let (mut service, calls) = google(&db)?;
|
||||
let limits = searchi_service::MediaLimits::default();
|
||||
service.index_embeddings_with_media(&tools, &limits)?;
|
||||
let first = chunks(&db)?;
|
||||
for name in [
|
||||
"agreement.txt",
|
||||
"retry.rs",
|
||||
"color.bmp",
|
||||
"tone.wav",
|
||||
"scene.mp4",
|
||||
] {
|
||||
assert!(
|
||||
first.iter().any(|chunk| chunk.0.ends_with(name)),
|
||||
"Missing actual source {name}"
|
||||
);
|
||||
}
|
||||
let initial = calls.lock().unwrap().clone();
|
||||
assert_eq!(initial.documents, 2);
|
||||
assert!(initial.images >= 2 && initial.audio >= 1);
|
||||
service.index_embeddings_with_media(&tools, &limits)?;
|
||||
assert_eq!(
|
||||
*calls.lock().unwrap(),
|
||||
initial,
|
||||
"All five complete source groups must avoid repeated inference"
|
||||
);
|
||||
assert_eq!(chunks(&db)?, first);
|
||||
|
||||
let metadata = fs::metadata(&image)?;
|
||||
let before = fs::read(&image)?;
|
||||
ffmpeg(
|
||||
&tools,
|
||||
&image,
|
||||
&[
|
||||
"-f",
|
||||
"lavfi",
|
||||
"-i",
|
||||
"color=c=green:s=64x64",
|
||||
"-frames:v",
|
||||
"1",
|
||||
"-c:v",
|
||||
"bmp",
|
||||
],
|
||||
)?;
|
||||
assert_eq!(fs::metadata(&image)?.len(), metadata.len());
|
||||
assert_ne!(fs::read(&image)?, before);
|
||||
fs::File::options()
|
||||
.write(true)
|
||||
.open(&image)?
|
||||
.set_times(fs::FileTimes::new().set_modified(metadata.modified()?))?;
|
||||
assert_eq!(fs::metadata(&image)?.modified()?, metadata.modified()?);
|
||||
// Core intentionally sees unchanged size/mtime: semantic reuse must still
|
||||
// hash actual media bytes before accepting a completion manifest.
|
||||
lexical.index_folder(&root)?;
|
||||
service.index_embeddings_with_media(&tools, &limits)?;
|
||||
let changed = calls.lock().unwrap().clone();
|
||||
assert_eq!(changed.documents, initial.documents);
|
||||
assert_eq!(changed.audio, initial.audio);
|
||||
assert_eq!(changed.images, initial.images + 1);
|
||||
let latest = chunks(&db)?;
|
||||
for chunk in first.iter().filter(|chunk| !chunk.0.ends_with("color.bmp")) {
|
||||
assert!(
|
||||
latest.contains(chunk),
|
||||
"Unaffected vector and citation changed"
|
||||
);
|
||||
}
|
||||
assert_ne!(
|
||||
latest
|
||||
.iter()
|
||||
.find(|chunk| chunk.0.ends_with("color.bmp"))
|
||||
.unwrap()
|
||||
.6,
|
||||
first
|
||||
.iter()
|
||||
.find(|chunk| chunk.0.ends_with("color.bmp"))
|
||||
.unwrap()
|
||||
.6
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Requires verified Google assets and normal native worker; CPU only"]
|
||||
fn genuine_google_activation_failure_resumes_completed_vectors_after_reopen() -> Result<()> {
|
||||
let temp = tempfile::tempdir()?;
|
||||
let root = temp.path().join("documents");
|
||||
fs::create_dir(&root)?;
|
||||
let file = root.join("agreement.txt");
|
||||
fs::write(
|
||||
&file,
|
||||
"The supplier agreement ends with sixty days written notice.",
|
||||
)?;
|
||||
let db = temp.path().join("index.db");
|
||||
let mut lexical = SearchIndex::open(&db)?;
|
||||
lexical.index_folder(&root)?;
|
||||
let (mut service, calls) = google(&db)?;
|
||||
service.index_embeddings()?;
|
||||
let connection = Connection::open(&db)?;
|
||||
let active = connection.query_row(
|
||||
"SELECT active_generation FROM semantic_meta WHERE singleton=1",
|
||||
[],
|
||||
|r| r.get::<_, i64>(0),
|
||||
)?;
|
||||
fs::write(
|
||||
&file,
|
||||
"The supplier agreement ends with ninety days written notice.",
|
||||
)?;
|
||||
lexical.index_folder(&root)?;
|
||||
connection.execute_batch("CREATE TRIGGER reject_incremental_activation BEFORE UPDATE OF active_generation ON semantic_meta BEGIN SELECT RAISE(ABORT,'injected incremental activation failure'); END;")?;
|
||||
let failure = service.index_embeddings().unwrap_err();
|
||||
assert!(format!("{failure:#}").contains("injected incremental activation failure"));
|
||||
assert_eq!(
|
||||
connection.query_row(
|
||||
"SELECT active_generation FROM semantic_meta WHERE singleton=1",
|
||||
[],
|
||||
|r| r.get::<_, i64>(0)
|
||||
)?,
|
||||
active
|
||||
);
|
||||
assert_eq!(calls.lock().unwrap().documents, 2);
|
||||
assert!(
|
||||
service
|
||||
.search_semantic("supplier agreement", 10, &SearchOptions::default())?
|
||||
.is_empty(),
|
||||
"Changed old source cannot remain readable"
|
||||
);
|
||||
drop(service);
|
||||
connection.execute_batch("DROP TRIGGER reject_incremental_activation;")?;
|
||||
let (mut reopened, resumed_calls) = google(&db)?;
|
||||
reopened.index_embeddings()?;
|
||||
assert_eq!(
|
||||
resumed_calls.lock().unwrap().documents,
|
||||
0,
|
||||
"Durable completed source must survive reopen after failed activation"
|
||||
);
|
||||
let results = reopened.search_semantic(
|
||||
"How much notice ends the supplier agreement?",
|
||||
10,
|
||||
&SearchOptions::default(),
|
||||
)?;
|
||||
assert_eq!(results.len(), 1);
|
||||
assert!(results[0].snippet.contains("ninety days"));
|
||||
assert!(results[0].source.is_some());
|
||||
Ok(())
|
||||
}
|
||||
Reference in new issue
Block a user