use async_trait::async_trait;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use crate::memory::clock::Clock;
use crate::memory::config::MemoryConfig;
use crate::memory::embedding::EmbeddingProvider;
use crate::memory::error::{EmbeddingError, MemoryError};
use crate::memory::retrieval::recall;
use crate::memory::store::{Memory, VectorStore};
use crate::memory::tokens::estimate_tokens;
use crate::memory::MemoryKind;
#[derive(Debug, Clone)]
pub struct BenchFact {
pub id: String,
pub text: String,
pub supersedes: Option<String>,
}
#[derive(Debug, Clone)]
pub struct BenchProbe {
pub id: String,
pub query: String,
pub expected_fact_ids: Vec<String>,
pub staleness_fact_ids: Vec<String>,
}
#[derive(Debug)]
pub struct BenchmarkDataset {
pub facts: Vec<BenchFact>,
pub probes: Vec<BenchProbe>,
}
impl BenchmarkDataset {
#[allow(dead_code)]
fn fact_text(&self, fact_id: &str) -> &str {
self.facts
.iter()
.find(|f| f.id == fact_id)
.map(|f| f.text.as_str())
.unwrap_or("")
}
}
#[derive(Debug, Clone)]
pub struct BenchProbeResult {
pub probe_id: String,
pub retrieved_fact_ids: Vec<String>,
pub context_tokens: usize,
}
#[derive(Debug, Clone)]
pub struct ArmReport {
pub recall_accuracy: f64,
pub staleness_rate: f64,
pub mean_context_tokens: f64,
}
#[derive(Debug, Clone)]
pub struct BenchmarkReport {
pub load_all: ArmReport,
pub selective: ArmReport,
}
pub struct DeterministicEmbedder {
dim: usize,
model: String,
}
impl DeterministicEmbedder {
pub fn new(dim: usize, model: impl Into<String>) -> Self {
Self {
dim,
model: model.into(),
}
}
fn word_bag(&self, text: &str) -> Vec<f32> {
let mut v = vec![0.0f32; self.dim];
for word in text.split_whitespace() {
let hash = Sha256::digest(word.to_lowercase().as_bytes());
let idx = u64::from_le_bytes(hash[..8].try_into().expect("sha256 ≥ 8 bytes")) as usize
% self.dim;
v[idx] += 1.0;
}
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 1e-9 {
v.iter_mut().for_each(|x| *x /= norm);
}
v
}
}
#[async_trait]
impl EmbeddingProvider for DeterministicEmbedder {
async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
Ok(texts.iter().map(|t| self.word_bag(t)).collect())
}
fn model_id(&self) -> &str {
&self.model
}
fn dim(&self) -> usize {
self.dim
}
fn query_prefix(&self) -> &str {
""
}
fn document_prefix(&self) -> &str {
""
}
}
pub fn compute_recall_accuracy(dataset: &BenchmarkDataset, results: &[BenchProbeResult]) -> f64 {
if results.is_empty() {
return 0.0;
}
let hits = results
.iter()
.filter(|r| {
let Some(probe) = dataset.probes.iter().find(|p| p.id == r.probe_id) else {
return false;
};
probe
.expected_fact_ids
.iter()
.all(|eid| r.retrieved_fact_ids.contains(eid))
})
.count();
hits as f64 / results.len() as f64
}
pub fn compute_staleness_rate(dataset: &BenchmarkDataset, results: &[BenchProbeResult]) -> f64 {
if results.is_empty() {
return 0.0;
}
let stale = results
.iter()
.filter(|r| {
let Some(probe) = dataset.probes.iter().find(|p| p.id == r.probe_id) else {
return false;
};
probe
.staleness_fact_ids
.iter()
.any(|sid| r.retrieved_fact_ids.contains(sid))
})
.count();
stale as f64 / results.len() as f64
}
pub fn compute_mean_context_tokens(results: &[BenchProbeResult]) -> f64 {
if results.is_empty() {
return 0.0;
}
let total: usize = results.iter().map(|r| r.context_tokens).sum();
total as f64 / results.len() as f64
}
async fn populate_store(
dataset: &BenchmarkDataset,
store: &dyn VectorStore,
embedder: &dyn EmbeddingProvider,
clock: &dyn Clock,
) -> Result<(), MemoryError> {
let base_ts = clock.now();
let mut id_map: HashMap<String, String> = HashMap::new();
for (i, fact) in dataset.facts.iter().enumerate() {
let mem_id = format!("bench-{}", fact.id);
id_map.insert(fact.id.clone(), mem_id.clone());
let doc_prefix = embedder.document_prefix().to_string();
let prefixed = if doc_prefix.is_empty() {
fact.text.clone()
} else {
format!("{doc_prefix}{}", fact.text)
};
let embedding = embedder
.embed(&[prefixed])
.await
.map_err(MemoryError::Embedding)?
.pop()
.unwrap_or_default();
let memory = Memory {
id: mem_id.clone(),
session_id: "bench-session".into(),
kind: MemoryKind::Episodic,
text: fact.text.clone(),
embedding,
model_id: embedder.model_id().to_string(),
dim: embedder.dim(),
created_at: base_ts + i as i64,
salience: 0.5,
access_count: 0,
last_accessed_at: base_ts + i as i64,
superseded_by: None,
evicted_at: None,
scope: "root".into(),
distilled_at: None,
};
match store.insert(&memory).await {
Ok(()) => {}
Err(MemoryError::Storage(ref e)) if e.contains("UNIQUE constraint") => {}
Err(e) => return Err(e),
}
}
for fact in &dataset.facts {
if let Some(old_id) = &fact.supersedes {
if let (Some(old_mem_id), Some(new_mem_id)) =
(id_map.get(old_id.as_str()), id_map.get(&fact.id))
{
store.set_superseded(old_mem_id, new_mem_id).await?;
}
}
}
Ok(())
}
fn run_load_all_arm(dataset: &BenchmarkDataset, cfg: &MemoryConfig) -> Vec<BenchProbeResult> {
let all_texts: String = dataset
.facts
.iter()
.map(|f| f.text.as_str())
.collect::<Vec<_>>()
.join(" ");
let total_tokens = estimate_tokens(&all_texts, cfg.chars_per_token);
dataset
.probes
.iter()
.map(|probe| {
let retrieved_fact_ids = dataset.facts.iter().map(|f| f.id.clone()).collect();
BenchProbeResult {
probe_id: probe.id.clone(),
retrieved_fact_ids,
context_tokens: total_tokens,
}
})
.collect()
}
async fn run_selective_arm(
dataset: &BenchmarkDataset,
store: &dyn VectorStore,
embedder: &dyn EmbeddingProvider,
clock: &dyn Clock,
cfg: &MemoryConfig,
) -> Result<Vec<BenchProbeResult>, MemoryError> {
let mut results = Vec::new();
for probe in &dataset.probes {
let ranked = recall(
store,
embedder,
clock,
cfg,
&probe.query,
cfg.context_budget_tokens,
"root",
)
.await?;
let context_text: String = ranked
.iter()
.map(|rm| rm.memory.text.as_str())
.collect::<Vec<_>>()
.join(" ");
let context_tokens = estimate_tokens(&context_text, cfg.chars_per_token);
let retrieved_fact_ids: Vec<String> = ranked
.iter()
.filter_map(|rm| {
dataset
.facts
.iter()
.find(|f| f.text == rm.memory.text)
.map(|f| f.id.clone())
})
.collect();
results.push(BenchProbeResult {
probe_id: probe.id.clone(),
retrieved_fact_ids,
context_tokens,
});
}
Ok(results)
}
pub async fn run_full_benchmark(
dataset: &BenchmarkDataset,
store: &dyn VectorStore,
embedder: &dyn EmbeddingProvider,
clock: &dyn Clock,
cfg: &MemoryConfig,
) -> Result<BenchmarkReport, MemoryError> {
populate_store(dataset, store, embedder, clock).await?;
let load_all_results = run_load_all_arm(dataset, cfg);
let selective_results = run_selective_arm(dataset, store, embedder, clock, cfg).await?;
Ok(BenchmarkReport {
load_all: ArmReport {
recall_accuracy: compute_recall_accuracy(dataset, &load_all_results),
staleness_rate: compute_staleness_rate(dataset, &load_all_results),
mean_context_tokens: compute_mean_context_tokens(&load_all_results),
},
selective: ArmReport {
recall_accuracy: compute_recall_accuracy(dataset, &selective_results),
staleness_rate: compute_staleness_rate(dataset, &selective_results),
mean_context_tokens: compute_mean_context_tokens(&selective_results),
},
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use tempfile::NamedTempFile;
use crate::memory::clock::FixedClock;
use crate::memory::config::MemoryConfig;
use crate::memory::store::SqliteVectorStore;
use crate::system::database::EncryptedSqliteMemory;
fn make_small_dataset() -> BenchmarkDataset {
BenchmarkDataset {
facts: vec![
BenchFact {
id: "pref_rust".into(),
text: "I prefer Rust over Python for systems programming".into(),
supersedes: None,
},
BenchFact {
id: "pref_dark".into(),
text: "I use dark mode in all editors and terminals".into(),
supersedes: None,
},
BenchFact {
id: "editor_old".into(),
text: "I use vim as my primary editor".into(),
supersedes: None,
},
BenchFact {
id: "editor_new".into(),
text: "I switched from vim to neovim".into(),
supersedes: Some("editor_old".into()),
},
BenchFact {
id: "distractor_1".into(),
text: "The capital of France is Paris".into(),
supersedes: None,
},
BenchFact {
id: "distractor_2".into(),
text: "Water boils at one hundred degrees Celsius".into(),
supersedes: None,
},
],
probes: vec![
BenchProbe {
id: "q_rust".into(),
query: "Rust Python systems programming".into(),
expected_fact_ids: vec!["pref_rust".into()],
staleness_fact_ids: vec![],
},
BenchProbe {
id: "q_editor".into(),
query: "editor vim neovim".into(),
expected_fact_ids: vec!["editor_new".into()],
staleness_fact_ids: vec!["editor_old".into()],
},
BenchProbe {
id: "q_dark".into(),
query: "dark mode editors terminals".into(),
expected_fact_ids: vec!["pref_dark".into()],
staleness_fact_ids: vec![],
},
],
}
}
fn make_small_cfg() -> MemoryConfig {
MemoryConfig {
top_k: 2,
chars_per_token: 4.0,
..Default::default()
}
}
fn open_store() -> (NamedTempFile, Arc<SqliteVectorStore>) {
let tmp = NamedTempFile::new().unwrap();
let mem = EncryptedSqliteMemory::new(tmp.path().to_path_buf(), "benchpw".into()).unwrap();
let store = Arc::new(SqliteVectorStore::new(mem.shared_conn(), mem.data_key()).unwrap());
(tmp, store)
}
#[test]
fn test_recall_accuracy_hand_calculated() {
let dataset = BenchmarkDataset {
facts: vec![
BenchFact {
id: "f1".into(),
text: "fact one".into(),
supersedes: None,
},
BenchFact {
id: "f2".into(),
text: "fact two".into(),
supersedes: None,
},
],
probes: vec![
BenchProbe {
id: "q1".into(),
query: "one".into(),
expected_fact_ids: vec!["f1".into()],
staleness_fact_ids: vec![],
},
BenchProbe {
id: "q2".into(),
query: "two".into(),
expected_fact_ids: vec!["f2".into()],
staleness_fact_ids: vec![],
},
],
};
let results = vec![
BenchProbeResult {
probe_id: "q1".into(),
retrieved_fact_ids: vec!["f1".into()], context_tokens: 10,
},
BenchProbeResult {
probe_id: "q2".into(),
retrieved_fact_ids: vec![], context_tokens: 5,
},
];
let acc = compute_recall_accuracy(&dataset, &results);
assert!(
(acc - 0.5).abs() < 1e-9,
"expected recall_accuracy=0.5 (1 hit / 2 probes), got {acc}"
);
}
#[test]
fn test_staleness_rate_hand_calculated() {
let dataset = BenchmarkDataset {
facts: vec![
BenchFact {
id: "old".into(),
text: "old fact".into(),
supersedes: None,
},
BenchFact {
id: "new".into(),
text: "new fact".into(),
supersedes: Some("old".into()),
},
],
probes: vec![
BenchProbe {
id: "q1".into(),
query: "fact".into(),
expected_fact_ids: vec!["new".into()],
staleness_fact_ids: vec!["old".into()],
},
BenchProbe {
id: "q2".into(),
query: "other".into(),
expected_fact_ids: vec![],
staleness_fact_ids: vec![],
},
],
};
let results = vec![
BenchProbeResult {
probe_id: "q1".into(),
retrieved_fact_ids: vec!["old".into()], context_tokens: 8,
},
BenchProbeResult {
probe_id: "q2".into(),
retrieved_fact_ids: vec![], context_tokens: 4,
},
];
let rate = compute_staleness_rate(&dataset, &results);
assert!(
(rate - 0.5).abs() < 1e-9,
"expected staleness_rate=0.5 (1 stale / 2 probes), got {rate}"
);
}
#[test]
fn test_mean_context_tokens_hand_calculated() {
let results = vec![
BenchProbeResult {
probe_id: "q1".into(),
retrieved_fact_ids: vec![],
context_tokens: 10,
},
BenchProbeResult {
probe_id: "q2".into(),
retrieved_fact_ids: vec![],
context_tokens: 20,
},
];
let mean = compute_mean_context_tokens(&results);
assert!(
(mean - 15.0).abs() < 1e-9,
"expected mean_context_tokens=15.0, got {mean}"
);
}
#[tokio::test]
async fn test_benchmark_uses_deterministic_embedder() {
let emb = DeterministicEmbedder::new(64, "det-test");
assert_eq!(emb.model_id(), "det-test");
assert_eq!(emb.dim(), 64);
let texts = vec!["hello world".to_string(), "rust programming".to_string()];
let vecs = emb
.embed(&texts)
.await
.expect("embed must succeed without network");
assert_eq!(vecs.len(), 2, "one vector per input text");
assert_eq!(vecs[0].len(), 64, "vector length = dim");
assert_eq!(vecs[1].len(), 64, "vector length = dim");
let norm: f32 = vecs[0].iter().map(|x| x * x).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 1e-5,
"word-bag vector must be L2-normalised; got norm={norm}"
);
let v1 = emb.embed(&["rust".to_string()]).await.unwrap().remove(0);
let v2 = emb.embed(&["rust".to_string()]).await.unwrap().remove(0);
assert_eq!(
v1, v2,
"embedding must be deterministic under the same input"
);
}
#[tokio::test]
async fn test_benchmark_runs_both_arms_and_emits_report() {
let dataset = make_small_dataset();
let cfg = make_small_cfg();
let emb = Arc::new(DeterministicEmbedder::new(128, "det-bench"));
let clk = Arc::new(FixedClock::new(1_000_000));
let (_tmp, store) = open_store();
let report = run_full_benchmark(&dataset, &*store, &*emb, &*clk, &cfg)
.await
.expect("run_full_benchmark must succeed with DeterministicEmbedder");
assert!(
report.load_all.recall_accuracy >= 0.0 && report.load_all.recall_accuracy <= 1.0,
"load_all recall out of range: {}",
report.load_all.recall_accuracy
);
assert!(
report.selective.recall_accuracy >= 0.0 && report.selective.recall_accuracy <= 1.0,
"selective recall out of range: {}",
report.selective.recall_accuracy
);
assert!(
report.load_all.staleness_rate >= 0.0 && report.load_all.staleness_rate <= 1.0,
"load_all staleness out of range: {}",
report.load_all.staleness_rate
);
assert!(
report.selective.staleness_rate >= 0.0 && report.selective.staleness_rate <= 1.0,
"selective staleness out of range: {}",
report.selective.staleness_rate
);
assert!(
report.load_all.mean_context_tokens > 0.0,
"load_all must have non-zero context tokens (all facts injected)"
);
}
#[tokio::test]
async fn test_benchmark_is_reproducible_under_same_seed() {
let dataset = make_small_dataset();
let cfg = make_small_cfg();
let emb = Arc::new(DeterministicEmbedder::new(128, "det-bench"));
let clk = Arc::new(FixedClock::new(1_000_000));
let (_tmp, store) = open_store();
let r1 = run_full_benchmark(&dataset, &*store, &*emb, &*clk, &cfg)
.await
.unwrap();
let r2 = run_full_benchmark(&dataset, &*store, &*emb, &*clk, &cfg)
.await
.unwrap();
assert!(
(r1.selective.recall_accuracy - r2.selective.recall_accuracy).abs() < 1e-9,
"selective recall must be identical across runs: {} vs {}",
r1.selective.recall_accuracy,
r2.selective.recall_accuracy
);
assert!(
(r1.selective.staleness_rate - r2.selective.staleness_rate).abs() < 1e-9,
"selective staleness must be identical across runs"
);
}
#[tokio::test]
async fn test_selective_outperforms_load_all_on_dataset() {
let dataset = make_small_dataset();
let cfg = MemoryConfig {
top_k: 2,
chars_per_token: 4.0,
..Default::default()
};
let emb = Arc::new(DeterministicEmbedder::new(128, "det-bench"));
let clk = Arc::new(FixedClock::new(1_000_000));
let (_tmp, store) = open_store();
let report = run_full_benchmark(&dataset, &*store, &*emb, &*clk, &cfg)
.await
.expect("run_full_benchmark must succeed");
assert!(
report.selective.recall_accuracy >= report.load_all.recall_accuracy,
"selective recall ({}) must be ≥ load_all recall ({})",
report.selective.recall_accuracy,
report.load_all.recall_accuracy
);
assert!(
report.selective.mean_context_tokens < report.load_all.mean_context_tokens,
"selective context tokens ({:.1}) must be < load_all ({:.1}) — bounded window",
report.selective.mean_context_tokens,
report.load_all.mean_context_tokens
);
assert!(
report.selective.staleness_rate <= report.load_all.staleness_rate,
"selective staleness ({}) must be ≤ load_all ({}) — supersession excludes stale facts",
report.selective.staleness_rate,
report.load_all.staleness_rate
);
assert!(
report.load_all.staleness_rate > 0.0,
"load_all must exhibit staleness for the superseded editor fact; got 0.0"
);
assert!(
report.selective.staleness_rate == 0.0,
"selective must have zero staleness since superseded facts are excluded from active(); \
got {}",
report.selective.staleness_rate
);
}
}