feat: implement provenance tagging and poisoning protection (issue #33)
This commit is contained in:
@@ -8,6 +8,7 @@ use std::collections::HashMap;
|
||||
|
||||
use crate::db::DbError;
|
||||
use crate::embedding::store::format_vector_literal;
|
||||
use crate::provenance::store as provenance_store;
|
||||
use crate::retrieval::RetrievalCandidate;
|
||||
use duckdb::Connection;
|
||||
|
||||
@@ -67,6 +68,14 @@ pub fn load_and_rerank(
|
||||
// Compute corpus embedding scores
|
||||
let corpus_scores = compute_corpus_scores(conn, query_vector, &id_list)?;
|
||||
|
||||
// Load provenance metadata for all candidates
|
||||
let all_ids: Vec<String> = memories.iter().map(|m| m.id.clone()).collect();
|
||||
let provenance_records = provenance_store::get_provenance_batch(conn, &all_ids)?;
|
||||
let provenance_map: HashMap<String, crate::provenance::ProvenanceRecord> = provenance_records
|
||||
.into_iter()
|
||||
.map(|r| (r.memory_id.clone(), r))
|
||||
.collect();
|
||||
|
||||
// Build candidates with all scores
|
||||
let mut candidates: Vec<RetrievalCandidate> = memories
|
||||
.into_iter()
|
||||
@@ -91,6 +100,7 @@ pub fn load_and_rerank(
|
||||
description_score: ds,
|
||||
corpus_score: cs,
|
||||
final_score,
|
||||
provenance_metadata: provenance_map.get(&mem.id).cloned(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Reference in New Issue
Block a user