use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use mentedb::MenteDb;
use mentedb::core::memory::{MemoryNode, MemoryType};
use mentedb_context::{AssemblyConfig, ContextAssembler, ScoredMemory};
use mentedb_core::types::AgentId;
use tempfile::TempDir;
fn generate_embedding(dim: usize, seed: u64) -> Vec<f32> {
(0..dim)
.map(|i| ((seed.wrapping_mul(31).wrapping_add(i as u64)) % 1000) as f32 / 1000.0)
.collect()
}
fn setup_db(count: usize) -> (MenteDb, TempDir) {
let tmp = TempDir::new().unwrap();
let mut db = MenteDb::open(tmp.path()).unwrap();
for i in 0..count {
let emb = generate_embedding(128, i as u64);
let node = MemoryNode::new(
AgentId::new(),
MemoryType::Semantic,
format!("Memory content number {i} with meaningful context for assembly"),
emb,
);
db.store(node).unwrap();
}
(db, tmp)
}
fn bench_context_assembly(c: &mut Criterion) {
let mut group = c.benchmark_group("context_assembly");
let config = AssemblyConfig {
token_budget: 4096,
..AssemblyConfig::default()
};
for count in [100, 1_000, 10_000] {
let (mut db, _tmp) = setup_db(count);
let query = generate_embedding(128, 999_999);
group.bench_with_input(
BenchmarkId::new("recall_and_assemble", count),
&count,
|b, _| {
b.iter(|| {
let results = db.recall_similar(&query, 50).unwrap();
let scored: Vec<ScoredMemory> = results
.iter()
.filter_map(|&(id, score)| {
db.get_memory(id)
.ok()
.map(|memory| ScoredMemory { memory, score })
})
.collect();
black_box(ContextAssembler::assemble(scored, vec![], &config));
});
},
);
}
group.finish();
}
criterion_group!(benches, bench_context_assembly);
criterion_main!(benches);