diff --git a/Cargo.toml b/Cargo.toml index 492a2b8..b3ab10a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,19 +29,9 @@ unsafe_code = "forbid" [lints.clippy] all = "deny" - [[bench]] name = "lifecycle_deletion_repairs" -harness = false [[bench]] name = "claims_nplus1" harness = false - -[[bench]] -name = "store_review_evidence_nplus1" -harness = false - -[[bench]] -name = "validate_evidence" -harness = false diff --git a/benches/store_review_evidence_nplus1.rs b/benches/store_review_evidence_nplus1.rs deleted file mode 100644 index fcab2a0..0000000 --- a/benches/store_review_evidence_nplus1.rs +++ /dev/null @@ -1,98 +0,0 @@ -use rusqlite::{Connection, params}; - -use std::time::Instant; - -fn main() -> Result<(), rusqlite::Error> { - let mut connection = Connection::open_in_memory()?; - - connection.execute_batch( - "CREATE TABLE evidence( - id TEXT, - tenant_id TEXT, - person_id TEXT, - deleted_at INTEGER, - PRIMARY KEY (id, tenant_id, person_id) - );", - )?; - - let num_items = 1000; - let mut ids = Vec::new(); - - let tx = connection.transaction()?; - let mut stmt = tx.prepare("INSERT INTO evidence VALUES(?, 't1', 'p1', NULL)")?; - for i in 0..num_items { - stmt.execute(params![i.to_string()])?; - ids.push(i.to_string()); - } - drop(stmt); - tx.commit()?; - - // N+1 baseline (Happy path) - let start_nplus1 = Instant::now(); - let mut read_stmt = connection.prepare("SELECT EXISTS(SELECT 1 FROM evidence WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL)")?; - - let mut missing_ids_nplus1 = Vec::new(); - for id in &ids { - let found: bool = read_stmt.query_row(params![id, "t1", "p1"], |row| row.get(0))?; - if !found { - missing_ids_nplus1.push(id.clone()); - } - } - let nplus1_duration = start_nplus1.elapsed(); - - // Using EXCEPT approach - let start_except = Instant::now(); - let ids_json = serde_json::to_string(&ids).unwrap(); - let mut missing_stmt = connection.prepare( - "SELECT value FROM json_each(?1) \ - EXCEPT \ - SELECT id FROM evidence WHERE tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL", - )?; - let missing_ids: Vec = missing_stmt - .query_map(params![ids_json, "t1", "p1"], |row| row.get(0))? - .collect::>()?; - let except_duration = start_except.elapsed(); - - assert_eq!(missing_ids_nplus1.len(), 0); - assert_eq!(missing_ids.len(), 0); - - println!("N+1 Duration (Happy Path): {:?}", nplus1_duration); - println!("EXCEPT Duration (Happy Path): {:?}", except_duration); - - // Bench with one missing - let mut ids_with_missing = ids.clone(); - ids_with_missing.push("missing_id".to_string()); - - // N+1 baseline (Missing) - let start_nplus1_miss = Instant::now(); - let mut missing_ids_nplus1_miss = Vec::new(); - for id in &ids_with_missing { - let found: bool = read_stmt.query_row(params![id, "t1", "p1"], |row| row.get(0))?; - if !found { - missing_ids_nplus1_miss.push(id.clone()); - break; // Stop at first miss, like the code - } - } - let nplus1_duration_miss = start_nplus1_miss.elapsed(); - - // Using EXCEPT approach (Missing) - let start_except_miss = Instant::now(); - let ids_json_miss = serde_json::to_string(&ids_with_missing).unwrap(); - let mut missing_ids_miss = missing_stmt - .query_map(params![ids_json_miss, "t1", "p1"], |row| { - row.get::<_, String>(0) - })?; - let mut found_missing = false; - if let Some(_missing_id) = missing_ids_miss.next() { - found_missing = true; - } - let except_duration_miss = start_except_miss.elapsed(); - - assert!(missing_ids_nplus1_miss.contains(&"missing_id".to_string())); - assert!(found_missing); - - println!("N+1 Duration (1 Missing): {:?}", nplus1_duration_miss); - println!("EXCEPT Duration (1 Missing): {:?}", except_duration_miss); - - Ok(()) -} diff --git a/benches/validate_evidence.rs b/benches/validate_evidence.rs deleted file mode 100644 index f963570..0000000 --- a/benches/validate_evidence.rs +++ /dev/null @@ -1,81 +0,0 @@ -use rusqlite::{Connection, params}; -use std::time::Instant; - -fn main() -> Result<(), rusqlite::Error> { - let mut connection = Connection::open_in_memory()?; - - connection.execute_batch( - "CREATE TABLE evidence( - id TEXT, - tenant_id TEXT, - person_id TEXT, - deleted_at TEXT, - PRIMARY KEY (id, tenant_id, person_id) - );", - )?; - - let num_items = 900; // Chunk size is 900 - - // Insert test data (leave the last one out to trigger the error path) - let tx = connection.transaction()?; - let mut stmt = - tx.prepare("INSERT INTO evidence (id, tenant_id, person_id) VALUES(?, 't1', 'p1')")?; - for i in 0..(num_items - 1) { - stmt.execute(params![i.to_string()])?; - } - drop(stmt); - tx.commit()?; - - let chunk: Vec = (0..num_items).map(|i| i.to_string()).collect(); - - // N+1 baseline - let start_nplus1 = Instant::now(); - for _ in 0..100 { - // Loop multiple times to measure small differences - for id in &chunk { - let live: bool = connection.query_row( - "SELECT EXISTS(SELECT 1 FROM evidence WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL)", - params![id, "t1", "p1"], - |row| row.get(0), - )?; - if !live { - // println!("missing: {}", id); - break; - } - } - } - let nplus1_duration = start_nplus1.elapsed(); - - // Single Query optimized - let start_optimized = Instant::now(); - for _ in 0..100 { - let json_arr = serde_json::to_string(&chunk).unwrap(); - - // This is SQLite's json_each to find missing elements - let mut stmt = connection.prepare( - "SELECT value - FROM json_each(?1) - WHERE NOT EXISTS ( - SELECT 1 FROM evidence - WHERE id = value - AND tenant_id = ?2 - AND person_id = ?3 - AND deleted_at IS NULL - ) - LIMIT 1", - )?; - - let missing_id: rusqlite::Result = - stmt.query_row(params![json_arr, "t1", "p1"], |row| row.get(0)); - - if let Ok(_id) = missing_id { - // Found a missing element - } - } - let optimized_duration = start_optimized.elapsed(); - - println!("N+1 Duration: {:?}", nplus1_duration); - println!("Optimized (json_each) Duration: {:?}", optimized_duration); - - Ok(()) -} diff --git a/plugins/openclaw/cli.ts b/plugins/openclaw/cli.ts index 28e6bf8..b909b9e 100644 --- a/plugins/openclaw/cli.ts +++ b/plugins/openclaw/cli.ts @@ -86,15 +86,9 @@ export async function runZkr( child.stdout.on("data", (chunk: Buffer) => capture(output.stdout, chunk)); child.stderr.on("data", (chunk: Buffer) => capture(output.stderr, chunk)); - child.stdin.on("error", () => { - // Ignore EPIPE errors which occur when the child process closes its stdin before we finish writing - }); child.on("error", () => { fail(false); }); - child.stdin.on("error", () => { - // Ignore EPIPE errors if the process closes stdin early - }); child.on("close", (code) => { if (settled) return; clearTimeout(timeout); diff --git a/src/main.rs b/src/main.rs index f081250..12d7f74 100644 --- a/src/main.rs +++ b/src/main.rs @@ -45,15 +45,7 @@ fn run() -> Result, Box> { return Err("usage: zkr --db PATH COMMAND (use --help)".into()); } let mut database = MemoryDb::open(&arguments[1])?; - let value = dispatch_command(&mut database, arguments[2].as_str())?; - Ok(Some(value)) -} - -fn dispatch_command( - database: &mut MemoryDb, - command: &str, -) -> Result> { - let value = match command { + let value = match arguments[2].as_str() { "remember" => { let request = read_json::()?; serde_json::to_value(database.remember_with_locator(request.memory, request.locator)?)? @@ -93,7 +85,7 @@ fn dispatch_command( "apply" => serde_json::to_value(database.apply(read_json::()?)?)?, command => return Err(format!("unknown command {command:?}").into()), }; - Ok(value) + Ok(Some(value)) } fn read_json() -> Result> { diff --git a/src/model.rs b/src/model.rs index f94d470..d1c9e17 100644 --- a/src/model.rs +++ b/src/model.rs @@ -481,37 +481,4 @@ mod tests { _ => panic!("Expected EmptyText error for evidence id"), } } - - #[test] - fn test_validate_text_valid() { - assert_eq!(validate_text("test_field", "valid text"), Ok(())); - assert_eq!(validate_text("test_field", " leading whitespace"), Ok(())); - assert_eq!(validate_text("test_field", "trailing whitespace "), Ok(())); - assert_eq!(validate_text("test_field", " both "), Ok(())); - assert_eq!(validate_text("test_field", "a"), Ok(())); - } - - #[test] - fn test_validate_text_invalid() { - assert_eq!( - validate_text("test_field", ""), - Err(ValidationError::EmptyText("test_field")) - ); - assert_eq!( - validate_text("test_field", " "), - Err(ValidationError::EmptyText("test_field")) - ); - assert_eq!( - validate_text("test_field", "\t"), - Err(ValidationError::EmptyText("test_field")) - ); - assert_eq!( - validate_text("test_field", "\n"), - Err(ValidationError::EmptyText("test_field")) - ); - assert_eq!( - validate_text("test_field", " \t\n "), - Err(ValidationError::EmptyText("test_field")) - ); - } } diff --git a/src/personality.rs b/src/personality.rs index d9a1495..428187a 100644 --- a/src/personality.rs +++ b/src/personality.rs @@ -2095,17 +2095,4 @@ mod tests { .unwrap(); assert_eq!(augmented, "Base prompt."); } - - #[test] - fn search_personality_propagates_db_error() { - let tmp = tempfile::tempdir().unwrap(); - let db = MemoryDb::open(tmp.path().join("personality.db")).unwrap(); - // Use an empty TenantId, which makes db.search fail validation - let tenant_id = TenantId("".into()); - let person_id = PersonId("p1".into()); - let personality = Personality::new(db, tenant_id, person_id); - - let result = personality.search_personality("query", 5); - assert!(result.is_err()); - } } diff --git a/src/store/apply.rs b/src/store/apply.rs index 1c25109..cd5e3e3 100644 --- a/src/store/apply.rs +++ b/src/store/apply.rs @@ -450,20 +450,6 @@ fn apply_evidence( Ok(()) } -struct StoredClaim { - subject: String, - predicate: String, - value: String, - kind: String, - valid_from: Timestamp, - recorded_from: Timestamp, - valid_until: Option, - recorded_until: Option, - status: String, - tier: String, - processing_state: String, -} - fn apply_claim(transaction: &Transaction<'_>, record: &Claim, applied_at: Timestamp) -> Result<()> { assert_legal_state(&record.tier, &record.status, &record.processing_state) .map_err(|error| Error::Invalid(error.to_string()))?; @@ -476,19 +462,19 @@ fn apply_claim(transaction: &Transaction<'_>, record: &Claim, applied_at: Timest "SELECT subject, predicate, value, kind, valid_from, recorded_from, valid_until, recorded_until, status, tier, processing_state FROM claims WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3", params![record.id.0, record.tenant_id.0, record.person_id.0], |row| { - Ok(StoredClaim { - subject: row.get::<_, String>(0)?, - predicate: row.get::<_, String>(1)?, - value: row.get::<_, String>(2)?, - kind: row.get::<_, String>(3)?, - valid_from: row.get::<_, Timestamp>(4)?, - recorded_from: row.get::<_, Timestamp>(5)?, - valid_until: row.get::<_, Option>(6)?, - recorded_until: row.get::<_, Option>(7)?, - status: row.get::<_, String>(8)?, - tier: row.get::<_, String>(9)?, - processing_state: row.get::<_, String>(10)?, - }) + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, String>(2)?, + row.get::<_, String>(3)?, + row.get::<_, Timestamp>(4)?, + row.get::<_, Timestamp>(5)?, + row.get::<_, Option>(6)?, + row.get::<_, Option>(7)?, + row.get::<_, String>(8)?, + row.get::<_, String>(9)?, + row.get::<_, String>(10)?, + )) }, ) .optional()?; @@ -499,23 +485,23 @@ fn apply_claim(transaction: &Transaction<'_>, record: &Claim, applied_at: Timest ).map_err(claim_interval_error)?; return Ok(()); }; - if stored.subject != record.subject - || stored.predicate != record.predicate - || stored.value != record.value - || stored.kind != kind - || stored.valid_from != record.valid_time.from - || stored.recorded_from != record.recorded_time.from + if stored.0 != record.subject + || stored.1 != record.predicate + || stored.2 != record.value + || stored.3 != kind + || stored.4 != record.valid_time.from + || stored.5 != record.recorded_time.from { return Err(Error::Invalid(format!( "applied claim {} conflicts with the stored claim payload", record.id.0 ))); } - if stored.valid_until == record.valid_time.until - && stored.recorded_until == record.recorded_time.until - && stored.status == status - && stored.tier == tier - && stored.processing_state == processing_state + if stored.6 == record.valid_time.until + && stored.7 == record.recorded_time.until + && stored.8 == status + && stored.9 == tier + && stored.10 == processing_state { return Ok(()); } diff --git a/src/store/embeddings.rs b/src/store/embeddings.rs index e651356..8a997be 100644 --- a/src/store/embeddings.rs +++ b/src/store/embeddings.rs @@ -76,25 +76,37 @@ pub(super) fn projection_input_from( target: EmbeddingTarget, ) -> Result { require_scope(tenant_id, person_id)?; - let (text, target_revision) = match &target { - EmbeddingTarget::Source(id) => connection.query_row( - "SELECT content, revision FROM sources WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL", - params![id.0, tenant_id.0, person_id.0], - |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + let (table, id, expression, revision, live) = match &target { + EmbeddingTarget::Source(id) => ( + "sources", + &id.0, + "content", + "revision", + "deleted_at IS NULL", ), - EmbeddingTarget::Evidence(id) => connection.query_row( - "SELECT quote, source_revision FROM evidence WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL", - params![id.0, tenant_id.0, person_id.0], - |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + EmbeddingTarget::Evidence(id) => ( + "evidence", + &id.0, + "quote", + "source_revision", + "deleted_at IS NULL", ), - EmbeddingTarget::Claim(id) => connection.query_row( - "SELECT subject || ' ' || predicate || ' ' || value, recorded_from FROM claims WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND status = 'accepted' AND valid_until IS NULL AND recorded_until IS NULL AND tier IN ('short_term', 'long_term') AND processing_state = 'processed'", - params![id.0, tenant_id.0, person_id.0], - |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + EmbeddingTarget::Claim(id) => ( + "claims", + &id.0, + "subject || ' ' || predicate || ' ' || value", + "recorded_from", + "status = 'accepted' AND valid_until IS NULL AND recorded_until IS NULL AND tier IN ('short_term', 'long_term') AND processing_state = 'processed'", ), - } - .optional()? - .ok_or(Error::NotFound)?; + }; + let (text, target_revision) = connection + .query_row( + &format!("SELECT {expression}, {revision} FROM {table} WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND {live}"), + params![id, tenant_id.0, person_id.0], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?)), + ) + .optional()? + .ok_or(Error::NotFound)?; Ok(ProjectionInput { target, input_hash: input_hash(&text), @@ -342,6 +354,7 @@ impl MemoryDb { }, )?; let mut scores = HashMap::::new(); + let mut candidate_scores: HashMap<(String, String), f32> = HashMap::new(); for row in rows { let ( target_kind, @@ -408,18 +421,27 @@ impl MemoryDb { ))); } let score = vector_score(&query.vector, &vector, &configuration.2)?; - for target in self.retrieval_targets_for_embedding( - tenant_id, - person_id, - &target_kind, - &target_id, - )? { - scores - .entry(target) - .and_modify(|existing| *existing = existing.max(score)) - .or_insert(score); + candidate_scores + .entry((target_kind, target_id)) + .and_modify(|existing| *existing = existing.max(score)) + .or_insert(score); + } + + let valid_candidates: Vec<(String, String)> = candidate_scores.keys().cloned().collect(); + let bulk_targets = + self.retrieval_targets_for_embeddings_bulk(tenant_id, person_id, &valid_candidates)?; + + for ((kind, id), score) in candidate_scores { + if let Some(targets) = bulk_targets.get(&(kind, id)) { + for target in targets { + scores + .entry(target.clone()) + .and_modify(|existing| *existing = existing.max(score)) + .or_insert(score); + } } } + let mut ranked = scores.into_iter().collect::>(); ranked.sort_by(|left, right| { right diff --git a/src/store/lifecycle.rs b/src/store/lifecycle.rs index fd67142..d74c278 100644 --- a/src/store/lifecycle.rs +++ b/src/store/lifecycle.rs @@ -211,68 +211,152 @@ impl MemoryDb { require_scope(&input.tenant_id, &input.person_id)?; require_text("correction text", &input.text)?; require_text("value", &input.value)?; - let transaction = self.connection.transaction()?; - - let (removed_profiles, old_claim, stale_evidence_ids) = find_correction_targets( - &transaction, - &input.tenant_id, - &input.person_id, - &input.claim_id, + let removed_profiles = transaction + .prepare( + "SELECT id FROM profile_entries WHERE tenant_id = ?1 AND person_id = ?2 AND claim_id = ?3 ORDER BY id", + )? + .query_map( + params![input.tenant_id.0, input.person_id.0, input.claim_id.0], + |row| row.get::<_, String>(0), + )? + .collect::, _>>()?; + let old = transaction + .query_row( + "SELECT subject, predicate, kind, valid_from, recorded_from FROM claims WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND status = 'accepted'", + params![input.claim_id.0, input.tenant_id.0, input.person_id.0], + |row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?, row.get::<_, String>(2)?, row.get::<_, i64>(3)?, row.get::<_, i64>(4)?)), + ) + .optional()? + .ok_or(Error::NotFound)?; + let stale_evidence_ids = transaction + .prepare( + "SELECT evidence_id FROM claim_evidence WHERE tenant_id = ?1 AND person_id = ?2 AND claim_id = ?3 AND relation = '\"supports\"' ORDER BY evidence_id", + )? + .query_map( + params![input.tenant_id.0, input.person_id.0, input.claim_id.0], + |row| row.get::<_, String>(0).map(EvidenceId), + )? + .collect::, _>>()?; + if input.valid_at <= old.3 || input.recorded_at <= old.4 { + return Err(Error::Invalid( + "correction timestamps must advance the original valid and recorded intervals" + .to_owned(), + )); + } + let source_id = SourceId(new_id(&transaction)?); + let evidence_id = EvidenceId(new_id(&transaction)?); + transaction.execute( + "INSERT INTO sources(id, tenant_id, person_id, revision, kind, content, captured_at, recorded_at) VALUES(?1, ?2, ?3, 1, '\"user_correction\"', ?4, ?5, ?6)", + params![source_id.0, input.tenant_id.0, input.person_id.0, input.text, input.valid_at, input.recorded_at], )?; - - validate_correction_time( - input.valid_at, - input.recorded_at, - old_claim.valid_from, - old_claim.recorded_from, + transaction.execute( + "INSERT INTO source_fts(source_id, tenant_id, person_id, content) VALUES(?1, ?2, ?3, ?4)", + params![source_id.0, input.tenant_id.0, input.person_id.0, input.text], )?; - - let (source_id, evidence_id) = insert_correction_source_and_evidence( - &transaction, - &input.tenant_id, - &input.person_id, - &input.text, - input.valid_at, - input.recorded_at, + transaction.execute( + "INSERT INTO evidence(id, tenant_id, person_id, source_id, source_revision, quote, recorded_at) VALUES(?1, ?2, ?3, ?4, 1, ?5, ?6)", + params![evidence_id.0, input.tenant_id.0, input.person_id.0, source_id.0, input.text, input.recorded_at], )?; - - supersede_stale_claim( + transaction.execute( + "UPDATE claims SET status = 'superseded', valid_until = ?1, recorded_until = ?2 WHERE id = ?3 AND tenant_id = ?4 AND person_id = ?5", + params![input.valid_at, input.recorded_at, input.claim_id.0, input.tenant_id.0, input.person_id.0], + )?; + invalidate_summaries_for_evidence( &transaction, &input.tenant_id, &input.person_id, - &input.claim_id, &stale_evidence_ids, - input.valid_at, input.recorded_at, )?; - - let (claim_id, correction) = apply_correction_claim( + enqueue_projection_repair( &transaction, &input.tenant_id, &input.person_id, - &input.claim_id, - &old_claim, - &input.value, - input.valid_at, + EmbeddingTarget::Claim(input.claim_id.clone()), + "superseded_sync", input.recorded_at, - &source_id, - &evidence_id, )?; - - let records = build_correction_records( + transaction.execute( + "DELETE FROM profile_entries WHERE tenant_id = ?1 AND person_id = ?2 AND claim_id = ?3", + params![input.tenant_id.0, input.person_id.0, input.claim_id.0], + )?; + let claim_id = insert_claim( &transaction, &input.tenant_id, &input.person_id, - &input.claim_id, - &source_id, &evidence_id, - &claim_id, - correction, - removed_profiles, + ClaimInput { + subject: old.0, + predicate: old.1, + value: input.value, + kind: claim_kind(&old.2)?, + valid_from: input.valid_at, + tier: MemoryTier::LongTerm, + processing_state: MemoryProcessingState::Processed, + }, input.recorded_at, )?; - + transaction.execute( + "UPDATE sources SET origin_evidence_id = ?1, origin_claim_id = ?2 WHERE id = ?3 AND tenant_id = ?4 AND person_id = ?5", + params![evidence_id.0, claim_id.0, source_id.0, input.tenant_id.0, input.person_id.0], + )?; + let correction = CorrectionRecord { + tenant_id: input.tenant_id.clone(), + person_id: input.person_id.clone(), + superseded_claim_id: input.claim_id.clone(), + claim_id: claim_id.clone(), + source_id: source_id.clone(), + evidence_id: evidence_id.clone(), + valid_at: input.valid_at, + recorded_at: input.recorded_at, + }; + transaction.execute( + "INSERT INTO corrections(tenant_id, person_id, superseded_claim_id, claim_id, source_id, evidence_id, valid_at, recorded_at) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![correction.tenant_id.0, correction.person_id.0, correction.superseded_claim_id.0, correction.claim_id.0, correction.source_id.0, correction.evidence_id.0, correction.valid_at, correction.recorded_at], + )?; + let mut records = vec![ + ExportRecord::Source(source_record( + &transaction, + &input.tenant_id, + &input.person_id, + &source_id, + )?), + ExportRecord::Evidence(evidence_record( + &transaction, + &input.tenant_id, + &input.person_id, + &evidence_id, + )?), + ExportRecord::Claim(claim_record( + &transaction, + &input.tenant_id, + &input.person_id, + &input.claim_id, + )?), + ExportRecord::Claim(claim_record( + &transaction, + &input.tenant_id, + &input.person_id, + &claim_id, + )?), + ExportRecord::ClaimEvidence(claim_evidence_record( + &transaction, + &input.tenant_id, + &input.person_id, + &claim_id, + &evidence_id, + )?), + ExportRecord::Correction(correction), + ]; + records.extend(removed_profiles.into_iter().map(|id| { + ExportRecord::Deletion(DeletionRecord { + tenant_id: input.tenant_id.clone(), + person_id: input.person_id.clone(), + target: MemoryRef::ProfileEntry(ProfileEntryId(id)), + deleted_at: input.recorded_at, + }) + })); append_commit( &transaction, &input.tenant_id, @@ -280,9 +364,7 @@ impl MemoryDb { input.recorded_at, records, )?; - transaction.commit()?; - Ok(Corrected { source_id, evidence_id, @@ -290,6 +372,7 @@ impl MemoryDb { superseded_claim_id: input.claim_id, }) } + pub fn delete_source(&mut self, input: DeleteInput) -> Result { require_scope(&input.tenant_id, &input.person_id)?; let transaction = self.connection.transaction()?; @@ -682,30 +765,19 @@ impl MemoryDb { return Err(Error::Invalid("review needs evidence_ids".to_owned())); } let transaction = self.connection.transaction()?; - let evidence_ids_json = serde_json::to_string(&input.evidence_ids)?; - - let missing_id = { - let mut missing_stmt = transaction.prepare_cached( - "SELECT value FROM json_each(?1) \ - EXCEPT \ - SELECT id FROM evidence WHERE tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL", - )?; - let mut missing_ids = missing_stmt.query_map( - params![evidence_ids_json, input.tenant_id.0, input.person_id.0], - |row| row.get::<_, String>(0), + for evidence_id in &input.evidence_ids { + let found: bool = transaction.query_row( + "SELECT EXISTS(SELECT 1 FROM evidence WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL)", + params![evidence_id.0, input.tenant_id.0, input.person_id.0], + |row| row.get(0), )?; - - if let Some(missing_id_result) = missing_ids.next() { - Some(missing_id_result?) - } else { - None + if !found { + return Err(Error::Invalid(format!( + "evidence {} is unavailable", + evidence_id.0 + ))); } - }; - - if let Some(id) = missing_id { - return Err(Error::Invalid(format!("evidence {} is unavailable", id))); } - let id = DailyReviewId(new_id(&transaction)?); transaction.execute( "INSERT INTO daily_reviews(id, tenant_id, person_id, day, summary, evidence_ids, recorded_at) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7)", @@ -863,229 +935,6 @@ fn new_id(transaction: &Transaction<'_>) -> Result { Ok(transaction.query_row("SELECT lower(hex(randomblob(16)))", [], |row| row.get(0))?) } -struct OldClaim { - subject: String, - predicate: String, - kind: String, - valid_from: i64, - recorded_from: i64, -} - -#[allow(clippy::type_complexity)] -fn find_correction_targets( - transaction: &rusqlite::Transaction<'_>, - tenant_id: &TenantId, - person_id: &PersonId, - claim_id: &ClaimId, -) -> Result<(Vec, OldClaim, Vec)> { - let removed_profiles = transaction - .prepare( - "SELECT id FROM profile_entries WHERE tenant_id = ?1 AND person_id = ?2 AND claim_id = ?3 ORDER BY id", - )? - .query_map( - params![tenant_id.0, person_id.0, claim_id.0], - |row| row.get::<_, String>(0), - )? - .collect::, _>>()?; - let old = transaction - .query_row( - "SELECT subject, predicate, kind, valid_from, recorded_from FROM claims WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND status = 'accepted'", - params![claim_id.0, tenant_id.0, person_id.0], - |row| Ok(OldClaim { - subject: row.get(0)?, - predicate: row.get(1)?, - kind: row.get(2)?, - valid_from: row.get(3)?, - recorded_from: row.get(4)?, - }), - ) - .optional()? - .ok_or(Error::NotFound)?; - let stale_evidence_ids = transaction - .prepare( - "SELECT evidence_id FROM claim_evidence WHERE tenant_id = ?1 AND person_id = ?2 AND claim_id = ?3 AND relation = '\"supports\"' ORDER BY evidence_id", - )? - .query_map( - params![tenant_id.0, person_id.0, claim_id.0], - |row| row.get::<_, String>(0).map(EvidenceId), - )? - .collect::, _>>()?; - Ok((removed_profiles, old, stale_evidence_ids)) -} - -fn validate_correction_time( - input_valid_at: i64, - input_recorded_at: i64, - old_valid_from: i64, - old_recorded_from: i64, -) -> Result<()> { - if input_valid_at <= old_valid_from || input_recorded_at <= old_recorded_from { - return Err(Error::Invalid( - "correction timestamps must advance the original valid and recorded intervals" - .to_owned(), - )); - } - Ok(()) -} - -fn insert_correction_source_and_evidence( - transaction: &rusqlite::Transaction<'_>, - tenant_id: &TenantId, - person_id: &PersonId, - text: &str, - valid_at: i64, - recorded_at: i64, -) -> Result<(SourceId, EvidenceId)> { - let source_id = SourceId(new_id(transaction)?); - let evidence_id = EvidenceId(new_id(transaction)?); - transaction.execute( - "INSERT INTO sources(id, tenant_id, person_id, revision, kind, content, captured_at, recorded_at) VALUES(?1, ?2, ?3, 1, '\"user_correction\"', ?4, ?5, ?6)", - params![source_id.0, tenant_id.0, person_id.0, text, valid_at, recorded_at], - )?; - transaction.execute( - "INSERT INTO source_fts(source_id, tenant_id, person_id, content) VALUES(?1, ?2, ?3, ?4)", - params![source_id.0, tenant_id.0, person_id.0, text], - )?; - transaction.execute( - "INSERT INTO evidence(id, tenant_id, person_id, source_id, source_revision, quote, recorded_at) VALUES(?1, ?2, ?3, ?4, 1, ?5, ?6)", - params![evidence_id.0, tenant_id.0, person_id.0, source_id.0, text, recorded_at], - )?; - Ok((source_id, evidence_id)) -} - -fn supersede_stale_claim( - transaction: &rusqlite::Transaction<'_>, - tenant_id: &TenantId, - person_id: &PersonId, - claim_id: &ClaimId, - stale_evidence_ids: &[EvidenceId], - valid_at: i64, - recorded_at: i64, -) -> Result<()> { - transaction.execute( - "UPDATE claims SET status = 'superseded', valid_until = ?1, recorded_until = ?2 WHERE id = ?3 AND tenant_id = ?4 AND person_id = ?5", - params![valid_at, recorded_at, claim_id.0, tenant_id.0, person_id.0], - )?; - invalidate_summaries_for_evidence( - transaction, - tenant_id, - person_id, - stale_evidence_ids, - recorded_at, - )?; - enqueue_projection_repair( - transaction, - tenant_id, - person_id, - EmbeddingTarget::Claim(claim_id.clone()), - "superseded_sync", - recorded_at, - )?; - transaction.execute( - "DELETE FROM profile_entries WHERE tenant_id = ?1 AND person_id = ?2 AND claim_id = ?3", - params![tenant_id.0, person_id.0, claim_id.0], - )?; - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -fn apply_correction_claim( - transaction: &rusqlite::Transaction<'_>, - tenant_id: &TenantId, - person_id: &PersonId, - superseded_claim_id: &ClaimId, - old: &OldClaim, - value: &str, - valid_at: i64, - recorded_at: i64, - source_id: &SourceId, - evidence_id: &EvidenceId, -) -> Result<(ClaimId, CorrectionRecord)> { - let claim_id = insert_claim( - transaction, - tenant_id, - person_id, - evidence_id, - ClaimInput { - subject: old.subject.clone(), - predicate: old.predicate.clone(), - value: value.to_owned(), - kind: claim_kind(&old.kind)?, - valid_from: valid_at, - tier: MemoryTier::LongTerm, - processing_state: MemoryProcessingState::Processed, - }, - recorded_at, - )?; - transaction.execute( - "UPDATE sources SET origin_evidence_id = ?1, origin_claim_id = ?2 WHERE id = ?3 AND tenant_id = ?4 AND person_id = ?5", - params![evidence_id.0, claim_id.0, source_id.0, tenant_id.0, person_id.0], - )?; - let correction = CorrectionRecord { - tenant_id: tenant_id.clone(), - person_id: person_id.clone(), - superseded_claim_id: superseded_claim_id.clone(), - claim_id: claim_id.clone(), - source_id: source_id.clone(), - evidence_id: evidence_id.clone(), - valid_at, - recorded_at, - }; - transaction.execute( - "INSERT INTO corrections(tenant_id, person_id, superseded_claim_id, claim_id, source_id, evidence_id, valid_at, recorded_at) VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", - params![correction.tenant_id.0, correction.person_id.0, correction.superseded_claim_id.0, correction.claim_id.0, correction.source_id.0, correction.evidence_id.0, correction.valid_at, correction.recorded_at], - )?; - Ok((claim_id, correction)) -} - -#[allow(clippy::too_many_arguments)] -fn build_correction_records( - transaction: &rusqlite::Transaction<'_>, - tenant_id: &TenantId, - person_id: &PersonId, - superseded_claim_id: &ClaimId, - source_id: &SourceId, - evidence_id: &EvidenceId, - claim_id: &ClaimId, - correction: CorrectionRecord, - removed_profiles: Vec, - recorded_at: i64, -) -> Result> { - let mut records = vec![ - ExportRecord::Source(source_record(transaction, tenant_id, person_id, source_id)?), - ExportRecord::Evidence(evidence_record( - transaction, - tenant_id, - person_id, - evidence_id, - )?), - ExportRecord::Claim(claim_record( - transaction, - tenant_id, - person_id, - superseded_claim_id, - )?), - ExportRecord::Claim(claim_record(transaction, tenant_id, person_id, claim_id)?), - ExportRecord::ClaimEvidence(claim_evidence_record( - transaction, - tenant_id, - person_id, - claim_id, - evidence_id, - )?), - ExportRecord::Correction(correction), - ]; - records.extend(removed_profiles.into_iter().map(|id| { - ExportRecord::Deletion(DeletionRecord { - tenant_id: tenant_id.clone(), - person_id: person_id.clone(), - target: MemoryRef::ProfileEntry(ProfileEntryId(id)), - deleted_at: recorded_at, - }) - })); - Ok(records) -} - #[allow(clippy::type_complexity)] fn find_deletion_targets( transaction: &rusqlite::Transaction<'_>, @@ -1328,27 +1177,19 @@ fn build_deletion_records( claim_id, )?)); } - if !profile_ids.is_empty() { - let profile_ids_json = serde_json::to_string(&profile_ids).unwrap(); - let mut stmt = transaction.prepare_cached( - "SELECT id FROM profile_entries WHERE tenant_id = ?1 AND person_id = ?2 AND id IN (SELECT value FROM json_each(?3))" + for profile_id in profile_ids { + let remains: bool = transaction.query_row( + "SELECT EXISTS(SELECT 1 FROM profile_entries WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3)", + params![profile_id, tenant_id.0, person_id.0], + |row| row.get(0), )?; - - let existing_profile_ids: std::collections::HashSet = stmt - .query_map(params![tenant_id.0, person_id.0, profile_ids_json], |row| { - row.get(0) - })? - .collect::>()?; - - for profile_id in profile_ids { - if !existing_profile_ids.contains(&profile_id) { - records.push(ExportRecord::Deletion(DeletionRecord { - tenant_id: tenant_id.clone(), - person_id: person_id.clone(), - target: MemoryRef::ProfileEntry(ProfileEntryId(profile_id)), - deleted_at, - })); - } + if !remains { + records.push(ExportRecord::Deletion(DeletionRecord { + tenant_id: tenant_id.clone(), + person_id: person_id.clone(), + target: MemoryRef::ProfileEntry(ProfileEntryId(profile_id)), + deleted_at, + })); } } records.extend(review_ids.into_iter().map(|id| { diff --git a/src/store/repair.rs b/src/store/repair.rs index 7109f4d..afe4e82 100644 --- a/src/store/repair.rs +++ b/src/store/repair.rs @@ -108,90 +108,6 @@ struct EmbeddingRow { vector: String, } -fn fetch_target_embeddings( - transaction: &Transaction<'_>, - input: &RepairInput, - rows: &[(String, String, String)], -) -> Result>> { - let mut embeddings_by_target: HashMap<(String, String), Vec> = HashMap::new(); - - if rows.is_empty() { - return Ok(embeddings_by_target); - } - - let mut query = String::from( - "SELECT target_kind, target_id, model, version, target_revision, input_hash, dimension, normalization, distance, vector FROM embeddings WHERE tenant_id = ? AND person_id = ? AND (target_kind, target_id) IN (", - ); - - for (i, _) in rows.iter().enumerate() { - if i > 0 { - query.push_str(", "); - } - query.push_str("(?, ?)"); - } - query.push(')'); - - let mut sql_params: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(2 + rows.len() * 2); - sql_params.push(&input.tenant_id.0); - sql_params.push(&input.person_id.0); - for (_, target_kind, target_id) in rows { - sql_params.push(target_kind); - sql_params.push(target_id); - } - - let mut statement = transaction.prepare(&query)?; - let embedding_rows = statement.query_map(rusqlite::params_from_iter(sql_params), |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, String>(1)?, - EmbeddingRow { - model: row.get::<_, String>(2)?, - version: row.get::<_, String>(3)?, - revision: row.get::<_, i64>(4)?, - hash: row.get::<_, String>(5)?, - dimension: row.get::<_, usize>(6)?, - normalization: row.get::<_, String>(7)?, - distance: row.get::<_, String>(8)?, - vector: row.get::<_, String>(9)?, - }, - )) - })?; - - for row in embedding_rows { - let (target_kind, target_id, embed_row) = row?; - embeddings_by_target - .entry((target_kind, target_id)) - .or_default() - .push(embed_row); - } - - Ok(embeddings_by_target) -} - -fn mark_repair_outbox_processed( - transaction: &Transaction<'_>, - processed_ids: &[String], - processed_at: Timestamp, -) -> Result<()> { - if processed_ids.is_empty() { - return Ok(()); - } - - let placeholders = vec!["?"; processed_ids.len()].join(", "); - let query = format!( - "UPDATE memory_repair_outbox SET processed_at = ?1 WHERE id IN ({})", - placeholders - ); - let mut params: Vec> = vec![Box::new(processed_at)]; - for id in processed_ids { - params.push(Box::new(id.clone())); - } - let param_refs: Vec<&dyn rusqlite::ToSql> = params.iter().map(|s| s.as_ref()).collect(); - transaction.execute(&query, param_refs.as_slice())?; - - Ok(()) -} - fn process_repair_target( transaction: &Transaction<'_>, input: &RepairInput, @@ -264,7 +180,56 @@ impl MemoryDb { let processed_at: Timestamp = transaction.query_row("SELECT unixepoch()", [], |row| row.get(0))?; - let embeddings_by_target = fetch_target_embeddings(&transaction, &input, &rows)?; + let mut embeddings_by_target: HashMap<(String, String), Vec> = HashMap::new(); + + if !rows.is_empty() { + let mut query = String::from( + "SELECT target_kind, target_id, model, version, target_revision, input_hash, dimension, normalization, distance, vector FROM embeddings WHERE tenant_id = ? AND person_id = ? AND (target_kind, target_id) IN (", + ); + + for (i, _) in rows.iter().enumerate() { + if i > 0 { + query.push_str(", "); + } + query.push_str("(?, ?)"); + } + query.push(')'); + + let mut sql_params: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(2 + rows.len() * 2); + sql_params.push(&input.tenant_id.0); + sql_params.push(&input.person_id.0); + for (_, target_kind, target_id) in &rows { + sql_params.push(target_kind); + sql_params.push(target_id); + } + + let mut statement = transaction.prepare(&query)?; + let embedding_rows = + statement.query_map(rusqlite::params_from_iter(sql_params), |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + EmbeddingRow { + model: row.get::<_, String>(2)?, + version: row.get::<_, String>(3)?, + revision: row.get::<_, i64>(4)?, + hash: row.get::<_, String>(5)?, + dimension: row.get::<_, usize>(6)?, + normalization: row.get::<_, String>(7)?, + distance: row.get::<_, String>(8)?, + vector: row.get::<_, String>(9)?, + }, + )) + })?; + + for row in embedding_rows { + let (target_kind, target_id, embed_row) = row?; + embeddings_by_target + .entry((target_kind, target_id)) + .or_default() + .push(embed_row); + } + } let mut processed = 0; let mut processed_ids = Vec::new(); @@ -292,8 +257,19 @@ impl MemoryDb { processed_ids.push(id); processed += 1; } - mark_repair_outbox_processed(&transaction, &processed_ids, processed_at)?; - + if !processed_ids.is_empty() { + let placeholders = vec!["?"; processed_ids.len()].join(", "); + let query = format!( + "UPDATE memory_repair_outbox SET processed_at = ?1 WHERE id IN ({})", + placeholders + ); + let mut params: Vec> = vec![Box::new(processed_at)]; + for id in &processed_ids { + params.push(Box::new(id.clone())); + } + let param_refs: Vec<&dyn rusqlite::ToSql> = params.iter().map(|s| s.as_ref()).collect(); + transaction.execute(&query, param_refs.as_slice())?; + } let summaries_stale = stale_summary_count(&transaction, &input.tenant_id, &input.person_id)?; record_operation( diff --git a/src/store/retrieval.rs b/src/store/retrieval.rs index 476fe10..90560a6 100644 --- a/src/store/retrieval.rs +++ b/src/store/retrieval.rs @@ -181,82 +181,121 @@ impl MemoryDb { self.retrieval_item(&input.tenant_id, &input.person_id, target, 10_000, None) } - pub(super) fn retrieval_targets_for_embedding( + pub(super) fn retrieval_targets_for_embeddings_bulk( &self, tenant_id: &TenantId, person_id: &PersonId, - target_kind: &str, - target_id: &str, - ) -> Result> { - let sql = match target_kind { - "claim" => { - "SELECT id FROM claims WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND status = 'accepted' AND valid_until IS NULL AND recorded_until IS NULL AND tier IN ('short_term', 'long_term') AND processing_state = 'processed'" - } - "evidence" => { - "SELECT c.id FROM evidence e JOIN sources s ON s.id = e.source_id AND s.tenant_id = e.tenant_id AND s.person_id = e.person_id LEFT JOIN claim_evidence ce ON ce.evidence_id = e.id AND ce.tenant_id = e.tenant_id AND ce.person_id = e.person_id AND ce.relation = '\"supports\"' LEFT JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ce.tenant_id AND c.person_id = ce.person_id AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' WHERE e.id = ?1 AND e.tenant_id = ?2 AND e.person_id = ?3 AND e.deleted_at IS NULL AND s.deleted_at IS NULL ORDER BY c.id" - } - "source" => { - "SELECT DISTINCT c.id FROM sources s JOIN evidence e ON e.source_id = s.id AND e.tenant_id = s.tenant_id AND e.person_id = s.person_id LEFT JOIN claim_evidence ce ON ce.evidence_id = e.id AND ce.tenant_id = e.tenant_id AND ce.person_id = e.person_id AND ce.relation = '\"supports\"' LEFT JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ce.tenant_id AND c.person_id = ce.person_id AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' WHERE s.id = ?1 AND s.tenant_id = ?2 AND s.person_id = ?3 AND s.deleted_at IS NULL AND e.deleted_at IS NULL ORDER BY c.id" - } - _ => { - return Err(Error::Invalid( - "stored embedding target is invalid".to_owned(), - )); + targets: &[(String, String)], + ) -> Result>> { + if targets.is_empty() { + return Ok(std::collections::HashMap::new()); + } + + let mut inputs_json = String::with_capacity(targets.len() * 40); + inputs_json.push('['); + for (i, (kind, id)) in targets.iter().enumerate() { + if i > 0 { + inputs_json.push(','); } - }; + // Use serde_json to safely escape strings for JSON + inputs_json.push_str(&format!( + "[{},{}]", + serde_json::to_string(kind).unwrap_or_else(|_| "\"\"".to_string()), + serde_json::to_string(id).unwrap_or_else(|_| "\"\"".to_string()) + )); + } + inputs_json.push(']'); + + let sql = " + WITH inputs AS ( + SELECT + json_extract(value, '$[0]') as kind, + json_extract(value, '$[1]') as id + FROM json_each(?3) + ) + SELECT DISTINCT i.kind, i.id as orig_id, 'claim' as target_kind, c.id as target_id + FROM inputs i + JOIN evidence e ON i.kind = 'evidence' AND e.id = i.id AND e.tenant_id = ?1 AND e.person_id = ?2 AND e.deleted_at IS NULL + JOIN sources s ON s.id = e.source_id AND s.tenant_id = ?1 AND s.person_id = ?2 AND s.deleted_at IS NULL + JOIN claim_evidence ce ON ce.evidence_id = e.id AND ce.tenant_id = ?1 AND ce.person_id = ?2 AND ce.relation = '\"supports\"' + JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ?1 AND c.person_id = ?2 AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' + + UNION ALL + + SELECT DISTINCT i.kind, i.id as orig_id, 'claim' as target_kind, c.id as target_id + FROM inputs i + JOIN sources s ON i.kind = 'source' AND s.id = i.id AND s.tenant_id = ?1 AND s.person_id = ?2 AND s.deleted_at IS NULL + JOIN evidence e ON e.source_id = s.id AND e.tenant_id = ?1 AND e.person_id = ?2 AND e.deleted_at IS NULL + JOIN claim_evidence ce ON ce.evidence_id = e.id AND ce.tenant_id = ?1 AND ce.person_id = ?2 AND ce.relation = '\"supports\"' + JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ?1 AND c.person_id = ?2 AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' + + UNION ALL + + SELECT DISTINCT i.kind, i.id as orig_id, 'evidence' as target_kind, e.id as target_id + FROM inputs i + JOIN evidence e ON i.kind = 'evidence' AND e.id = i.id AND e.tenant_id = ?1 AND e.person_id = ?2 AND e.deleted_at IS NULL + JOIN sources s ON s.id = e.source_id AND s.tenant_id = ?1 AND s.person_id = ?2 AND s.deleted_at IS NULL + WHERE NOT EXISTS ( + SELECT 1 + FROM claim_evidence ce + JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ?1 AND c.person_id = ?2 AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' + WHERE ce.evidence_id = i.id AND ce.tenant_id = ?1 AND ce.person_id = ?2 AND ce.relation = '\"supports\"' + ) + + UNION ALL + + SELECT DISTINCT i.kind, i.id as orig_id, 'source' as target_kind, s.id as target_id + FROM inputs i + JOIN sources s ON i.kind = 'source' AND s.id = i.id AND s.tenant_id = ?1 AND s.person_id = ?2 AND s.deleted_at IS NULL + WHERE NOT EXISTS ( + SELECT 1 + FROM evidence e + JOIN claim_evidence ce ON ce.evidence_id = e.id AND ce.tenant_id = ?1 AND ce.person_id = ?2 AND ce.relation = '\"supports\"' + JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ?1 AND c.person_id = ?2 AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' + WHERE e.source_id = i.id AND e.tenant_id = ?1 AND e.person_id = ?2 AND e.deleted_at IS NULL + ) + + UNION ALL + + SELECT DISTINCT i.kind, i.id as orig_id, 'claim' as target_kind, c.id as target_id + FROM inputs i + JOIN claims c ON i.kind = 'claim' AND c.id = i.id AND c.tenant_id = ?1 AND c.person_id = ?2 AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' + "; + let mut statement = self.connection.prepare(sql)?; - let rows = statement.query_map(params![target_id, tenant_id.0, person_id.0], |row| { - row.get::<_, Option>(0) + let rows = statement.query_map(params![tenant_id.0, person_id.0, inputs_json], |row| { + Ok(( + row.get::<_, String>(0)?, + row.get::<_, String>(1)?, + row.get::<_, String>(2)?, + row.get::<_, String>(3)?, + )) })?; - let rows = rows.collect::, _>>()?; - if rows.is_empty() { - return Ok(Vec::new()); - } - let claims = rows - .into_iter() - .flatten() - .map(|id| RetrievalTarget::Claim(ClaimId(id))) - .collect::>(); - if !claims.is_empty() { - return Ok(claims); + + let mut results: std::collections::HashMap<(String, String), Vec> = + std::collections::HashMap::new(); + for row in rows { + let (orig_kind, orig_id, target_kind, target_id) = row?; + let target = match target_kind.as_str() { + "claim" => RetrievalTarget::Claim(ClaimId(target_id)), + "evidence" => RetrievalTarget::Evidence(EvidenceId(target_id)), + "source" => RetrievalTarget::Source(SourceId(target_id)), + _ => continue, + }; + results + .entry((orig_kind, orig_id)) + .or_default() + .push(target); } - if self.target_has_claim(tenant_id, person_id, target_kind, target_id)? { - return Ok(Vec::new()); + + // Return empty vectors for any target that had no results. + // It acts exactly like the original method when no results are found. + for (kind, id) in targets { + let key = (kind.clone(), id.clone()); + results.entry(key).or_default(); } - Ok(match target_kind { - "source" => vec![RetrievalTarget::Source(SourceId(target_id.to_owned()))], - "evidence" => vec![RetrievalTarget::Evidence(EvidenceId(target_id.to_owned()))], - "claim" => Vec::new(), - _ => unreachable!(), - }) - } - fn target_has_claim( - &self, - tenant_id: &TenantId, - person_id: &PersonId, - target_kind: &str, - target_id: &str, - ) -> Result { - let sql = match target_kind { - "source" => { - "SELECT EXISTS(SELECT 1 FROM claim_evidence ce JOIN evidence e ON e.id = ce.evidence_id AND e.tenant_id = ce.tenant_id AND e.person_id = ce.person_id JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ce.tenant_id AND c.person_id = ce.person_id AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' WHERE ce.relation = '\"supports\"' AND e.source_id = ?1 AND e.tenant_id = ?2 AND e.person_id = ?3)" - } - "evidence" => { - "SELECT EXISTS(SELECT 1 FROM claim_evidence ce JOIN claims c ON c.id = ce.claim_id AND c.tenant_id = ce.tenant_id AND c.person_id = ce.person_id AND c.status = 'accepted' AND c.valid_until IS NULL AND c.recorded_until IS NULL AND c.tier IN ('short_term', 'long_term') AND c.processing_state = 'processed' WHERE ce.relation = '\"supports\"' AND ce.evidence_id = ?1 AND ce.tenant_id = ?2 AND ce.person_id = ?3)" - } - "claim" => return Ok(true), - _ => { - return Err(Error::Invalid( - "stored embedding target is invalid".to_owned(), - )); - } - }; - Ok(self - .connection - .query_row(sql, params![target_id, tenant_id.0, person_id.0], |row| { - row.get(0) - })?) + Ok(results) } fn retrieval_item( diff --git a/src/store/schema.rs b/src/store/schema.rs index 07c2fef..007fe7d 100644 --- a/src/store/schema.rs +++ b/src/store/schema.rs @@ -418,8 +418,8 @@ fn ensure_column( } let exists = transaction.query_row( - "SELECT EXISTS(SELECT 1 FROM pragma_table_info(?1) WHERE name = ?2)", - [table, column], + &format!("SELECT EXISTS(SELECT 1 FROM pragma_table_info('{table}') WHERE name = ?1)"), + [column], |row| row.get::<_, bool>(0), )?; if !exists { diff --git a/src/store/summaries.rs b/src/store/summaries.rs index 659db61..8798f91 100644 --- a/src/store/summaries.rs +++ b/src/store/summaries.rs @@ -313,28 +313,19 @@ fn validate_evidence( if count != chunk.len() { // Find the specific missing evidence to produce the exact error message - let chunk_json = - serde_json::to_string(&chunk.iter().map(|id| &id.0).collect::>()) - .map_err(|e| Error::Invalid(e.to_string()))?; - - let missing_id: String = transaction.query_row( - "SELECT value FROM json_each(?1) - WHERE NOT EXISTS ( - SELECT 1 FROM evidence - WHERE id = value - AND tenant_id = ?2 - AND person_id = ?3 - AND deleted_at IS NULL - ) - LIMIT 1", - params![chunk_json, tenant_id.0, person_id.0], - |row| row.get(0), - )?; - - return Err(Error::Invalid(format!( - "evidence {} is unavailable", - missing_id - ))); + for evidence_id in chunk { + let live: bool = transaction.query_row( + "SELECT EXISTS(SELECT 1 FROM evidence WHERE id = ?1 AND tenant_id = ?2 AND person_id = ?3 AND deleted_at IS NULL)", + params![evidence_id.0, tenant_id.0, person_id.0], + |row| row.get(0), + )?; + if !live { + return Err(Error::Invalid(format!( + "evidence {} is unavailable", + evidence_id.0 + ))); + } + } } } diff --git a/tests/main.rs b/tests/main.rs index 569c86e..b6dbe3d 100644 --- a/tests/main.rs +++ b/tests/main.rs @@ -31,19 +31,6 @@ fn test_help_argument() { } } -#[test] -fn test_db_help_argument() { - let output = Command::new(get_binary_path()) - .args(["--db", "mock.db", "help"]) - .output() - .expect("Failed to execute command"); - - assert!(output.status.success()); - let stdout = String::from_utf8(output.stdout).unwrap(); - assert!(stdout.contains("zkr --db PATH COMMAND")); - assert!(stdout.contains("Commands")); -} - #[test] fn test_invalid_arguments_error() { let output = Command::new(get_binary_path())