273 lines
10 KiB
Rust
273 lines
10 KiB
Rust
//! Durable bounded semantic source tasks. Only completed source groups may
|
|
//! supply vectors; tokenizer/inference work stays outside SQLite transactions.
|
|
use super::*;
|
|
use searchi_core::IndexedPassage;
|
|
use store::{content_hash, SourceTask};
|
|
|
|
pub(super) struct BuildWork {
|
|
pub report: EmbeddingReport,
|
|
pub keys: Vec<String>,
|
|
pub media: Vec<searchi_media::MediaIdentity>,
|
|
documents: HashMap<String, searchi_media::MediaIdentity>,
|
|
}
|
|
|
|
impl BuildWork {
|
|
pub fn track(&mut self, task: &SourceTask) -> Result<()> {
|
|
ensure!(
|
|
self.keys.len() < MAX_SCAN_CHUNKS,
|
|
"Semantic source task budget exceeded"
|
|
);
|
|
self.keys.push(task.key.clone());
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
pub(super) fn task_key(parts: &[&str]) -> String {
|
|
let mut input = String::new();
|
|
for part in parts {
|
|
input.push_str(&format!("{}:", part.len()));
|
|
input.push_str(part);
|
|
}
|
|
content_hash(&input)
|
|
}
|
|
|
|
pub(super) fn text_task(passage: &IndexedPassage, policy: &str) -> Result<SourceTask> {
|
|
let key = task_key(&[
|
|
"text",
|
|
&passage.path,
|
|
&passage.source.kind,
|
|
&passage.source.number.to_string(),
|
|
]);
|
|
Ok(SourceTask {
|
|
key,
|
|
path: passage.path.clone(),
|
|
source_type: "text".into(),
|
|
source_hash: content_hash(&passage.text),
|
|
title_hash: content_hash(bounded_title(&passage.filename)),
|
|
policy: policy.into(),
|
|
passage_id: passage.id,
|
|
})
|
|
}
|
|
|
|
impl HybridSearch {
|
|
pub(super) fn rebuild_generation(
|
|
&mut self,
|
|
fingerprint: String,
|
|
build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
|
|
) -> Result<EmbeddingReport> {
|
|
let lock = self.store.build_lock()?;
|
|
let (generation, resumed) = self.store.start_or_resume(&fingerprint, &lock)?;
|
|
// This revision covers the entire discovery/build/finalization pass.
|
|
// A concurrent lexical update aborts publication; next invocation
|
|
// reconciles and resumes all still-compatible completed groups.
|
|
let revision = self.store.revision()?;
|
|
let previous_fingerprint = std::mem::replace(&mut self.generation_fingerprint, fingerprint);
|
|
let mut work = BuildWork {
|
|
report: EmbeddingReport {
|
|
media_files: 0,
|
|
media_failures: 0,
|
|
failures: Vec::new(),
|
|
passages: 0,
|
|
chunks: 0,
|
|
generation,
|
|
model_fingerprint: self.generation_fingerprint.clone(),
|
|
reused_sources: 0,
|
|
reused_chunks: 0,
|
|
encoded_sources: 0,
|
|
document_inference_calls: 0,
|
|
image_inference_calls: 0,
|
|
audio_inference_calls: 0,
|
|
resumed,
|
|
},
|
|
keys: Vec::new(),
|
|
media: Vec::new(),
|
|
documents: HashMap::new(),
|
|
};
|
|
let result = build(self, generation, &mut work).and_then(|()| {
|
|
for identity in work.media.iter().chain(work.documents.values()) {
|
|
// Full bytes, not just stat metadata: in-place changes can
|
|
// preserve size and timestamps, including on mounted volumes.
|
|
let current = searchi_media::MediaIdentity::read(&identity.path, identity.bytes)?;
|
|
ensure!(
|
|
current.sha256 == identity.sha256
|
|
&& current.modified_ns == identity.modified_ns,
|
|
"Source changed during semantic indexing; retry required"
|
|
);
|
|
}
|
|
self.store
|
|
.reconcile_tasks(generation, &work.keys, revision)?;
|
|
self.store.activate_checked(generation, revision)?;
|
|
Ok(work.report)
|
|
});
|
|
if result.is_err() {
|
|
self.generation_fingerprint = previous_fingerprint;
|
|
}
|
|
result
|
|
}
|
|
|
|
pub(super) fn build(&mut self, generation: i64, work: &mut BuildWork) -> Result<()> {
|
|
let mut after_id = 0;
|
|
loop {
|
|
let passages = self.lexical.passage_batch(after_id, 64)?;
|
|
if passages.is_empty() {
|
|
break;
|
|
}
|
|
for passage in passages {
|
|
after_id = passage.id;
|
|
if passage.text.trim().is_empty() {
|
|
continue;
|
|
}
|
|
if !work.documents.contains_key(&passage.path) {
|
|
ensure!(
|
|
work.documents.len() < MAX_SCAN_CHUNKS,
|
|
"Document source budget exceeded"
|
|
);
|
|
let identity = searchi_media::MediaIdentity::read(
|
|
Path::new(&passage.path),
|
|
32 * 1024 * 1024,
|
|
)?;
|
|
work.documents.insert(passage.path.clone(), identity);
|
|
}
|
|
let task = text_task(&passage, &self.generation_fingerprint)?;
|
|
work.track(&task)?;
|
|
self.store.stage_task(generation, &task)?;
|
|
let result = self.build_text_source(generation, &task, &passage, work);
|
|
if let Err(error) = result {
|
|
if let Err(cleanup) =
|
|
self.store
|
|
.fail_source(generation, &task, &format!("{error:#}"))
|
|
{
|
|
return Err(error.context(format!(
|
|
"Failed to discard incomplete text source: {cleanup:#}"
|
|
)));
|
|
}
|
|
return Err(error);
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn build_text_source(
|
|
&mut self,
|
|
generation: i64,
|
|
task: &SourceTask,
|
|
passage: &IndexedPassage,
|
|
work: &mut BuildWork,
|
|
) -> Result<()> {
|
|
if let Some(count) = self.store.reuse_text_source(generation, task, passage)? {
|
|
ensure!(
|
|
work.report.chunks + count <= MAX_SCAN_CHUNKS,
|
|
"Corpus exceeds ANN chunk budget"
|
|
);
|
|
work.report.chunks += count;
|
|
work.report.passages += 1;
|
|
work.report.reused_chunks += count;
|
|
work.report.reused_sources += 1;
|
|
return Ok(());
|
|
}
|
|
let title = bounded_title(&passage.filename);
|
|
let windows = passage_windows(&mut *self.embedder, &passage.text, title)?;
|
|
ensure!(
|
|
work.report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
|
|
"Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
|
|
);
|
|
for (batch_index, batch) in windows.chunks(MAX_BATCH).enumerate() {
|
|
let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
|
|
work.report.document_inference_calls += 1;
|
|
let vectors = self.embedder.embed_documents(&texts, title)?;
|
|
self.store.insert(
|
|
generation,
|
|
passage,
|
|
batch_index * MAX_BATCH,
|
|
batch,
|
|
&vectors,
|
|
)?;
|
|
}
|
|
self.store
|
|
.complete_source(generation, task, windows.len())?;
|
|
work.report.chunks += windows.len();
|
|
work.report.passages += 1;
|
|
work.report.encoded_sources += 1;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// Fault-only provider: never returns a vector or claims semantic quality.
|
|
struct RejectInference;
|
|
impl EmbeddingProvider for RejectInference {
|
|
fn fingerprint(&self) -> &str {
|
|
"fault-injection-only"
|
|
}
|
|
fn dimensions(&self) -> usize {
|
|
GEMMA_DIMENSIONS
|
|
}
|
|
fn split_windows(&mut self, text: &str, _: &str) -> Result<Vec<TextWindow>> {
|
|
Ok(vec![TextWindow {
|
|
text: text.into(),
|
|
start_byte: 0,
|
|
end_byte: text.len(),
|
|
}])
|
|
}
|
|
fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
|
|
anyhow::bail!("injected primary inference failure")
|
|
}
|
|
fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
|
|
anyhow::bail!("fault provider never supplies query vectors")
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cleanup_failure_keeps_primary_error_reader_policy_and_staging() -> Result<()> {
|
|
let temp = tempfile::tempdir()?;
|
|
let root = temp.path().join("corpus");
|
|
std::fs::create_dir(&root)?;
|
|
std::fs::write(
|
|
root.join("notes.txt"),
|
|
"A real source whose inference is deliberately rejected.",
|
|
)?;
|
|
let db = temp.path().join("index.db");
|
|
SearchIndex::open(&db)?.index_folder(&root)?;
|
|
let mut service = HybridSearch::with_provider(&db, Box::new(RejectInference))?;
|
|
let prior_policy = service.generation_fingerprint.clone();
|
|
let connection = rusqlite::Connection::open(&db)?;
|
|
connection.execute_batch("CREATE TRIGGER reject_task_cleanup BEFORE UPDATE OF state ON semantic_tasks WHEN new.state='failed' BEGIN SELECT RAISE(ABORT,'injected task cleanup failure'); END;")?;
|
|
let error = service.index_embeddings().unwrap_err();
|
|
let message = format!("{error:#}");
|
|
assert!(
|
|
message.contains("injected primary inference failure"),
|
|
"{message}"
|
|
);
|
|
assert!(
|
|
message.contains("injected task cleanup failure"),
|
|
"{message}"
|
|
);
|
|
assert_eq!(service.generation_fingerprint, prior_policy);
|
|
let staging: usize =
|
|
connection.query_row("SELECT count(*) FROM semantic_generations", [], |r| {
|
|
r.get(0)
|
|
})?;
|
|
assert_eq!(staging, 1);
|
|
let state: String =
|
|
connection.query_row("SELECT state FROM semantic_tasks", [], |r| r.get(0))?;
|
|
assert_eq!(
|
|
state, "running",
|
|
"Failed cleanup transaction must roll back"
|
|
);
|
|
connection.execute_batch("DROP TRIGGER reject_task_cleanup;")?;
|
|
let error = service.index_embeddings().unwrap_err();
|
|
assert!(format!("{error:#}").contains("injected primary inference failure"));
|
|
assert_eq!(service.generation_fingerprint, prior_policy);
|
|
let staging: usize =
|
|
connection.query_row("SELECT count(*) FROM semantic_generations", [], |r| {
|
|
r.get(0)
|
|
})?;
|
|
assert_eq!(staging, 1, "Resume must retain the same staging generation");
|
|
Ok(())
|
|
}
|
|
}
|