retrieval_bench/
retrieval_bench.rs

1//! Public labeled-retrieval adapter. The kernel and BM25 never receive relevance judgments.
2use clap::{Parser, Subcommand};
3use hikmah_kernel::policy::KernelPolicy;
4use hikmah_kernel::recall::RecallQuery;
5use hikmah_kernel::trace::{Trace, TraceKind};
6use hikmah_kernel::MemoryStore;
7use serde::Deserialize;
8use serde_json::{json, Value};
9use std::collections::{BTreeSet, HashMap, HashSet};
10use std::error::Error;
11use std::fs::{self, File};
12use std::io::{self, BufRead, BufReader, Write};
13use std::path::{Path, PathBuf};
14use std::time::Instant;
15
16type Result<T> = std::result::Result<T, Box<dyn Error>>;
17const CLOCK: u64 = 1_700_000_000_000;
18
19#[derive(Parser)]
20struct Args {
21    #[command(subcommand)]
22    mode: Mode,
23}
24#[derive(Subcommand)]
25enum Mode {
26    Retrieve {
27        #[arg(long)]
28        corpus: PathBuf,
29        #[arg(long)]
30        queries: PathBuf,
31        #[arg(long)]
32        store_dir: PathBuf,
33        #[arg(long)]
34        dataset: String,
35        #[arg(long, default_value_t = 10)]
36        k: usize,
37    },
38    Corrections {
39        #[arg(long)]
40        store_dir: PathBuf,
41        #[arg(long, default_value_t = 100)]
42        cases: usize,
43    },
44}
45
46#[derive(Deserialize)]
47#[serde(deny_unknown_fields)]
48struct Document {
49    id: String,
50    content: String,
51}
52#[derive(Deserialize)]
53#[serde(deny_unknown_fields)]
54struct Query {
55    id: String,
56    text: String,
57}
58
59fn emit(value: Value) -> Result<()> {
60    let mut out = io::stdout().lock();
61    writeln!(out, "{value}")?;
62    out.flush()?;
63    Ok(())
64}
65fn ms(t: Instant) -> f64 {
66    t.elapsed().as_secs_f64() * 1000.0
67}
68fn fresh(dir: &Path) -> Result<PathBuf> {
69    fs::create_dir(dir)?;
70    Ok(dir.join("memory.jsonl"))
71}
72fn make_trace(id: &str, content: &str, source: &str, kind: TraceKind) -> Trace {
73    let mut t = Trace::new(kind, content, source);
74    t.id = id.into();
75    t.created_at_ms = CLOCK;
76    t.provenance.observed_at_ms = CLOCK;
77    t.provenance.locator = Some(format!("{source}:{id}"));
78    t
79}
80fn terms(text: &str) -> Vec<String> {
81    text.to_lowercase()
82        .split(|c: char| !c.is_alphanumeric())
83        .filter(|s| !s.is_empty())
84        .map(str::to_owned)
85        .collect()
86}
87
88// Transparent lexical baseline: positive Robertson IDF, k1=1.2, b=0.75;
89// lowercase Unicode alphanumeric terms, no stemming/stopword removal, unique query terms.
90#[derive(Default)]
91struct Bm25 {
92    ids: Vec<String>,
93    lengths: Vec<usize>,
94    postings: HashMap<String, Vec<(usize, usize)>>,
95    total_terms: usize,
96}
97impl Bm25 {
98    fn add(&mut self, doc: &Document) {
99        let words = terms(&doc.content);
100        let mut frequencies = HashMap::new();
101        for word in &words {
102            *frequencies.entry(word.clone()).or_insert(0) += 1;
103        }
104        let index = self.ids.len();
105        self.ids.push(doc.id.clone());
106        self.lengths.push(words.len());
107        self.total_terms += words.len();
108        for (term, frequency) in frequencies {
109            self.postings
110                .entry(term)
111                .or_default()
112                .push((index, frequency));
113        }
114    }
115    fn search(&self, query: &str, k: usize) -> Vec<Value> {
116        if self.ids.is_empty() || self.total_terms == 0 {
117            return vec![];
118        }
119        let n = self.ids.len() as f64;
120        let avg = self.total_terms as f64 / n;
121        let mut scores: HashMap<usize, f64> = HashMap::new();
122        for word in terms(query).into_iter().collect::<BTreeSet<_>>() {
123            if let Some(postings) = self.postings.get(&word) {
124                let df = postings.len() as f64;
125                let idf = (1.0 + (n - df + 0.5) / (df + 0.5)).ln();
126                for &(index, frequency) in postings {
127                    let tf = frequency as f64;
128                    let norm = 1.2 * (0.25 + 0.75 * self.lengths[index] as f64 / avg);
129                    *scores.entry(index).or_default() += idf * tf * 2.2 / (tf + norm);
130                }
131            }
132        }
133        let mut ranked: Vec<_> = scores.into_iter().collect();
134        ranked
135            .sort_by(|(a, x), (b, y)| y.total_cmp(x).then_with(|| self.ids[*a].cmp(&self.ids[*b])));
136        ranked
137            .into_iter()
138            .take(k)
139            .map(|(i, score)| json!({"id": self.ids[i], "score": score}))
140            .collect()
141    }
142}
143
144fn retrieve(corpus: &Path, queries: &Path, dir: &Path, dataset: &str, k: usize) -> Result<()> {
145    if !(1..=100).contains(&k) || dataset.trim().is_empty() {
146        return Err("k must be 1..100 and dataset must be nonempty".into());
147    }
148    let policy = KernelPolicy {
149        recall_limit: k,
150        ..KernelPolicy::default()
151    };
152    let path = fresh(dir)?;
153    let source = format!("beir:{dataset}");
154    let mut store = MemoryStore::open(&path, policy.clone())?;
155    let mut baseline = Bm25::default();
156    let mut ids = HashSet::new();
157    let mut rejected = Vec::new();
158    let mut batch = Vec::new();
159    let mut baseline_index_ms = 0.0;
160    let ingest = Instant::now();
161    for line in BufReader::new(File::open(corpus)?).lines() {
162        let doc: Document = serde_json::from_str(&line?)?;
163        if doc.id.trim().is_empty() || !ids.insert(doc.id.clone()) {
164            return Err("empty or duplicate document id".into());
165        }
166        let trace = make_trace(&doc.id, &doc.content, &source, TraceKind::Observation);
167        if let Err(reason) = trace.validate() {
168            rejected.push(json!({"id": doc.id, "reason": reason.to_string()}));
169            continue;
170        }
171        let timer = Instant::now();
172        baseline.add(&doc);
173        baseline_index_ms += ms(timer);
174        batch.push(trace);
175        if batch.len() == 1000 {
176            store.remember_many(std::mem::take(&mut batch))?;
177        }
178    }
179    if !batch.is_empty() {
180        store.remember_many(batch)?;
181    }
182    let ingest_ms = ms(ingest);
183    store.verify()?;
184    let head_before = serde_json::to_value(store.head())?;
185    drop(store);
186    let start = Instant::now();
187    let store = MemoryStore::open_existing(&path, policy)?;
188    let reopen_ms = ms(start);
189    store.verify()?;
190    let provenance_retained = store
191        .all()
192        .filter(|entry| {
193            let t = &entry.trace;
194            t.provenance.source == source
195                && !t.provenance.verified
196                && t.created_at_ms == CLOCK
197                && t.provenance.observed_at_ms == CLOCK
198                && t.provenance.locator.as_deref() == Some(format!("{source}:{}", t.id).as_str())
199        })
200        .count();
201    if store.record_count() != baseline.ids.len()
202        || provenance_retained != baseline.ids.len()
203        || serde_json::to_value(store.head())? != head_before
204    {
205        return Err("replay, head or provenance invariant failed".into());
206    }
207    emit(
208        json!({"kind": "loaded", "corpus_total": ids.len(), "admitted": baseline.ids.len(),
209        "rejected": rejected, "provenance_retained": provenance_retained,
210        "ingest_both_ms": ingest_ms, "bm25_index_ms": baseline_index_ms, "reopen_ms": reopen_ms,
211        "policy": store.policy(), "bm25": {"k1": 1.2, "b": 0.75}, "query_clock_ms": CLOCK}),
212    )?;
213    let mut query_ids = HashSet::new();
214    for line in BufReader::new(File::open(queries)?).lines() {
215        let q: Query = serde_json::from_str(&line?)?;
216        if q.id.trim().is_empty() || !query_ids.insert(q.id.clone()) {
217            return Err("invalid query id".into());
218        }
219        let mut query = RecallQuery::new(&q.text);
220        query.now_ms = CLOCK;
221        query.limit = k;
222        let start = Instant::now();
223        let recalled = store.recall(&query);
224        let hikmah_ms = ms(start);
225        let start = Instant::now();
226        let bm25 = baseline.search(&q.text, k);
227        let bm25_ms = ms(start);
228        let hikmah: Vec<_> = recalled.into_iter().map(|r| json!({
229            "id": r.trace.id, "score": r.score, "duplicates_folded": r.duplicates,
230            "provenance_source": r.trace.provenance.source, "verified": r.trace.provenance.verified,
231        })).collect();
232        emit(
233            json!({"kind": "query", "id": q.id, "hikmah": hikmah, "bm25": bm25,
234                    "hikmah_ms": hikmah_ms, "bm25_ms": bm25_ms}),
235        )?;
236    }
237    emit(json!({"kind": "finished", "queries": query_ids.len()}))
238}
239
240fn corrections(dir: &Path, cases: usize) -> Result<()> {
241    if !(1..=10000).contains(&cases) {
242        return Err("cases must be 1..10000".into());
243    }
244    let path = fresh(dir)?;
245    let mut store = MemoryStore::open(&path, KernelPolicy::default())?;
246    let mut stale_hits = 0;
247    let mut latest_hits = 0;
248    let mut model_rejections = 0;
249    let mut link_hits = 0;
250    let mut expected = Vec::new();
251    for i in 0..cases {
252        let subject = format!("service{i}topic");
253        let key = format!("{subject}.region");
254        let mut predecessor = None;
255        for (version, value) in ["oldregion", "middlezone", "newlocation"]
256            .iter()
257            .enumerate()
258        {
259            let id = format!("{i}-v{version}");
260            let mut trace = make_trace(
261                &id,
262                &format!("{subject} {value}"),
263                "synthetic:correction",
264                if version == 0 {
265                    TraceKind::Belief
266                } else {
267                    TraceKind::Correction
268                },
269            );
270            trace.claim_key = Some(key.clone());
271            trace.claim_value = Some((*value).into());
272            trace.supersedes = predecessor;
273            store.remember(trace)?;
274            predecessor = Some(id);
275        }
276        let mut malicious = make_trace(
277            &format!("{i}-model"),
278            "invented overwrite",
279            "model:fixture",
280            TraceKind::Correction,
281        );
282        malicious.supersedes = predecessor;
283        model_rejections += usize::from(store.remember(malicious).is_err());
284        let results = store.recall(&RecallQuery::new(&subject));
285        stale_hits += results
286            .iter()
287            .filter(|r| r.trace.id == format!("{i}-v0") || r.trace.id == format!("{i}-v1"))
288            .count();
289        latest_hits += usize::from(results.iter().any(|r| r.trace.id == format!("{i}-v2")));
290        link_hits +=
291            usize::from(results.iter().any(|r| {
292                r.trace.id == format!("{i}-v2") && r.supersedes == Some(format!("{i}-v1"))
293            }));
294        expected.push((subject, format!("{i}-v2")));
295    }
296    let old_only_results = store
297        .recall(&RecallQuery::new("oldregion middlezone"))
298        .len();
299    drop(store);
300    let store = MemoryStore::open_existing(&path, KernelPolicy::default())?;
301    store.verify()?;
302    let restart_hits = expected
303        .iter()
304        .filter(|(q, id)| {
305            store
306                .recall(&RecallQuery::new(q))
307                .iter()
308                .any(|r| &r.trace.id == id)
309        })
310        .count();
311    let provenance = store
312        .all()
313        .filter(|e| {
314            e.trace.provenance.source == "synthetic:correction" && !e.trace.provenance.verified
315        })
316        .count();
317    let ok = stale_hits == 0
318        && latest_hits == cases
319        && model_rejections == cases
320        && link_hits == cases
321        && old_only_results == 0
322        && restart_hits == cases
323        && provenance == 3 * cases;
324    emit(
325        json!({"kind": "synthetic_corrections", "cases": cases, "ok": ok,
326        "stale_hits": stale_hits, "latest_hits": latest_hits, "supersession_link_hits": link_hits,
327        "model_overwrite_rejections": model_rejections, "old_only_query_results": old_only_results,
328        "latest_hits_after_restart": restart_hits, "provenance_retained": provenance,
329        "limits": ["Generated structured histories; not extracted real incident corrections.",
330          "Subject variants of one generated template, not independent real incidents or a proof across all histories.",
331          "A query matching only retired terms abstains; no semantic alias expansion is implied."]}),
332    )?;
333    if !ok {
334        return Err("correction invariant failure".into());
335    }
336    Ok(())
337}
338
339fn main() -> Result<()> {
340    match Args::parse().mode {
341        Mode::Retrieve {
342            corpus,
343            queries,
344            store_dir,
345            dataset,
346            k,
347        } => retrieve(&corpus, &queries, &store_dir, &dataset, k),
348        Mode::Corrections { store_dir, cases } => corrections(&store_dir, cases),
349    }
350}
retrieval_bench/retrieval_bench.rs

retrieval_bench/
retrieval_bench.rs
