419 lines
16 KiB
Rust
419 lines
16 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_base(policy).into(),
|
|
passage_id: passage.id,
|
|
})
|
|
}
|
|
|
|
impl HybridSearch {
|
|
pub(super) fn rebuild_generation(
|
|
&mut self,
|
|
fingerprint: String,
|
|
media_limits: Option<&searchi_media::MediaLimits>,
|
|
build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
|
|
) -> Result<EmbeddingReport> {
|
|
self.rebuild_generation_guarded(fingerprint, media_limits, false, build)?
|
|
.ok_or_else(|| anyhow::anyhow!("Explicit semantic rebuild was unexpectedly deferred"))
|
|
}
|
|
|
|
pub(super) fn rebuild_generation_preserving_media(
|
|
&mut self,
|
|
fingerprint: String,
|
|
build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
|
|
) -> Result<Option<EmbeddingReport>> {
|
|
self.rebuild_generation_guarded(fingerprint, None, true, build)
|
|
}
|
|
|
|
fn rebuild_generation_guarded(
|
|
&mut self,
|
|
fingerprint: String,
|
|
media_limits: Option<&searchi_media::MediaLimits>,
|
|
preserve_full_media: bool,
|
|
build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
|
|
) -> Result<Option<EmbeddingReport>> {
|
|
let lock = self.store.build_lock()?;
|
|
if preserve_full_media && self.store.has_compatible_full_generation(&fingerprint)? {
|
|
return Ok(None);
|
|
}
|
|
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 = self
|
|
.discover_tasks(media_limits)
|
|
.and_then(|tasks| {
|
|
self.store
|
|
.prune_before_inference(generation, &tasks, revision)?;
|
|
build(self, generation, &mut work)
|
|
})
|
|
.and_then(|()| {
|
|
self.store
|
|
.reconcile_tasks(generation, &work.keys, revision)?;
|
|
self.store
|
|
.activate_checked_with_validation(generation, revision, || {
|
|
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"
|
|
);
|
|
}
|
|
Ok(())
|
|
})?;
|
|
Ok(Some(work.report))
|
|
});
|
|
if result.is_err() {
|
|
self.generation_fingerprint = previous_fingerprint;
|
|
}
|
|
result
|
|
}
|
|
|
|
/// Discover source identities in keyset batches before spending the encoder
|
|
/// budget. This extra byte-verification pass enables safe staging pruning;
|
|
/// it is indexing I/O, not a foreground latency optimization.
|
|
fn discover_tasks(
|
|
&self,
|
|
media_limits: Option<&searchi_media::MediaLimits>,
|
|
) -> Result<Vec<SourceTask>> {
|
|
let mut tasks = Vec::new();
|
|
let mut after = 0;
|
|
loop {
|
|
let passages = self.lexical.passage_batch(after, 64)?;
|
|
if passages.is_empty() {
|
|
break;
|
|
}
|
|
for passage in passages {
|
|
after = passage.id;
|
|
if passage.text.trim().is_empty() {
|
|
continue;
|
|
}
|
|
ensure!(
|
|
tasks.len() < MAX_SCAN_CHUNKS,
|
|
"Semantic source task budget exceeded"
|
|
);
|
|
tasks.push(text_task(&passage, &self.generation_fingerprint)?);
|
|
}
|
|
}
|
|
if let Some(limits) = media_limits {
|
|
let mut after = 0;
|
|
let mut count = 0;
|
|
loop {
|
|
let files = self.store.media_batch(after)?;
|
|
if files.is_empty() {
|
|
break;
|
|
}
|
|
for (id, path, _, _) in files {
|
|
after = id;
|
|
let Some(kind) = searchi_media::classify(Path::new(&path)) else {
|
|
continue;
|
|
};
|
|
count += 1;
|
|
ensure!(
|
|
count <= 10000,
|
|
"Media generation exceeds 10000-file resource budget"
|
|
);
|
|
ensure!(
|
|
tasks.len() < MAX_SCAN_CHUNKS,
|
|
"Semantic source task budget exceeded"
|
|
);
|
|
let identity = searchi_media::MediaIdentity::read(
|
|
Path::new(&path),
|
|
if kind == searchi_media::MediaKind::Image {
|
|
64 * 1024 * 1024
|
|
} else {
|
|
limits.max_file_bytes
|
|
},
|
|
);
|
|
tasks.push(SourceTask {
|
|
key: task_key(&["media", &path]),
|
|
path,
|
|
source_type: "media".into(),
|
|
source_hash: identity.map(|identity| identity.sha256).unwrap_or_default(),
|
|
title_hash: String::new(),
|
|
policy: self.generation_fingerprint.clone(),
|
|
passage_id: -id,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
Ok(tasks)
|
|
}
|
|
|
|
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,
|
|
)?;
|
|
let extracted = self.lexical.extracted_source_hash(&passage.path)?;
|
|
ensure!(extracted.as_deref() == Some(identity.sha256.as_str()),
|
|
"Document bytes differ from extracted passages; reindex the folder before semantic indexing");
|
|
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"
|
|
);
|
|
let batch_limit = self.embedder.document_batch_limit().clamp(1, MAX_BATCH);
|
|
for (batch_index, batch) in windows.chunks(batch_limit).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 * batch_limit,
|
|
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(())
|
|
}
|
|
#[test]
|
|
fn discovery_and_encoding_use_the_same_normalized_text_component_identity() -> Result<()> {
|
|
let temp = tempfile::tempdir()?;
|
|
let root = temp.path().join("docs");
|
|
std::fs::create_dir(&root)?;
|
|
std::fs::write(
|
|
root.join("notes.txt"),
|
|
"Exact source bytes before a deliberately failed encoder.",
|
|
)?;
|
|
let db = temp.path().join("index.db");
|
|
SearchIndex::open(&db)?.index_folder(&root)?;
|
|
let mut service = HybridSearch::with_provider(&db, Box::new(RejectInference))?;
|
|
let base = generation_fingerprint(service.embedder.fingerprint());
|
|
service.generation_fingerprint = format!("{base}:text-only");
|
|
let discovered = service.discover_tasks(None)?.remove(0);
|
|
assert_eq!(discovered.policy, base);
|
|
let error = service
|
|
.index_text_embeddings_preserving_media()
|
|
.unwrap_err();
|
|
assert!(format!("{error:#}").contains("injected primary inference failure"));
|
|
let connection = rusqlite::Connection::open(&db)?;
|
|
let encoded: (String, String, String, String) = connection.query_row(
|
|
"SELECT task_key,source_hash,title_hash,policy FROM semantic_tasks",
|
|
[],
|
|
|row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
|
|
)?;
|
|
assert_eq!(
|
|
encoded,
|
|
(
|
|
discovered.key,
|
|
discovered.source_hash,
|
|
discovered.title_hash,
|
|
discovered.policy
|
|
)
|
|
);
|
|
Ok(())
|
|
}
|
|
}
|