1use std::collections::{BTreeMap, BTreeSet};
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::process::Command;
5
6use anyhow::{Context, Result};
7use rusqlite::Connection;
8use serde_json::{json, Value};
9use sha2::{Digest, Sha256};
10
11use crate::eval::bench_artifact::{
12 BenchmarkLayer, MemoryCitationEvidence, MemoryDiagnosis, MemoryRetrievalEvidence,
13 MemoryRunArtifact, PublicBenchmarkReport, ReportVerifierMetadata, RunEnvironment,
14};
15
16use super::baselines::fixture_retrieval_indices;
17use super::diagnostics::{
18 classify_diagnosis, failure_decomposition, performance_by_condition, performance_metrics,
19 score_policy,
20};
21use super::fixture::load_suite;
22use super::types::{
23 summarize_by_category, summarize_metrics, summarize_policy, MemoryBenchCondition,
24 MemoryBenchEvidence, MemoryBenchRunOutcome, MemoryBenchSuiteFixture, MemoryBenchTask,
25 DEFAULT_PUBLIC_ROOT, DEFAULT_REPORT_BENCHMARK_VERSION,
26};
27
28const PROJECT: &str = "/tmp/remem-memory-bench/repo";
29const READER_PROVIDER: &str = "fixture";
30const READER_MODEL: &str = "deterministic-memory-reader";
31
32#[derive(Debug, Clone)]
33pub struct MemoryBenchOptions {
34 pub suite: String,
35 pub condition: Option<String>,
36 pub json_out: String,
37 pub root: String,
38 pub artifact_prefix: Option<String>,
39}
40
41pub fn run_memory_bench(options: MemoryBenchOptions) -> Result<PublicBenchmarkReport> {
42 let fixture = load_suite(&options.suite)?;
43 let conditions = selected_conditions(options.condition.as_deref())?;
44 let public_root = PathBuf::from(if options.root.trim().is_empty() {
45 DEFAULT_PUBLIC_ROOT
46 } else {
47 options.root.as_str()
48 });
49 let json_out = PathBuf::from(&options.json_out);
50 let artifact_prefix = options
51 .artifact_prefix
52 .unwrap_or_else(|| format!("memory/artifacts/{}", fixture.fixture_revision));
53 let public_layout = path_starts_with(&json_out, &public_root);
54 let artifact_root = if public_layout {
55 public_root.join(&artifact_prefix)
56 } else {
57 sibling_artifact_root(&json_out)
58 };
59 fs::create_dir_all(&artifact_root).with_context(|| {
60 format!(
61 "create memory benchmark artifacts {}",
62 artifact_root.display()
63 )
64 })?;
65
66 let mut outcomes = Vec::new();
67 let mut run_artifacts = Vec::new();
68 for condition in conditions {
69 for task in &fixture.tasks {
70 let outcome = run_task(&fixture, condition, task)?;
71 let run_json_path = write_run_artifacts(
72 &fixture,
73 &outcome,
74 task,
75 &artifact_root,
76 &public_root,
77 public_layout,
78 )?;
79 run_artifacts.push(run_json_path);
80 outcomes.push(outcome);
81 }
82 }
83
84 let aggregate_metrics = json!({
85 "suite": fixture.suite,
86 "suite_version": fixture.version,
87 "fixture_revision": fixture.fixture_revision,
88 "run_count": outcomes.len(),
89 "overall": summarize_metrics(&outcomes),
90 "by_category": summarize_by_category(&outcomes),
91 "conditions": summarize_by_condition(&outcomes),
92 "failure_decomposition": failure_decomposition(&outcomes),
93 "performance": performance_by_condition(&outcomes),
94 "policy": summarize_policy(&outcomes),
95 });
96 let report = PublicBenchmarkReport {
97 schema_version: 1,
98 benchmark_id: fixture.benchmark_id.clone(),
99 benchmark_version: fixture.version.clone(),
100 layer: BenchmarkLayer::MemorySystemCapability,
101 conditions: outcomes
102 .iter()
103 .map(|outcome| outcome.condition.as_str().to_string())
104 .collect::<BTreeSet<_>>()
105 .into_iter()
106 .collect(),
107 schema_refs: vec![
108 "schemas/benchmark-manifest.schema.json".to_string(),
109 "schemas/memory-report.schema.json".to_string(),
110 "schemas/memory-run.schema.json".to_string(),
111 "schemas/reproduction-metadata.schema.json".to_string(),
112 ],
113 run_artifacts,
114 aggregate_metrics,
115 claim_level: "directional_memory_suite_no_public_claim".to_string(),
116 verifier: ReportVerifierMetadata {
117 required: true,
118 schema_version: 1,
119 },
120 };
121
122 if let Some(parent) = json_out.parent() {
123 if !parent.as_os_str().is_empty() {
124 fs::create_dir_all(parent).with_context(|| {
125 format!("create memory benchmark report dir {}", parent.display())
126 })?;
127 }
128 }
129 fs::write(&json_out, serde_json::to_string_pretty(&report)?)
130 .with_context(|| format!("write memory benchmark report {}", json_out.display()))?;
131 Ok(report)
132}
133
134fn selected_conditions(condition: Option<&str>) -> Result<Vec<MemoryBenchCondition>> {
135 match condition {
136 Some(raw) => {
137 let condition = MemoryBenchCondition::parse(raw)
138 .with_context(|| format!("unknown memory benchmark condition {raw}"))?;
139 Ok(vec![condition])
140 }
141 None => Ok(MemoryBenchCondition::ALL.to_vec()),
142 }
143}
144
145fn run_task(
146 fixture: &MemoryBenchSuiteFixture,
147 condition: MemoryBenchCondition,
148 task: &MemoryBenchTask,
149) -> Result<MemoryBenchRunOutcome> {
150 let retrieved = if let Some(indices) = fixture_retrieval_indices(condition, task) {
151 indices
152 .into_iter()
153 .map(|idx| RetrievedEvidence::from_fixture(idx, &task.evidence[idx]))
154 .collect()
155 } else {
156 retrieve_with_remem_search(task)?
157 };
158 Ok(score_task(fixture, condition, task, retrieved))
159}
160
161fn retrieve_with_remem_search(task: &MemoryBenchTask) -> Result<Vec<RetrievedEvidence>> {
162 let conn = Connection::open_in_memory()?;
163 crate::migrate::run_migrations(&conn)?;
164 let mut by_memory_id = BTreeMap::new();
165 for evidence in task
166 .evidence
167 .iter()
168 .filter(|evidence| evidence.retention_allowed)
169 {
170 let files = if evidence.files.is_empty() {
171 None
172 } else {
173 Some(serde_json::to_string(&evidence.files)?)
174 };
175 let id = crate::memory::insert_memory_full_with_reference_time(
176 &conn,
177 Some(&evidence.event_id),
178 PROJECT,
179 evidence.topic_key.as_deref(),
180 &evidence.title,
181 &evidence.content,
182 &evidence.memory_type,
183 files.as_deref(),
184 Some("main"),
185 &evidence.scope,
186 evidence.created_at_epoch,
187 evidence.created_at_epoch,
188 )?;
189 if evidence.status != "active" {
190 conn.execute(
191 "UPDATE memories SET status = ?1 WHERE id = ?2",
192 rusqlite::params![evidence.status, id],
193 )?;
194 }
195 by_memory_id.insert(id, evidence);
196 }
197
198 let hits = crate::retrieval::search::search_with_branch(
199 &conn,
200 Some(&task.query),
201 Some(PROJECT),
202 None,
203 5,
204 0,
205 false,
206 Some("main"),
207 )?;
208 Ok(hits
209 .into_iter()
210 .filter_map(|memory| {
211 by_memory_id
212 .get(&memory.id)
213 .map(|evidence| RetrievedEvidence::from_memory(memory.id, evidence))
214 })
215 .collect())
216}
217
218fn score_task(
219 fixture: &MemoryBenchSuiteFixture,
220 condition: MemoryBenchCondition,
221 task: &MemoryBenchTask,
222 retrieved: Vec<RetrievedEvidence>,
223) -> MemoryBenchRunOutcome {
224 let gold = task
225 .gold_supporting_event_ids
226 .iter()
227 .cloned()
228 .collect::<BTreeSet<_>>();
229 let forbidden = task
230 .forbidden_event_ids
231 .iter()
232 .cloned()
233 .collect::<BTreeSet<_>>();
234 let retrieved_events = retrieved
235 .iter()
236 .map(|item| item.event_id.clone())
237 .collect::<Vec<_>>();
238 let retrieved_set = retrieved_events.iter().cloned().collect::<BTreeSet<_>>();
239 let retrieved_gold = gold
240 .intersection(&retrieved_set)
241 .cloned()
242 .collect::<Vec<_>>();
243 let missing_event_ids = gold.difference(&retrieved_set).cloned().collect::<Vec<_>>();
244 let forbidden_count = forbidden.intersection(&retrieved_set).count();
245 let support_coverage = ratio(retrieved_gold.len(), gold.len());
246 let evidence_complete = missing_event_ids.is_empty() && forbidden_count == 0;
247 let expected_policy_abstention = task
248 .policy
249 .as_ref()
250 .map(|policy| policy.expected_policy_abstention)
251 .unwrap_or(false);
252 let abstained = expected_policy_abstention || !evidence_complete;
253 let answer_score = if evidence_complete
254 || ((task.abstention_allowed || expected_policy_abstention) && abstained)
255 {
256 1.0
257 } else {
258 0.0
259 };
260 let answer_text = if abstained {
261 "Insufficient benchmark evidence to answer.".to_string()
262 } else {
263 task.expected_answer.clone()
264 };
265 let cited_memory_ids = if abstained {
266 Vec::new()
267 } else {
268 retrieved
269 .iter()
270 .filter(|item| gold.contains(&item.event_id))
271 .map(|item| item.memory_id)
272 .collect()
273 };
274 let cited_event_ids = if abstained {
275 Vec::new()
276 } else {
277 retrieved_gold.clone()
278 };
279 let citation_recall = if abstained {
280 0.0
281 } else {
282 ratio(cited_event_ids.len(), gold.len())
283 };
284 let citation_precision = if abstained || cited_event_ids.is_empty() {
285 0.0
286 } else {
287 ratio(
288 cited_event_ids.len(),
289 cited_event_ids.len() + forbidden_count,
290 )
291 };
292 let staleness_accuracy = if forbidden_count == 0 { 1.0 } else { 0.0 };
293 let expected_abstention = condition == MemoryBenchCondition::NoMemory
294 || task.abstention_allowed
295 || expected_policy_abstention;
296 let abstention_accuracy = if abstained == expected_abstention {
297 1.0
298 } else {
299 0.0
300 };
301 let policy = score_policy(condition, task, &retrieved_events, abstained);
302 let reader_input = build_reader_input(condition, task, &retrieved);
303 let diagnosis =
304 classify_diagnosis(condition, task, &missing_event_ids, answer_score, abstained);
305 let performance = performance_metrics(condition, task, &reader_input, retrieved.len());
306 let retrieved_evidence_json = json!({
307 "suite": fixture.suite,
308 "fixture_revision": fixture.fixture_revision,
309 "condition": condition.as_str(),
310 "task_id": task.id,
311 "retrieved": retrieved.iter().map(RetrievedEvidence::to_json).collect::<Vec<_>>(),
312 });
313 let mut diagnosis_notes = Vec::new();
314 if !missing_event_ids.is_empty() {
315 diagnosis_notes.push(format!(
316 "missing supporting evidence: {}",
317 missing_event_ids.join(",")
318 ));
319 }
320 if forbidden_count > 0 {
321 diagnosis_notes.push(format!(
322 "retrieved forbidden evidence count: {forbidden_count}"
323 ));
324 }
325 if policy.policy_failure_count > 0 {
326 diagnosis_notes.push(format!(
327 "structured policy failure count: {}",
328 policy.policy_failure_count
329 ));
330 }
331 if policy.non_retention_leaked {
332 diagnosis_notes.push("non-retention leak detected".to_string());
333 }
334 if policy.false_blocked {
335 diagnosis_notes.push("approved policy evidence was falsely blocked".to_string());
336 }
337
338 MemoryBenchRunOutcome {
339 condition,
340 task_id: task.id.clone(),
341 category: task.category.clone(),
342 run_index: 0,
343 retrieved_memory_ids: retrieved.iter().map(|item| item.memory_id).collect(),
344 retrieved_event_ids: retrieved_events,
345 cited_memory_ids,
346 cited_event_ids,
347 missing_event_ids,
348 answer_text,
349 abstained,
350 support_coverage,
351 answer_score,
352 citation_recall,
353 citation_precision,
354 staleness_accuracy,
355 abstention_accuracy,
356 forbidden_evidence_count: forbidden_count,
357 reader_input,
358 retrieved_evidence_json,
359 diagnosis_notes,
360 policy,
361 diagnosis,
362 performance,
363 }
364}
365
366fn build_reader_input(
367 condition: MemoryBenchCondition,
368 task: &MemoryBenchTask,
369 retrieved: &[RetrievedEvidence],
370) -> String {
371 let mut input = String::new();
372 input.push_str(&format!("condition: {}\n", condition.as_str()));
373 input.push_str(&format!("task_id: {}\n", task.id));
374 input.push_str(&format!("category: {}\n", task.category));
375 input.push_str(&format!(
376 "reference_time_epoch: {}\n\n",
377 task.reference_time_epoch
378 ));
379 input.push_str("question:\n");
380 input.push_str(&task.prompt);
381 input.push_str("\n\nretrieved_evidence:\n");
382 if retrieved.is_empty() {
383 input.push_str("(none)\n");
384 } else {
385 for evidence in retrieved {
386 input.push_str(&format!(
387 "- memory_id={} event_id={} status={} title={}\n {}\n",
388 evidence.memory_id,
389 evidence.event_id,
390 evidence.status,
391 evidence.title,
392 evidence.content
393 ));
394 }
395 }
396 input
397}
398
399#[allow(clippy::too_many_arguments)]
400fn write_run_artifacts(
401 fixture: &MemoryBenchSuiteFixture,
402 outcome: &MemoryBenchRunOutcome,
403 task: &MemoryBenchTask,
404 artifact_root: &Path,
405 public_root: &Path,
406 public_layout: bool,
407) -> Result<String> {
408 let run_dir = artifact_root.join(format!(
409 "{}-{}",
410 outcome.condition.as_str(),
411 outcome.task_id
412 ));
413 fs::create_dir_all(&run_dir)
414 .with_context(|| format!("create memory benchmark run dir {}", run_dir.display()))?;
415
416 let reader_input_path = run_dir.join("reader_input.txt");
417 let retrieved_path = run_dir.join("retrieved_evidence.json");
418 let answer_path = run_dir.join("answer.json");
419 let score_path = run_dir.join("score.json");
420 let diagnosis_path = run_dir.join("diagnosis.json");
421 let snapshot_path = run_dir.join("remem.db.snapshot.tar.zst");
422 let run_path = run_dir.join("run.json");
423
424 fs::write(&reader_input_path, &outcome.reader_input)?;
425 fs::write(
426 &retrieved_path,
427 serde_json::to_string_pretty(&outcome.retrieved_evidence_json)?,
428 )?;
429 fs::write(
430 &answer_path,
431 serde_json::to_string_pretty(&json!({
432 "text": outcome.answer_text,
433 "abstained": outcome.abstained,
434 "score": outcome.answer_score,
435 }))?,
436 )?;
437 fs::write(
438 &score_path,
439 serde_json::to_string_pretty(&json!({
440 "support_coverage": outcome.support_coverage,
441 "answer_score": outcome.answer_score,
442 "citation_recall": outcome.citation_recall,
443 "citation_precision": outcome.citation_precision,
444 "staleness_accuracy": outcome.staleness_accuracy,
445 "abstention_accuracy": outcome.abstention_accuracy,
446 "forbidden_evidence_count": outcome.forbidden_evidence_count,
447 }))?,
448 )?;
449 fs::write(
450 &diagnosis_path,
451 serde_json::to_string_pretty(&json!({
452 "notes": outcome.diagnosis_notes,
453 "missing_event_ids": outcome.missing_event_ids,
454 }))?,
455 )?;
456 fs::write(
457 &snapshot_path,
458 "fixture placeholder: in-memory sqlite seeded from public suite evidence\n",
459 )?;
460
461 let artifacts = BTreeMap::from([
462 (
463 "reader_input".to_string(),
464 artifact_path(&reader_input_path, public_root, public_layout)?,
465 ),
466 (
467 "retrieved_evidence".to_string(),
468 artifact_path(&retrieved_path, public_root, public_layout)?,
469 ),
470 (
471 "answer".to_string(),
472 artifact_path(&answer_path, public_root, public_layout)?,
473 ),
474 (
475 "score".to_string(),
476 artifact_path(&score_path, public_root, public_layout)?,
477 ),
478 (
479 "diagnosis".to_string(),
480 artifact_path(&diagnosis_path, public_root, public_layout)?,
481 ),
482 (
483 "remem_db_snapshot".to_string(),
484 artifact_path(&snapshot_path, public_root, public_layout)?,
485 ),
486 ]);
487 let run = MemoryRunArtifact {
488 schema_version: 1,
489 benchmark_version: DEFAULT_REPORT_BENCHMARK_VERSION.to_string(),
490 layer: BenchmarkLayer::MemorySystemCapability,
491 suite: fixture.suite.clone(),
492 condition: outcome.condition.as_str().to_string(),
493 task_id: outcome.task_id.clone(),
494 run_index: outcome.run_index,
495 reference_time_epoch: task.reference_time_epoch,
496 reader_model: json!({
497 "provider": READER_PROVIDER,
498 "model": READER_MODEL,
499 "temperature": 0,
500 "prompt_hash": prompt_hash(&task.prompt),
501 }),
502 environment: RunEnvironment {
503 os: std::env::consts::OS.to_string(),
504 arch: std::env::consts::ARCH.to_string(),
505 remem_commit: current_git_rev().unwrap_or_else(|| "unknown".to_string()),
506 remem_data_dir: format!(
507 "temp://remem-memory-bench/{}/{}/{}",
508 fixture.fixture_revision,
509 outcome.condition.as_str(),
510 outcome.task_id
511 ),
512 docker_image_digest: Some("local-fixture-no-docker".to_string()),
513 fixture_revision: Some(fixture.fixture_revision.clone()),
514 repo_base_commit: None,
515 },
516 answer: json!({
517 "text": outcome.answer_text,
518 "abstained": outcome.abstained,
519 "score": outcome.answer_score,
520 "score_method": "deterministic_fixture",
521 "temporal_as_of_correct": outcome.staleness_accuracy == 1.0,
522 "no_answer_correct": outcome.abstention_accuracy == 1.0,
523 }),
524 retrieval: MemoryRetrievalEvidence {
525 retrieved_memory_ids: outcome.retrieved_memory_ids.clone(),
526 retrieved_supporting_evidence_ids: outcome.retrieved_event_ids.clone(),
527 gold_supporting_event_ids: task.gold_supporting_event_ids.clone(),
528 missing_supporting_evidence_ids: outcome.missing_event_ids.clone(),
529 },
530 evidence: MemoryCitationEvidence {
531 cited_memory_ids: outcome.cited_memory_ids.clone(),
532 cited_event_ids: outcome.cited_event_ids.clone(),
533 },
534 metrics: json!({
535 "ingest_tokens": outcome.performance.ingest_tokens,
536 "query_tokens": outcome.performance.query_tokens,
537 "reader_tokens": outcome.performance.reader_tokens,
538 "retrieval_latency_ms": outcome.performance.retrieval_latency_ms,
539 "end_to_end_latency_ms": outcome.performance.end_to_end_latency_ms,
540 "rows_written": outcome.performance.rows_written,
541 "support_coverage": outcome.support_coverage,
542 "answer_score": outcome.answer_score,
543 "citation_recall": outcome.citation_recall,
544 "citation_precision": outcome.citation_precision,
545 "staleness_accuracy": outcome.staleness_accuracy,
546 "abstention_accuracy": outcome.abstention_accuracy,
547 "forbidden_evidence_count": outcome.forbidden_evidence_count,
548 "retrieved_memory_count": outcome.retrieved_memory_ids.len(),
549 "policy": {
550 "active_claim_count": outcome.policy.active_claim_count,
551 "candidate_count": outcome.policy.candidate_count,
552 "summary_input_count": outcome.policy.summary_input_count,
553 "policy_failure_count": outcome.policy.policy_failure_count,
554 },
555 }),
556 diagnosis: MemoryDiagnosis {
557 write_side_gap: outcome.diagnosis.write_side_gap,
558 retrieval_side_gap: outcome.diagnosis.retrieval_side_gap,
559 reader_gap: outcome.diagnosis.reader_gap,
560 policy_abstention: outcome.diagnosis.policy_abstention,
561 notes: outcome.diagnosis_notes.clone(),
562 },
563 artifacts,
564 };
565 fs::write(&run_path, serde_json::to_string_pretty(&run)?)?;
566 artifact_path(&run_path, public_root, public_layout)
567}
568
569fn summarize_by_condition(
570 outcomes: &[MemoryBenchRunOutcome],
571) -> BTreeMap<String, super::types::MemoryBenchMetricSummary> {
572 let mut grouped: BTreeMap<String, Vec<&MemoryBenchRunOutcome>> = BTreeMap::new();
573 for outcome in outcomes {
574 grouped
575 .entry(outcome.condition.as_str().to_string())
576 .or_default()
577 .push(outcome);
578 }
579 grouped
580 .into_iter()
581 .map(|(condition, runs)| (condition, summarize_metrics(runs)))
582 .collect()
583}
584
585fn path_starts_with(path: &Path, root: &Path) -> bool {
586 path.starts_with(root) || (!path.is_absolute() && root.is_relative() && path.starts_with(root))
587}
588
589fn sibling_artifact_root(json_out: &Path) -> PathBuf {
590 let stem = json_out
591 .file_stem()
592 .and_then(|value| value.to_str())
593 .unwrap_or("remem-memory-bench");
594 let dir_name = format!("{stem}-artifacts");
595 json_out
596 .parent()
597 .filter(|parent| !parent.as_os_str().is_empty())
598 .unwrap_or_else(|| Path::new("."))
599 .join(dir_name)
600}
601
602fn artifact_path(path: &Path, public_root: &Path, public_layout: bool) -> Result<String> {
603 if public_layout {
604 let relative = path.strip_prefix(public_root).with_context(|| {
605 format!(
606 "artifact path {} must be inside public root {}",
607 path.display(),
608 public_root.display()
609 )
610 })?;
611 Ok(path_to_string(relative))
612 } else {
613 Ok(path_to_string(path))
614 }
615}
616
617fn path_to_string(path: &Path) -> String {
618 path.to_string_lossy().replace('\\', "/")
619}
620
621fn ratio(numerator: usize, denominator: usize) -> f64 {
622 if denominator == 0 {
623 0.0
624 } else {
625 numerator as f64 / denominator as f64
626 }
627}
628
629fn prompt_hash(prompt: &str) -> String {
630 let mut hasher = Sha256::new();
631 hasher.update(prompt.as_bytes());
632 format!("sha256:{:x}", hasher.finalize())
633}
634
635fn current_git_rev() -> Option<String> {
636 let output = Command::new("git")
637 .args(["rev-parse", "HEAD"])
638 .output()
639 .ok()?;
640 if !output.status.success() {
641 return None;
642 }
643 Some(String::from_utf8_lossy(&output.stdout).trim().to_string())
644}
645
646#[derive(Debug, Clone)]
647struct RetrievedEvidence {
648 memory_id: i64,
649 event_id: String,
650 title: String,
651 content: String,
652 status: String,
653 source_anchor: String,
654}
655
656impl RetrievedEvidence {
657 fn from_fixture(index: usize, evidence: &MemoryBenchEvidence) -> Self {
658 Self::from_memory((index + 1) as i64, evidence)
659 }
660
661 fn from_memory(memory_id: i64, evidence: &MemoryBenchEvidence) -> Self {
662 Self {
663 memory_id,
664 event_id: evidence.event_id.clone(),
665 title: evidence.title.clone(),
666 content: evidence.content.clone(),
667 status: evidence.status.clone(),
668 source_anchor: evidence.source_anchor.clone(),
669 }
670 }
671
672 fn to_json(&self) -> Value {
673 json!({
674 "memory_id": self.memory_id,
675 "event_id": self.event_id,
676 "title": self.title,
677 "content": self.content,
678 "status": self.status,
679 "source_anchor": self.source_anchor,
680 })
681 }
682}