use criterion::{Criterion, black_box, criterion_group, criterion_main};
use mentedb::context::{AssemblyConfig, ContextAssembler, ScoredMemory};
use mentedb::core::memory::{MemoryNode, MemoryType};
use mentedb::index::hnsw::{HnswConfig, HnswIndex};
use mentedb::query::mql::Mql;
use mentedb::storage::StorageEngine;
use mentedb_core::types::{AgentId, MemoryId};
fn random_embedding(dim: usize) -> Vec<f32> {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::time::SystemTime;
let seed = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap()
.as_nanos();
let mut hasher = DefaultHasher::new();
seed.hash(&mut hasher);
let mut state = hasher.finish();
(0..dim)
.map(|i| {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(i as u64);
(state as f32 / u64::MAX as f32) * 2.0 - 1.0
})
.collect()
}
fn make_memory(i: usize) -> MemoryNode {
MemoryNode::new(
AgentId::new(),
MemoryType::Episodic,
format!("benchmark memory content for item number {i}"),
random_embedding(128),
)
}
fn bench_storage_write(c: &mut Criterion) {
c.bench_function("storage_write_1000", |b| {
b.iter_with_setup(
|| {
let dir = tempfile::tempdir().unwrap();
let engine = StorageEngine::open(dir.path()).unwrap();
let memories: Vec<MemoryNode> = (0..1000).map(make_memory).collect();
(dir, engine, memories)
},
|(_dir, mut engine, memories)| {
for mem in &memories {
black_box(engine.store_memory(mem).unwrap());
}
engine.close().unwrap();
},
);
});
}
fn bench_hnsw_search(c: &mut Criterion) {
let index = HnswIndex::new(HnswConfig::default());
for _ in 0..10_000 {
index
.insert(MemoryId::new(), &random_embedding(128))
.unwrap();
}
c.bench_function("hnsw_search_top10", |b| {
b.iter(|| {
let query = random_embedding(128);
black_box(index.search(&query, 10));
});
});
}
fn bench_context_assembly(c: &mut Criterion) {
let memories: Vec<ScoredMemory> = (0..100)
.map(|i| {
let score = (100 - i) as f32 / 100.0;
ScoredMemory {
memory: make_memory(i),
score,
}
})
.collect();
let config = AssemblyConfig::default();
c.bench_function("context_assembly_100", |b| {
b.iter(|| {
let mems = memories.clone();
black_box(ContextAssembler::assemble(mems, vec![], &config));
});
});
}
fn bench_mql_parse(c: &mut Criterion) {
let queries: Vec<String> = (0..100)
.map(|i| match i % 5 {
0 => format!(
"RECALL memories NEAR [{}] LIMIT 10",
(0..10)
.map(|j| format!("{:.2}", (i * 10 + j) as f32 / 100.0))
.collect::<Vec<_>>()
.join(", ")
),
1 => format!("RECALL memories WHERE type = episodic AND tag = \"bench{i}\" LIMIT 5"),
2 => format!(r#"RECALL memories WHERE content ~> "topic number {i}" LIMIT 10"#),
3 => format!(
"TRAVERSE {} DEPTH 3 WHERE edge_type = caused",
AgentId::new()
),
_ => format!("RECALL memories WHERE salience > 0.{i:02} LIMIT 20"),
})
.collect();
c.bench_function("mql_parse_100", |b| {
b.iter(|| {
for q in &queries {
black_box(Mql::parse(q).unwrap());
}
});
});
}
criterion_group!(
benches,
bench_storage_write,
bench_hnsw_search,
bench_context_assembly,
bench_mql_parse,
);
criterion_main!(benches);