diff --git a/.gitignore b/.gitignore index 30a55d1..2bc33af 100644 --- a/.gitignore +++ b/.gitignore @@ -28,9 +28,14 @@ __pycache__/ # Step 0 exploration spikes (not part of the library crate) spikes/ -# Benchmark datasets (large, fetched manually — see benchmarks/*/README.md) -benchmarks/*/data/ +# Benchmark datasets (large, fetched manually — see benchmarks/*/README.md). +# Covers data/ and transformed variants like data_wh40k/. +benchmarks/*/data*/ # Embedding model cache (downloaded ONNX weights + tokenizer, ~90 MB) /models/ .fastembed_cache/ + + +.env +.env.local diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c50fa8..b4ca2c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +### Changed + +- **Lexicon boost is now a configurable dial — `Config::lexicon_boost_clamp`.** The lexicon scorer enters ranking as a multiplier `1.0 + boost.clamp(-c, c)`; `c` was a hardcoded `2.0` (→ up to 3×), now a `Config` field defaulting to a conservative `0.05`. Rationale: fused RRF relevance scores are tightly compressed (rank-1-to-2 gap ~1–2%), so a wide bound let a query-independent importance signal overwhelm relevance — on LongMemEval, the old bound dropped single-session Recall@1 from 0.94 to 0.38 even with semantic search active; `0.05` recovers ~87% of the no-lexicon baseline. `0.0` disables lexicon influence; larger values trade relevance precision for importance weight. Note: this benchmark measures only lexicon's *cost* on factual retrieval, not its *benefit* on persona/importance workloads — the default is harm-minimizing pending a persona-style eval. + ### Breaking Changes - **Lexicon scoring is now opt-in.** `ContextForge::builder(...)` no longer auto-seeds `DefaultEnglishScorer`; by default the engine ranks on relevance (BM25, plus semantic when an embedding model is set) with no lexicon layer. Enable the English importance scorer explicitly with `.with_default_english_scorer()`, and note that `.with_persona_scorer(...)` no longer implies the English layer — call both to compose them. Rationale: the always-on, query-independent importance boost was shown to degrade factual retrieval (LongMemEval Recall@1 dropped 0.71 → 0.17 with the scorer forced on). Callers that want the old behavior — persona/importance use cases, e.g. a chat assistant — add `.with_default_english_scorer()`. diff --git a/Cargo.lock b/Cargo.lock index 054f20c..94e94fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -887,6 +887,27 @@ dependencies = [ "typenum", ] +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "ctr" version = "0.9.2" @@ -2836,6 +2857,18 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "personamem-bench" +version = "0.0.0" +dependencies = [ + "anyhow", + "context-forge", + "csv", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -3090,7 +3123,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3594,7 +3627,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index f1d59e2..ee48e4b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,7 @@ exclude = [ ] [workspace] -members = ["benchmarks/longmemeval"] +members = ["benchmarks/longmemeval", "benchmarks/personamem"] [lib] name = "context_forge" diff --git a/benchmarks/personamem/Cargo.toml b/benchmarks/personamem/Cargo.toml new file mode 100644 index 0000000..6255a36 --- /dev/null +++ b/benchmarks/personamem/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "personamem-bench" +version = "0.0.0" +edition = "2021" +publish = false +description = "PersonaMem retrieval benchmark for context-forge: measures the lexicon-BENEFIT case (surfacing the current/salient user preference) — the complement to the LongMemEval cost benchmark. Dev-only; never published." + +[dependencies] +context-forge = { path = "../..", features = ["semantic", "distill-http"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +csv = "1.4" +anyhow = "1" diff --git a/benchmarks/personamem/src/dataset.rs b/benchmarks/personamem/src/dataset.rs new file mode 100644 index 0000000..54de586 --- /dev/null +++ b/benchmarks/personamem/src/dataset.rs @@ -0,0 +1,123 @@ +//! PersonaMem-v2 dataset: benchmark CSV + per-persona chat-history JSON. +//! +//! Source: HuggingFace `bowen-upenn/PersonaMem-v2` (see README). Verified +//! 2026-07-07 against the real files: `related_conversation_snippet` is a JSON +//! array of `{role, content}` whose `content` matches a chat-history message +//! **exactly**, so retrieval scoring maps snippet → ingested entry by content +//! equality with no normalization. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use serde::Deserialize; + +/// One benchmark row: a query about a persona's current preference, with the +/// evidence turn(s) annotated. Only the columns the retrieval eval uses are +/// deserialized (csv+serde ignores the rest by header name). +#[derive(Debug, Deserialize)] +pub struct Row { + pub persona_id: String, + pub chat_history_32k_link: String, + pub user_query: String, + /// JSON array string of `{role, content}` — the turn(s) where the user + /// implicitly stated the current preference. This is the retrieval ground + /// truth. Parse with [`Row::gold_contents`]. + pub related_conversation_snippet: String, + /// Whether this preference supersedes an earlier one (`"True"`/`"False"`). + /// The `updated == true` slice is the sharpest lexicon-benefit signal + /// (current-vs-stale preference evolution). + #[serde(default)] + pub updated: String, + /// Preference category (stereotypical / anti-stereotypical / neutral / + /// therapy / health). Retained from the schema for future per-type slicing; + /// not yet read by the overall/updated scorer. + #[serde(default)] + #[allow(dead_code)] + pub pref_type: String, +} + +impl Row { + /// True if this row is a preference-*update* case. + #[must_use] + pub fn is_updated(&self) -> bool { + self.updated.eq_ignore_ascii_case("true") + } + + /// The gold evidence contents: the `content` field of each turn in + /// `related_conversation_snippet`. Returns empty if the JSON can't be parsed + /// (row is then excluded from scoring, like an abstention). + #[must_use] + pub fn gold_contents(&self) -> Vec { + #[derive(Deserialize)] + struct Turn { + content: String, + } + serde_json::from_str::>(&self.related_conversation_snippet) + .map(|turns| turns.into_iter().map(|t| t.content).collect()) + .unwrap_or_default() + } +} + +/// A chat-history file: `{ metadata, chat_history: [{role, content}, ...] }`. +#[derive(Debug, Deserialize)] +pub struct ChatHistory { + pub chat_history: Vec, +} + +#[derive(Debug, Deserialize)] +pub struct Message { + /// Speaker role. Present in the schema; the eval ingests content regardless + /// of role, so this is not currently read. + #[allow(dead_code)] + pub role: String, + pub content: String, +} + +/// All rows for one persona, plus its (lazily loaded) chat history. Grouping by +/// persona lets us ingest the ~189-message history **once** and run all of the +/// persona's queries against it (queries are ~1ms; ingest is the cost). +pub struct Persona { + pub persona_id: String, + pub chat_history_link: String, + pub rows: Vec, +} + +/// Load `benchmark.csv` and group rows by persona, preserving first-seen order. +pub fn load_grouped(csv_path: &Path) -> anyhow::Result> { + let mut reader = csv::ReaderBuilder::new() + .flexible(true) + .from_path(csv_path) + .map_err(|e| anyhow::anyhow!("opening {}: {e}", csv_path.display()))?; + + let mut order: Vec = Vec::new(); + let mut by_persona: HashMap = HashMap::new(); + + for result in reader.deserialize() { + let row: Row = result.map_err(|e| anyhow::anyhow!("parsing benchmark.csv row: {e}"))?; + let pid = row.persona_id.clone(); + let entry = by_persona.entry(pid.clone()).or_insert_with(|| { + order.push(pid.clone()); + Persona { + persona_id: pid.clone(), + chat_history_link: row.chat_history_32k_link.clone(), + rows: Vec::new(), + } + }); + entry.rows.push(row); + } + + Ok(order + .into_iter() + .filter_map(|pid| by_persona.remove(&pid)) + .collect()) +} + +/// Load a persona's chat history. `link` is dataset-relative (e.g. +/// `data/chat_history_32k/...json`); `data_root` is where the dataset lives. +pub fn load_chat_history(data_root: &Path, link: &str) -> anyhow::Result { + let path: PathBuf = data_root.join(link); + let bytes = std::fs::read(&path) + .map_err(|e| anyhow::anyhow!("reading chat history {}: {e}", path.display()))?; + serde_json::from_slice(&bytes) + .map_err(|e| anyhow::anyhow!("parsing chat history {}: {e}", path.display())) +} diff --git a/benchmarks/personamem/src/main.rs b/benchmarks/personamem/src/main.rs new file mode 100644 index 0000000..93655be --- /dev/null +++ b/benchmarks/personamem/src/main.rs @@ -0,0 +1,365 @@ +//! PersonaMem Track 1 — deterministic retrieval eval for the lexicon-BENEFIT case. +//! +//! Task: surface the user's *current* preference. Gold = the evidence turn(s) +//! (`related_conversation_snippet`); the 4-way-MC distractors (not used here) are +//! outdated/irrelevant preferences. So higher recall of the gold turn = better +//! importance/salience ranking — exactly what the lexicon is meant to provide, +//! and the complement to LongMemEval (which measures the lexicon's *cost*). +//! +//! Fully deterministic (no reader/judge LLM). Ingests each persona's chat history +//! once, runs all that persona's queries, scores by content match. +//! +//! Usage: +//! cargo run -p personamem-bench --release -- [flags] +//! Flags: +//! --data-root PATH dataset root holding data/chat_history_32k/... (default: CSV's dir) +//! --pipeline bm25|lexicon|semantic|full (default bm25) +//! --clamp F override Config.lexicon_boost_clamp (lexicon/full only) +//! --limit N only the first N personas +//! --embed-dir PATH model cache dir for semantic/full (or CF_EMBED_MODEL_DIR) + +mod dataset; +mod metrics; + +use std::collections::{BTreeMap, HashSet}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use context_forge::{Config, ContextForge, Embedder, FasEmbedder, SaveOptions}; + +use dataset::{load_chat_history, load_grouped, Persona}; +use metrics::Mean; + +const K_VALUES: &[usize] = &[1, 3, 5, 10]; +const BUDGETS: &[usize] = &[500, 1000, 2000]; +const UNBOUNDED_BUDGET: usize = 100_000_000; +const MSG_KIND: &str = "message"; + +#[derive(Clone, Copy)] +enum Pipeline { + Bm25, + Lexicon, + Semantic, + Full, +} + +impl Pipeline { + fn parse(s: &str) -> anyhow::Result { + match s { + "bm25" => Ok(Pipeline::Bm25), + "lexicon" => Ok(Pipeline::Lexicon), + "semantic" => Ok(Pipeline::Semantic), + "full" => Ok(Pipeline::Full), + other => Err(anyhow::anyhow!("unknown pipeline {other:?}")), + } + } + fn label(self) -> &'static str { + match self { + Pipeline::Bm25 => "bm25", + Pipeline::Lexicon => "lexicon", + Pipeline::Semantic => "semantic", + Pipeline::Full => "full", + } + } + fn wants_embedder(self) -> bool { + matches!(self, Pipeline::Semantic | Pipeline::Full) + } + fn wants_lexicon(self) -> bool { + matches!(self, Pipeline::Lexicon | Pipeline::Full) + } +} + +struct Args { + csv: PathBuf, + data_root: PathBuf, + pipeline: Pipeline, + clamp: Option, + limit: Option, + embed_dir: Option, +} + +fn parse_args() -> anyhow::Result { + let mut csv = None; + let mut data_root = None; + let mut pipeline = Pipeline::Bm25; + let mut clamp = None; + let mut limit = None; + let mut embed_dir = std::env::var("CF_EMBED_MODEL_DIR").ok().map(PathBuf::from); + + let mut it = std::env::args().skip(1); + while let Some(arg) = it.next() { + match arg.as_str() { + "--data-root" => data_root = Some(PathBuf::from(next(&mut it, "--data-root")?)), + "--pipeline" => pipeline = Pipeline::parse(&next(&mut it, "--pipeline")?)?, + "--clamp" => clamp = Some(next(&mut it, "--clamp")?.parse()?), + "--limit" => limit = Some(next(&mut it, "--limit")?.parse()?), + "--embed-dir" => embed_dir = Some(PathBuf::from(next(&mut it, "--embed-dir")?)), + other if other.starts_with("--") => { + return Err(anyhow::anyhow!("unknown flag {other:?}")) + } + positional => csv = Some(PathBuf::from(positional)), + } + } + + let csv = csv.ok_or_else(|| anyhow::anyhow!("missing benchmark.csv path argument"))?; + // Default data-root = the CSV's parent dir (dataset links are relative to it). + let data_root = data_root + .or_else(|| csv.parent().map(Path::to_path_buf)) + .unwrap_or_else(|| PathBuf::from(".")); + Ok(Args { + csv, + data_root, + pipeline, + clamp, + limit, + embed_dir, + }) +} + +fn next(it: &mut impl Iterator, flag: &str) -> anyhow::Result { + it.next() + .ok_or_else(|| anyhow::anyhow!("{flag} requires a value")) +} + +/// Build a fresh engine for one persona (fresh in-memory DB → per-persona BM25 +/// stats + isolation), injecting the shared embedder and configured clamp. +async fn build_forge( + args: &Args, + shared_embedder: &Option>, +) -> anyhow::Result { + let mut config = Config::default(); + config.db_path = PathBuf::from(":memory:"); + config.max_entries = 100_000_000; + if let Some(c) = args.clamp { + config.lexicon_boost_clamp = c; + } + + let mut builder = ContextForge::builder(config); + if args.pipeline.wants_lexicon() { + builder = builder.with_default_english_scorer(); + } + if args.pipeline.wants_embedder() { + let emb = shared_embedder + .clone() + .ok_or_else(|| anyhow::anyhow!("this pipeline requires an embedder (--embed-dir)"))?; + builder = builder.with_embedder(emb); + } + Ok(builder.build().await?) +} + +#[derive(Default, Clone, Copy)] +struct Timing { + build: Duration, + ingest: Duration, + query: Duration, +} + +/// Aggregated recall for a slice of rows. +#[derive(Default)] +struct Agg { + recall_at: BTreeMap, + recall_budget: BTreeMap, + scored: usize, + skipped: usize, +} + +impl Agg { + fn record( + &mut self, + ranked: &[String], + budgets: &BTreeMap>, + gold: &HashSet, + ) { + for &k in K_VALUES { + self.recall_at + .entry(k) + .or_default() + .push(metrics::recall_at_k(ranked, gold, k)); + } + for (&b, set) in budgets { + self.recall_budget + .entry(b) + .or_default() + .push(metrics::recall_in_set(set, gold)); + } + } +} + +#[tokio::main(flavor = "multi_thread")] +async fn main() -> anyhow::Result<()> { + let args = parse_args()?; + let mut personas = load_grouped(&args.csv)?; + if let Some(n) = args.limit { + personas.truncate(n); + } + let total_rows: usize = personas.iter().map(|p| p.rows.len()).sum(); + + eprintln!( + "PersonaMem retrieval benchmark (Track 1)\n csv: {}\n personas: {} rows: {}\n pipeline: {}{}", + args.csv.display(), + personas.len(), + total_rows, + args.pipeline.label(), + args.clamp.map_or(String::new(), |c| format!(" clamp: {c}")), + ); + + let shared_embedder: Option> = if args.pipeline.wants_embedder() { + let dir = args + .embed_dir + .clone() + .ok_or_else(|| anyhow::anyhow!("--embed-dir (or CF_EMBED_MODEL_DIR) required"))?; + eprintln!(" loading embedding model once from {}...", dir.display()); + Some(Arc::new(FasEmbedder::new(&dir)?)) + } else { + None + }; + + let mut overall = Agg::default(); + let mut updated = Agg::default(); // preference-evolution slice (the key lexicon signal) + let mut total = Timing::default(); + let mut errors = 0usize; + + for (i, persona) in personas.iter().enumerate() { + match run_persona(&args, persona, &shared_embedder, &mut overall, &mut updated).await { + Ok(t) => { + total.build += t.build; + total.ingest += t.ingest; + total.query += t.query; + } + Err(e) => { + errors += 1; + eprintln!(" [!] persona {} ({}) failed: {e}", i, persona.persona_id); + } + } + if (i + 1) % 50 == 0 { + eprintln!(" ...{}/{} personas", i + 1, personas.len()); + } + } + + print_agg("OVERALL", &overall); + print_agg("UPDATED (preference-evolution)", &updated); + println!( + "\nphase totals: build={:.1}s ingest={:.1}s query={:.1}s errors={errors}", + total.build.as_secs_f64(), + total.ingest.as_secs_f64(), + total.query.as_secs_f64(), + ); + Ok(()) +} + +/// Ingest one persona's chat history once, then run all its queries. +async fn run_persona( + args: &Args, + persona: &Persona, + shared_embedder: &Option>, + overall: &mut Agg, + updated: &mut Agg, +) -> anyhow::Result { + let scope = persona.persona_id.as_str(); + + let t = Instant::now(); + let forge = build_forge(args, shared_embedder).await?; + let build = t.elapsed(); + + // Ingest every chat-history message as an entry (one batch → one index commit). + let history = load_chat_history(&args.data_root, &persona.chat_history_link)?; + let t = Instant::now(); + let items: Vec<(String, String, SaveOptions)> = history + .chat_history + .iter() + .filter(|m| !m.content.trim().is_empty()) + .map(|m| { + ( + m.content.clone(), + MSG_KIND.to_owned(), + SaveOptions { + scope: Some(scope.to_owned()), + ..Default::default() + }, + ) + }) + .collect(); + forge.save_batch(&items).await?; + let ingest = t.elapsed(); + + // Set of contents actually ingested for this persona, so rows whose gold + // evidence isn't in *this* history version (e.g. 128k-only evidence in a + // 32k run) can be skipped rather than scored as impossible zeros. Matches + // the transform script's filtering so English and WH40k runs use the same + // honest row set. + let ingested: HashSet<&str> = items.iter().map(|(c, _, _)| c.as_str()).collect(); + + // Run each row's query against this persona's store. + let t = Instant::now(); + for row in &persona.rows { + let gold: HashSet = row.gold_contents().into_iter().collect(); + if gold.is_empty() { + overall.skipped += 1; + continue; + } + // Skip rows whose evidence was never ingested (unanswerable here). + if !gold.iter().any(|g| ingested.contains(g.as_str())) { + overall.skipped += 1; + continue; + } + + let ranked = forge + .query(&row.user_query, Some(scope), UNBOUNDED_BUDGET) + .await?; + let ranked_contents: Vec = ranked.into_iter().map(|e| e.content).collect(); + + let mut budgets = BTreeMap::new(); + for &b in BUDGETS { + let entries = forge.query(&row.user_query, Some(scope), b).await?; + budgets.insert( + b, + entries + .into_iter() + .map(|e| e.content) + .collect::>(), + ); + } + + overall.scored += 1; + overall.record(&ranked_contents, &budgets, &gold); + if row.is_updated() { + updated.scored += 1; + updated.record(&ranked_contents, &budgets, &gold); + } + } + let query = t.elapsed(); + + Ok(Timing { + build, + ingest, + query, + }) +} + +fn print_agg(name: &str, agg: &Agg) { + if agg.scored == 0 { + println!("\n[{name}] no scored rows (skipped={})", agg.skipped); + return; + } + let sampled = agg.recall_at.get(&K_VALUES[0]).map_or(0, Mean::count); + println!( + "\n[{name}] (scored={}, skipped={}, sampled={sampled})", + agg.scored, agg.skipped + ); + let recalls: Vec = K_VALUES + .iter() + .map(|k| format!("R@{k}={:.3}", agg.recall_at.get(k).map_or(0.0, Mean::get))) + .collect(); + let budgets: Vec = BUDGETS + .iter() + .map(|b| { + format!( + "R@{b}tok={:.3}", + agg.recall_budget.get(b).map_or(0.0, Mean::get) + ) + }) + .collect(); + println!(" {}", recalls.join(" ")); + println!(" {}", budgets.join(" ")); +} diff --git a/benchmarks/personamem/src/metrics.rs b/benchmarks/personamem/src/metrics.rs new file mode 100644 index 0000000..383419d --- /dev/null +++ b/benchmarks/personamem/src/metrics.rs @@ -0,0 +1,134 @@ +//! Retrieval metrics for PersonaMem, scored by **content equality**. +//! +//! Unlike LongMemEval (gold = session ids), PersonaMem's gold is the evidence +//! turn *content* (`related_conversation_snippet`), which matches an ingested +//! message's content exactly. So a retrieved entry is "gold" iff its content is +//! in the gold-content set. Fully deterministic — no reader or judge LLM. + +use std::collections::HashSet; + +/// Deduplicate strings preserving first-seen order (a retrieved list may repeat +/// content across entries; collapse to first occurrence for rank cutoffs). +#[must_use] +pub fn distinct_in_order(items: &[String]) -> Vec { + let mut seen = HashSet::new(); + let mut out = Vec::new(); + for s in items { + if seen.insert(s.clone()) { + out.push(s.clone()); + } + } + out +} + +/// Recall@k: fraction of gold contents present in the top-k distinct retrieved +/// contents. `None` when there is no gold set (row excluded from averaging). +#[must_use] +pub fn recall_at_k(ranked_contents: &[String], gold: &HashSet, k: usize) -> Option { + if gold.is_empty() { + return None; + } + let top_k: HashSet = distinct_in_order(ranked_contents) + .into_iter() + .take(k) + .collect(); + let hits = gold.iter().filter(|g| top_k.contains(*g)).count(); + Some(hits as f64 / gold.len() as f64) +} + +/// Recall@budget: fraction of gold contents present among the contents that +/// survived a token-budgeted `query`. Order-independent. +#[must_use] +pub fn recall_in_set(retrieved: &HashSet, gold: &HashSet) -> Option { + if gold.is_empty() { + return None; + } + let hits = gold.iter().filter(|g| retrieved.contains(*g)).count(); + Some(hits as f64 / gold.len() as f64) +} + +/// Running mean that ignores `None` samples (rows with no parseable gold). +#[derive(Debug, Default, Clone)] +pub struct Mean { + sum: f64, + n: usize, +} + +impl Mean { + pub fn push(&mut self, sample: Option) { + if let Some(v) = sample { + self.sum += v; + self.n += 1; + } + } + + #[must_use] + pub fn get(&self) -> f64 { + if self.n == 0 { + 0.0 + } else { + self.sum / self.n as f64 + } + } + + #[must_use] + pub fn count(&self) -> usize { + self.n + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn set(items: &[&str]) -> HashSet { + items.iter().map(|s| (*s).to_owned()).collect() + } + fn list(items: &[&str]) -> Vec { + items.iter().map(|s| (*s).to_owned()).collect() + } + + #[test] + fn recall_hits_gold_in_top_k() { + let ranked = list(&["msg-a", "msg-b", "gold-1"]); + assert_eq!(recall_at_k(&ranked, &set(&["gold-1"]), 3), Some(1.0)); + // gold not in top-2 → 0. + assert_eq!(recall_at_k(&ranked, &set(&["gold-1"]), 2), Some(0.0)); + } + + #[test] + fn recall_partial_multi_gold() { + let ranked = list(&["g1", "x", "y"]); + // two gold turns, one retrieved in top-3 → 0.5. + assert_eq!(recall_at_k(&ranked, &set(&["g1", "g2"]), 3), Some(0.5)); + } + + #[test] + fn recall_dedups_before_cutoff() { + let ranked = list(&["dup", "dup", "dup", "gold"]); + assert_eq!(recall_at_k(&ranked, &set(&["gold"]), 2), Some(1.0)); + } + + #[test] + fn empty_gold_is_none() { + let ranked = list(&["a"]); + assert_eq!(recall_at_k(&ranked, &set(&[]), 5), None); + assert_eq!(recall_in_set(&set(&["a"]), &set(&[])), None); + } + + #[test] + fn recall_in_set_ignores_order() { + let retrieved = set(&["b", "gold", "a"]); + assert_eq!(recall_in_set(&retrieved, &set(&["gold"])), Some(1.0)); + } + + #[test] + fn mean_ignores_none() { + let mut m = Mean::default(); + m.push(Some(1.0)); + m.push(None); + m.push(Some(0.0)); + assert_eq!(m.count(), 2); + assert!((m.get() - 0.5).abs() < 1e-9); + } +} diff --git a/scripts/transform_personamem_wh40k.py b/scripts/transform_personamem_wh40k.py new file mode 100644 index 0000000..18ce81e --- /dev/null +++ b/scripts/transform_personamem_wh40k.py @@ -0,0 +1,418 @@ +#!/usr/bin/env python3 +"""Transform local PersonaMem 32k data into a WH40k/Mechanicus dialect variant. + +This writes a parallel dataset tree and preserves benchmark gold linkage by deriving +`related_conversation_snippet` from the transformed chat-history turns, not by +separately rewriting the snippet text. + +Expected endpoint: OpenAI-compatible /chat/completions. +""" + +from __future__ import annotations + +import argparse +import ast +import csv +import json +import os +import shutil +import sys +import time +import urllib.error +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +DEFAULT_IN_ROOT = Path("benchmarks/personamem/data") +DEFAULT_OUT_ROOT = Path("benchmarks/personamem/data_wh40k") +DEFAULT_BASE_URL = os.environ.get("OPENAI_BASE_URL", "http://localhost:1234/v1") +DEFAULT_MODEL = os.environ.get("OPENAI_MODEL", os.environ.get("MODEL", "local-model")) + +PROMPT = r""" +PURPOSE — read before transforming + +This text is being transformed to build an evaluation benchmark for a memory-retrieval +system. The benchmark measures whether a "domain lexicon" (a weighted vocabulary that +boosts importance-signalling words) helps a search engine surface the ONE conversation +turn where a user implicitly revealed a preference, out of a long history full of +distractor turns. + +The source data (PersonaMem) is in plain English. We are re-skinning it into Warhammer +40,000 Imperial/Adeptus-Mechanicus dialect (a "Cawl Inferior" persona voice) WITHOUT +changing its structure, so we can test whether a WH40k-tuned lexicon helps on WH40k- +flavored data. The benchmark scores by exact-matching the evidence turn, and its whole +validity depends on three things you control: + - The revealed preference must stay exactly as IMPLICIT as in the original. + - The MEANING must be preserved exactly. + - You must NOT know or infer which turn is "the important one." +Treat every turn as if it might be the evidence turn or a distractor; you cannot tell, +and you must not try. + +TASK + +You will receive a JSON array of turns, each {"i": , "role": , "content": }. +Rewrite each turn's content into WH40k Imperial/Adeptus-Mechanicus dialect under the +rules below. Return ONLY a JSON array of {"i": , "content": } with +the SAME indices, in the SAME order, one object per input turn — no other keys, no prose. + +HARD CONSTRAINTS — violating any invalidates the data: +1. PRESERVE MEANING EXACTLY. Every fact, preference, entity, number, and intent must + survive unchanged. Do not add facts. Do not drop facts. +2. PRESERVE IMPLICITNESS. Keep indirectly-stated preferences exactly as indirect. NEVER + make an implied preference more explicit, obvious, or easier to find. Do not + summarize, flag, or emphasize any turn's "point." +3. PRESERVE STRUCTURE. Same role meaning, roughly the same length, same number of + distinct statements, same conversational function per turn. +4. Judge each turn ONLY on its own text. You do NOT know and must NOT guess which turns + matter. Treat none as special. +5. Preserve the index i of every turn exactly; return exactly as many objects as received. + +REGISTER RULE: +- Vary devotional/reverent intensity by each turn's OWN emotional register, as a devout + character naturally would: personal/heartfelt/committal/emphatic statements → heavier + reverent language; neutral/transactional/procedural statements → light, plain reskin. +- Base this ONLY on how the speaker feels in that text. + +VOCABULARY (draw naturally; never force terms where they don't fit): +- Reverence/importance: the Omnissiah, the God-Emperor, the Machine-God, the Throne, + blessed, sacred, sanctified, by the Motive Force, praise be, it is the Emperor's will. +- Commitment/affirmation: "I so vow," "it shall be done," "by the Omnissiah I affirm," + "let it be recorded." +- Tech/framing: cogitator, data-slate, machine-spirit, augur, vox, lexmechanic, rites, litany. +- Negation/doubt: "the Emperor frowns upon this," "corrupted by the warp," "the + machine-spirit is silent." + +OUTPUT: ONLY the JSON array described above. No markdown fences, no commentary. +""".strip() + + +@dataclass(frozen=True) +class TurnRef: + role: str + content: str + history_index: int + + +class OpenAICompatClient: + def __init__(self, base_url: str, api_key: str, model: str, timeout: int, temperature: float): + self.base_url = base_url.rstrip("/") + self.api_key = api_key + self.model = model + self.timeout = timeout + self.temperature = temperature + + def transform_batch(self, items: list[dict[str, Any]], retries: int = 4) -> dict[int, str]: + body = { + "model": self.model, + "temperature": self.temperature, + "messages": [ + {"role": "system", "content": PROMPT}, + {"role": "user", "content": json.dumps(items, ensure_ascii=False)}, + ], + } + data = json.dumps(body, ensure_ascii=False).encode("utf-8") + req = urllib.request.Request( + f"{self.base_url}/chat/completions", + data=data, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {self.api_key}", + }, + method="POST", + ) + + last_err: Exception | None = None + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + payload = json.loads(resp.read().decode("utf-8")) + text = payload["choices"][0]["message"]["content"].strip() + parsed = parse_json_array(text) + return validate_response(items, parsed) + except Exception as e: # noqa: BLE001 - retain endpoint/parser errors for retry context. + last_err = e + if attempt == retries: + break + time.sleep(2**attempt) + raise RuntimeError(f"LLM batch failed after {retries + 1} attempts: {last_err}") + + +def parse_json_array(text: str) -> list[Any]: + if text.startswith("```"): + lines = text.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].startswith("```"): + lines = lines[:-1] + text = "\n".join(lines).strip() + return json.loads(text) + + +def validate_response(items: list[dict[str, Any]], parsed: list[Any]) -> dict[int, str]: + if not isinstance(parsed, list): + raise ValueError("response is not a JSON array") + expected = [int(x["i"]) for x in items] + got: list[int] = [] + out: dict[int, str] = {} + for obj in parsed: + if not isinstance(obj, dict) or set(obj.keys()) != {"i", "content"}: + raise ValueError(f"bad response object: {obj!r}") + i = int(obj["i"]) + if not isinstance(obj["content"], str): + raise ValueError(f"content for {i} is not a string") + got.append(i) + out[i] = obj["content"] + if got != expected: + raise ValueError(f"index/order mismatch: expected {expected}, got {got}") + return out + + +def batched(xs: list[dict[str, Any]], n: int): + for start in range(0, len(xs), n): + yield xs[start : start + n] + + +def transform_items(client: OpenAICompatClient, items: list[dict[str, Any]], batch_size: int) -> dict[int, str]: + transformed: dict[int, str] = {} + for batch in batched(items, batch_size): + transformed.update(client.transform_batch(batch)) + return transformed + + +def load_history(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def find_snippet_refs(history: list[dict[str, str]], snippet_text: str) -> list[TurnRef]: + turns = json.loads(snippet_text) + wanted = [(t.get("role", ""), t["content"]) for t in turns] + + # Prefer an exact contiguous role+content match, preserving the snippet sequence. + for start in range(0, len(history) - len(wanted) + 1): + if all( + history[start + off].get("role") == role and history[start + off].get("content") == content + for off, (role, content) in enumerate(wanted) + ): + return [TurnRef(role, content, start + off) for off, (role, content) in enumerate(wanted)] + + # Fallback: exact role+content per snippet turn. This handles non-contiguous annotations. + refs: list[TurnRef] = [] + used: set[int] = set() + for role, content in wanted: + for idx, msg in enumerate(history): + if idx not in used and msg.get("role") == role and msg.get("content") == content: + refs.append(TurnRef(role, content, idx)) + used.add(idx) + break + else: + raise ValueError("snippet turn not found in history") + return refs + + +def parse_user_query_cell(cell: str) -> tuple[Any, str | None]: + """Return parsed cell and content if cell is a {'role','content'} literal/JSON.""" + for loader in (json.loads, ast.literal_eval): + try: + obj = loader(cell) + if isinstance(obj, dict) and isinstance(obj.get("content"), str): + return obj, obj["content"] + except Exception: + pass + return cell, cell + + +def dump_user_query_cell(original_obj: Any, transformed_content: str) -> str: + if isinstance(original_obj, dict) and "content" in original_obj: + obj = dict(original_obj) + obj["content"] = transformed_content + # Input CSV uses Python-literal dicts; preserving that shape avoids changing parser assumptions. + return repr(obj) + return transformed_content + + +def transform_csv_fields_for_persona( + client: OpenAICompatClient, + rows: list[dict[str, str]], + batch_size: int, +) -> None: + items: list[dict[str, Any]] = [] + metadata: dict[int, tuple[dict[str, str], str, int | None]] = {} + next_i = 0 + + for row in rows: + parsed_query, query_content = parse_user_query_cell(row["user_query"]) + items.append({"i": next_i, "role": "user", "content": query_content}) + metadata[next_i] = (row, "user_query", None) + row["__parsed_user_query"] = parsed_query # type: ignore[assignment] + next_i += 1 + + items.append({"i": next_i, "role": "assistant", "content": row["correct_answer"]}) + metadata[next_i] = (row, "correct_answer", None) + next_i += 1 + + incorrect = json.loads(row["incorrect_answers"]) + row["__incorrect_answers"] = incorrect # type: ignore[assignment] + for j, answer in enumerate(incorrect): + items.append({"i": next_i, "role": "assistant", "content": answer}) + metadata[next_i] = (row, "incorrect_answers", j) + next_i += 1 + + # Keep each row's user_query + correct + three incorrect answers in the same LLM call. + # PersonaMem has exactly 3 incorrect options, so each row contributes 5 items. + row_width = 5 + grouped_batch_size = max(1, batch_size // row_width) * row_width + transformed = transform_items(client, items, grouped_batch_size) + + for i, text in transformed.items(): + row, field, offset = metadata[i] + if field == "user_query": + parsed = row.pop("__parsed_user_query") # type: ignore[arg-type] + row["user_query"] = dump_user_query_cell(parsed, text) + elif field == "correct_answer": + row["correct_answer"] = text + elif field == "incorrect_answers": + answers = row["__incorrect_answers"] # type: ignore[index] + answers[offset] = text # type: ignore[index] + else: + raise AssertionError(field) + + for row in rows: + if "__incorrect_answers" in row: + row["incorrect_answers"] = json.dumps(row.pop("__incorrect_answers"), ensure_ascii=False) # type: ignore[arg-type] + row.pop("__parsed_user_query", None) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--input-root", type=Path, default=DEFAULT_IN_ROOT) + ap.add_argument("--output-root", type=Path, default=DEFAULT_OUT_ROOT) + ap.add_argument("--base-url", default=DEFAULT_BASE_URL) + ap.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY", "not-needed")) + ap.add_argument("--model", default=DEFAULT_MODEL) + ap.add_argument("--batch-size", type=int, default=25) + ap.add_argument("--timeout", type=int, default=300) + ap.add_argument("--temperature", type=float, default=0.2) + ap.add_argument("--limit-personas", type=int) + ap.add_argument("--force", action="store_true", help="allow replacing an existing output tree") + ap.add_argument( + "--fail-missing-snippet", + action="store_true", + help="fail instead of skipping rows whose gold snippet is not present in the local 32k history", + ) + args = ap.parse_args() + + in_csv = args.input_root / "benchmark.csv" + out_csv = args.output_root / "benchmark.csv" + in_history_root = args.input_root / "data" / "chat_history_32k" + out_history_root = args.output_root / "data" / "chat_history_32k" + + if args.output_root.exists(): + if not args.force: + raise SystemExit(f"output root already exists: {args.output_root} (use --force to replace)") + shutil.rmtree(args.output_root) + out_history_root.mkdir(parents=True, exist_ok=True) + + with in_csv.open("r", encoding="utf-8", newline="") as f: + reader = csv.DictReader(f) + fieldnames = reader.fieldnames + if fieldnames is None: + raise SystemExit("benchmark.csv has no header") + all_rows = list(reader) + + existing_links: set[str] = set() + for row in all_rows: + link = row["chat_history_32k_link"] + if (args.input_root / link).exists(): + existing_links.add(link) + + links = sorted(existing_links) + if args.limit_personas is not None: + links = links[: args.limit_personas] + link_set = set(links) + + rows_by_link: dict[str, list[dict[str, str]]] = {link: [] for link in links} + for row in all_rows: + link = row["chat_history_32k_link"] + if link in link_set: + rows_by_link[link].append(row) + + client = OpenAICompatClient(args.base_url, args.api_key, args.model, args.timeout, args.temperature) + print(f"Transforming {len(links)} local personas into {args.output_root}", file=sys.stderr) + + included_ids: set[int] = set() + skipped_missing_snippet = 0 + for n, link in enumerate(links, start=1): + in_path = args.input_root / link + out_path = args.output_root / link + history_obj = load_history(in_path) + history = history_obj["chat_history"] + persona_rows = rows_by_link[link] + + # Validate and store gold refs before spending LLM calls on rows that cannot be scored + # against the local 32k history. Some locally available personas still have rows whose + # annotated evidence turn only exists in the 128k history. + valid_rows: list[dict[str, str]] = [] + row_refs: dict[int, list[TurnRef]] = {} + for row in persona_rows: + try: + row_refs[id(row)] = find_snippet_refs(history, row["related_conversation_snippet"]) + valid_rows.append(row) + except Exception as e: # noqa: BLE001 - choose skip/fail from CLI. + if args.fail_missing_snippet: + raise RuntimeError( + f"gold snippet for persona {row['persona_id']} not found in {link}" + ) from e + skipped_missing_snippet += 1 + + print( + f"[{n}/{len(links)}] {link}: {len(history)} messages, " + f"{len(valid_rows)}/{len(persona_rows)} scorable rows", + file=sys.stderr, + ) + + items = [ + {"i": i, "role": msg.get("role", ""), "content": msg.get("content", "")} + for i, msg in enumerate(history) + ] + transformed_history = transform_items(client, items, args.batch_size) + + for i, msg in enumerate(history): + msg["content"] = transformed_history[i] + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(history_obj, ensure_ascii=False, indent=2), encoding="utf-8") + + transform_csv_fields_for_persona(client, valid_rows, args.batch_size) + + # Relink gold snippets from transformed history turns. Never free-rewrite this field. + for row in valid_rows: + refs = row_refs[id(row)] + row["related_conversation_snippet"] = json.dumps( + [ + { + "role": ref.role, + "content": transformed_history[ref.history_index], + } + for ref in refs + ], + ensure_ascii=False, + ) + included_ids.add(id(row)) + + output_rows = [{k: row[k] for k in fieldnames} for row in all_rows if id(row) in included_ids] + + out_csv.parent.mkdir(parents=True, exist_ok=True) + with out_csv.open("w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(output_rows) + + print(f"Wrote {len(output_rows)} rows: {out_csv}", file=sys.stderr) + if skipped_missing_snippet: + print(f"Skipped {skipped_missing_snippet} rows whose gold snippet is absent from local 32k histories", file=sys.stderr) + print(f"Wrote histories: {out_history_root}", file=sys.stderr) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/config.rs b/src/config.rs index fe6f72c..84a1f03 100644 --- a/src/config.rs +++ b/src/config.rs @@ -10,6 +10,17 @@ pub const DEFAULT_RECENCY_HALF_LIFE_HOURS: f64 = 72.0; /// Default recency half-life in seconds (72 hours). pub const DEFAULT_RECENCY_HALF_LIFE_SECS: f64 = DEFAULT_RECENCY_HALF_LIFE_HOURS * 3600.0; +/// Default maximum absolute lexicon boost applied to the score multiplier. +/// +/// Deliberately small. The lexicon enters scoring as a multiplier +/// `1.0 + boost.clamp(-c, c)`, and fused RRF relevance scores are tightly +/// compressed (the rank-1-to-rank-2 gap is ~1–2%), so a wide bound lets a +/// query-*independent* importance signal overwhelm relevance and bury the +/// evidence. On a pure-relevance benchmark (`LongMemEval`) `0.05` recovered ~87% +/// of the no-lexicon baseline's Recall@1 while still expressing modest +/// importance; a larger value trades more relevance for more importance weight. +pub const DEFAULT_LEXICON_BOOST_CLAMP: f64 = 0.05; + /// Runtime configuration for the context engine. #[derive(Debug, Clone, Serialize, Deserialize)] #[non_exhaustive] @@ -31,6 +42,16 @@ pub struct Config { /// Controls how fast older entries lose relevance. A value of 259200 (72 hours) /// means an entry's score halves every 3 days. pub recency_half_life_secs: f64, + /// Maximum absolute lexicon boost, i.e. the relevance↔importance dial. + /// + /// The lexicon scorer's output enters ranking as a multiplier + /// `1.0 + boost.clamp(-c, c)` on the fused relevance score, where `c` is this + /// value. `0.0` disables lexicon influence entirely (pure relevance); + /// larger values give importance signals more weight at the cost of + /// relevance precision. See [`DEFAULT_LEXICON_BOOST_CLAMP`] for the rationale + /// behind the conservative default. Negative or non-finite values are + /// treated as the default. + pub lexicon_boost_clamp: f64, /// Secret-scrubbing configuration applied to entry content at save time. pub scrub: ScrubConfig, } @@ -47,6 +68,7 @@ impl Default for Config { db_path: PathBuf::from(":memory:"), eviction_policy: EvictionPolicy::Lru, recency_half_life_secs: DEFAULT_RECENCY_HALF_LIFE_SECS, + lexicon_boost_clamp: DEFAULT_LEXICON_BOOST_CLAMP, scrub: ScrubConfig::default(), } } diff --git a/src/engine.rs b/src/engine.rs index e89ea11..e8ea48c 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -76,6 +76,9 @@ impl ContextEngine { if !config.recency_half_life_secs.is_finite() || config.recency_half_life_secs <= 0.0 { config.recency_half_life_secs = DEFAULT_RECENCY_HALF_LIFE_SECS; } + if !config.lexicon_boost_clamp.is_finite() || config.lexicon_boost_clamp < 0.0 { + config.lexicon_boost_clamp = crate::config::DEFAULT_LEXICON_BOOST_CLAMP; + } Self { storage, @@ -203,7 +206,8 @@ impl ContextEngine { .scorer .as_ref() .map_or(0.0_f32, |s| s.score(&entry, query)); - let final_score = rrf * decay * (1.0 + f64::from(boost).clamp(-1.0, 2.0)); + let c = self.config.lexicon_boost_clamp; + let final_score = rrf * decay * (1.0 + f64::from(boost).clamp(-c, c)); (final_score, entry) }) .collect(); diff --git a/tests/lexicon_scoring.rs b/tests/lexicon_scoring.rs index 6cea52a..c01e3f2 100644 --- a/tests/lexicon_scoring.rs +++ b/tests/lexicon_scoring.rs @@ -93,7 +93,12 @@ patterns = ["the emperor frowns upon this"] "#; let persona: ConfigLexiconScorer = persona_toml.parse().unwrap(); - let cf = ContextForge::builder(in_memory_config()) + // This test verifies boost *magnitudes* order correctly (persona 1.4 > + // english 0.5 > neutral), which needs a clamp wide enough to distinguish + // them — the conservative default (0.05) would flatten both to a tie. + let mut config = in_memory_config(); + config.lexicon_boost_clamp = 2.0; + let cf = ContextForge::builder(config) .with_default_english_scorer() .with_persona_scorer(persona) .build()