From c6abf9beebd06ca4bd48ce0f7c4832dd19650ee3 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Mon, 29 Jun 2026 11:37:11 +0000 Subject: [PATCH] implement session-first memory retrieval profile --- .beads/issues.jsonl | 2 +- docs/long-term-memory-arch.md | 5 ++++- docs/memory-implementation-status.md | 20 ++++++++++++++------ klbr-bench/src/longmemeval.rs | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++++-- klbr-core/src/evidence.rs | 26 ++++++++++++++++++++------ klbr-core/src/memory.rs | 2 +- klbr-core/src/pipeline.rs | 405 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------- 7 file(s) changed, 490 insertion(s)(+), 24 deletion(s)(-) diff --git a/.beads/issues.jsonl b/.beads/issues.jsonl --- a/.beads/issues.jsonl +++ b/.beads/issues.jsonl @@ -1,4 +1,4 @@ -{"_type":"issue","id":"klbr-u0q","title":"Implement session-first lexical-contained memory retrieval","description":"Build the next retrieval architecture from docs/long-term-memory-arch.md: stage-one session/event candidate generation, lexical containment, and language-agnostic operation planning experiments without adding english cue-word hacks.","design":"Keep EvidencePacket as the rank object. Use exact refs, dense episode/session cards, learned sparse or fts side channels, and graph expansion from strong seeds only. Keep QueryOp shape but replace cue-list policy with a structured classifier plus deterministic fallback.","acceptance_criteria":"A bench profile can retrieve from session/episode cards first and use chunk fts/sparse hits as packet enrichment; lexical-only policy paths are removed or contained behind candidate channels; traces expose candidate session recall and packet/rendered evidence metrics; multilingual fixtures cover non-English cue-free retrieval cases.","notes":"2026-06-29 partial: removed hardcoded english cue-list planner, english lane routing, english negation/entity-boundary packet filters, and wh/pronoun neighbor-expansion heuristics. OpPlan now comes from structured model JSON when an llm endpoint is configured, otherwise conservative lookup. Core and bench tests pass.","status":"in_progress","priority":1,"issue_type":"feature","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-28T22:06:32Z","created_by":"dawn","updated_at":"2026-06-28T22:20:45Z","started_at":"2026-06-28T22:10:12Z","dependency_count":0,"dependent_count":0,"comment_count":0} +{"_type":"issue","id":"klbr-u0q","title":"Implement session-first lexical-contained memory retrieval","description":"Build the next retrieval architecture from docs/long-term-memory-arch.md: stage-one session/event candidate generation, lexical containment, and language-agnostic operation planning experiments without adding english cue-word hacks.","design":"Keep EvidencePacket as the rank object. Use exact refs, dense episode/session cards, learned sparse or fts side channels, and graph expansion from strong seeds only. Keep QueryOp shape but replace cue-list policy with a structured classifier plus deterministic fallback.","acceptance_criteria":"A bench profile can retrieve from session/episode cards first and use chunk fts/sparse hits as packet enrichment; lexical-only policy paths are removed or contained behind candidate channels; traces expose candidate session recall and packet/rendered evidence metrics; multilingual fixtures cover non-English cue-free retrieval cases.","notes":"2026-06-29 partial: removed hardcoded english cue-list planner, english lane routing, english negation/entity-boundary packet filters, and wh/pronoun neighbor-expansion heuristics. OpPlan now comes from structured model JSON when an llm endpoint is configured, otherwise conservative lookup. Core and bench tests pass.\n2026-06-29 partial: landed initial session-first retrieval profile for klbr-full/dense-only/session-first profiles. Retrieval now groups archival candidates by session_id, prefers episodic/event-card anchors, adds chunk fts/sparse/dense hits as enrichment, serializes stage_one.session_candidates, and reports CandidateSessionRecall metrics in klbr-bench. fts-only now writes episode notes but skips embedded episode-memory rows so lexical ablations do not require the dense embedder at ingest. Verified core/bench tests, continuous-loop, and a one-row session-first/fts-only retrieval-only trace smoke.","status":"in_progress","priority":1,"issue_type":"feature","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-28T22:06:32Z","created_by":"dawn","updated_at":"2026-06-29T11:37:00Z","started_at":"2026-06-28T22:10:12Z","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-jwn","title":"Organize evolving memory architecture docs","description":"Make the memory architecture docs distinguish current truth, implementation status, proposals, research inputs, and archived rationale so future agents do not treat older reports as canonical.","acceptance_criteria":"Docs have a clear index and status taxonomy; current memory architecture direction points at docs/long-term-memory-arch.md; older research reports are marked as historical or supporting; AGENTS.md routes future agents through the doc index.","status":"closed","priority":1,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-28T22:01:20Z","created_by":"dawn","updated_at":"2026-06-28T22:07:09Z","started_at":"2026-06-28T22:01:23Z","closed_at":"2026-06-28T22:07:09Z","close_reason":"Completed docs index, current architecture review rewrite, status banners, AGENTS routing, and follow-up implementation issue klbr-u0q.","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-h4l","title":"Cache bench embeddings locally and run 25-sample QA bench","description":"Make LongMemEval bench embeddings use a project-local gitignored sqlite cache so reruns do not pay for the same embeddings twice. Then run a 25-sample stratified QA bench, inspect failures, try bounded non-overfit tweaks if evidence supports them, and prepare a report packet if results remain weak or tuning would be overfit.","design":"Prefer a global project cache path such as benchmarks/cache/*.db over temp/run-local caches. Keep benchmark polling sparse: start the run, wait for artifacts or process completion, then inspect once.","acceptance_criteria":"Embedding cache lives under the project and is ignored by git; cache hits are reused across bench runs; a 25-sample bench artifact exists with failure analysis; any code changes are tested, committed, and pushed.","notes":"Implemented project-local embedding cache at benchmarks/cache/embeddings.db and wired LongMemEval bench LlmClient construction through it. Ran stratified 25-sample QA on seed 4937553249516211047: current-code run benchmarks/runs/longmemeval-s/klbr-full/2026-06-28_210327.693783Z scored official accuracy 0.6800. Report packet written to report_packet.md in that run dir. High-budget targeted reruns fixed only dd2973ad and a1cc6108; remaining failures point to fact/timeline synthesis rather than safe cue tuning.","status":"closed","priority":1,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-28T20:08:32Z","created_by":"dawn","updated_at":"2026-06-28T21:21:27Z","started_at":"2026-06-28T20:08:35Z","closed_at":"2026-06-28T21:21:27Z","close_reason":"Implemented project-local dense/sparse embedding cache, ran 25-sample QA bench, inspected failures, tried bounded planner/high-budget experiments, and wrote report packet.","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-f0q.1","title":"Add operation-aware query planning","description":"Introduce a small OpPlan/QueryOp classifier for memory QA queries: lookup, aggregate_count, aggregate_sum, aggregate_avg, order_or_rank, update_resolution, preference_recommendation, and abstain_or_false_premise_check. Use it to switch only high-value classes away from plain lookup behavior.","design":"Keep the label set intentionally small and rule-based first; avoid growing English lexical hacks into the main ranking decision boundary.","acceptance_criteria":"Production memory retrieval traces expose the selected query op and plan; existing lookup behavior remains unchanged by default; aggregate/order/update/preference queries can be detected deterministically in tests.","status":"closed","priority":1,"issue_type":"feature","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-28T19:44:46Z","created_by":"dawn","updated_at":"2026-06-28T19:54:15Z","started_at":"2026-06-28T19:45:04Z","closed_at":"2026-06-28T19:54:15Z","close_reason":"Implemented OpPlan/QueryOp classification, serialized op_plan into retrieval traces, preserved lookup defaults, added deterministic planner and pipeline trace tests.","dependencies":[{"issue_id":"klbr-f0q.1","depends_on_id":"klbr-f0q","type":"parent-child","created_at":"2026-06-28T22:44:46Z","created_by":"dawn","metadata":"{}"}],"dependency_count":0,"dependent_count":1,"comment_count":0} diff --git a/docs/long-term-memory-arch.md b/docs/long-term-memory-arch.md --- a/docs/long-term-memory-arch.md +++ b/docs/long-term-memory-arch.md @@ -155,7 +155,10 @@ ## implementation order -1. add a session/event-first profile. +1. add a session/event-first profile. initial implementation landed: selected + profiles now group archival candidates by session, prefer episode/event-card + anchors, expose `stage_one.session_candidates`, and use chunk fts/sparse hits + as enrichment. dense retrieval over episode cards should seed sessions before chunk hits decide the evidence horizon. chunk fts/sparse hits should enrich those session packets. diff --git a/docs/memory-implementation-status.md b/docs/memory-implementation-status.md --- a/docs/memory-implementation-status.md +++ b/docs/memory-implementation-status.md @@ -54,6 +54,13 @@ - query-time english negation/entity-boundary packet filters were removed. lexical infrastructure remains candidate generation through sparse embeddings or sqlite fts, not handwritten query policy. +- session/event-first retrieval has an initial production profile: `klbr-full`, + `dense-only`, and explicit `session-first` profiles group archival candidates + by `session_id`, prefer episodic/event-card anchors when present, then add + chunk fts/sparse/dense hits as packet enrichment. retrieval traces now include + `stage_one.session_candidates`. +- `fts-only` keeps episode notes for fts but skips embedded episode-memory rows, + so lexical ablations no longer require the dense embedder during ingest. - `klbr-bench -- run` now uses the production pipeline instead of manually calling low-level retrieval: - longmemeval-s/m style data - flexible locomo adapter with fixture coverage and protocol-note artifacts @@ -65,11 +72,10 @@ ## open direction -- `klbr-u0q`: implement session/event-first retrieval and finish lexical - containment beyond the first cue-list removal. +- `klbr-u0q`: continue session/event-first retrieval beyond the initial profile: + improve fact-group coverage, operation-aware assembly, and multilingual + regression coverage. - harden the structured planner path and add model-backed planner fixtures. -- add a session/event-first retrieval profile where dense episode cards seed - sessions and chunk hits enrich evidence packets. - finish collect-mode fact-group sufficiency and gap diagnosis for aggregate, order, and update-resolution questions. - render fact rows/timeline rows for operation-aware synthesis. @@ -88,8 +94,10 @@ ```text cargo check -p klbr-bench: 0 errors, 1 pre-existing warning -klbr-core tests: 105 passed, 1 ignored -continuous-loop smoke: passed +klbr-core tests: 141 passed, 1 ignored +klbr-bench tests: 17 passed +continuous-loop smoke: passed with the pre-existing agent.rs dead-code warning +session-first/fts-only trace smoke: evaluated 1, CandidateSessionRecallAny@5 1.0000 ``` retrieval benchmark smoke: diff --git a/klbr-bench/src/longmemeval.rs b/klbr-bench/src/longmemeval.rs --- a/klbr-bench/src/longmemeval.rs +++ b/klbr-bench/src/longmemeval.rs @@ -51,6 +51,8 @@ #[derive(Debug, Clone, Serialize)] struct PacketSelectionMetrics { packet_k: usize, + candidate_session_recall_any_at_k: f64, + candidate_session_recall_all_at_k: f64, answer_bearing_ref_selected: bool, answer_bearing_ref_rendered: bool, answer_bearing_ref_in_context: bool, @@ -1395,6 +1397,8 @@ let mut answer_bearing_ref_rendered = 0usize; let mut answer_bearing_ref_in_context = 0usize; let mut answer_value_visible = 0usize; + let mut candidate_session_recall_any_at_k = 0usize; + let mut candidate_session_recall_all_at_k = 0usize; let mut gold_session_in_context = 0usize; let mut gold_session_plus_neighbor_in_context = 0usize; let mut packet_recall_any_at_k = 0usize; @@ -1548,6 +1552,12 @@ if packet_metrics.answer_value_visible { answer_value_visible += 1; } + if packet_metrics.candidate_session_recall_any_at_k > 0.0 { + candidate_session_recall_any_at_k += 1; + } + if packet_metrics.candidate_session_recall_all_at_k > 0.0 { + candidate_session_recall_all_at_k += 1; + } if packet_metrics.gold_session_in_context { gold_session_in_context += 1; } @@ -1644,6 +1654,8 @@ "answer_bearing_ref_rendered": if answerable == 0 { 0.0 } else { answer_bearing_ref_rendered as f64 / answerable_denominator }, "answer_bearing_ref_in_context": if answerable == 0 { 0.0 } else { answer_bearing_ref_in_context as f64 / answerable_denominator }, "answer_value_visible": if answerable == 0 { 0.0 } else { answer_value_visible as f64 / answerable_denominator }, + "candidate_session_recall_any_at_k": if answerable == 0 { 0.0 } else { candidate_session_recall_any_at_k as f64 / answerable_denominator }, + "candidate_session_recall_all_at_k": if answerable == 0 { 0.0 } else { candidate_session_recall_all_at_k as f64 / answerable_denominator }, "gold_session_in_context": if answerable == 0 { 0.0 } else { gold_session_in_context as f64 / answerable_denominator }, "gold_session_plus_neighbor_in_context": if answerable == 0 { 0.0 } else { gold_session_plus_neighbor_in_context as f64 / answerable_denominator }, "packet_recall_any_at_k": if answerable == 0 { 0.0 } else { packet_recall_any_at_k as f64 / answerable_denominator }, @@ -1663,7 +1675,7 @@ fs::write( out_dir.join("report.md"), format!( - "# klbr memory pipeline benchmark\n\n- suite: `{}`\n- profile: `{}`\n- evaluated: `{}`\n- answerable: `{}`\n- sample: `{}`\n- RecallAny@5: `{:.4}`\n- RecallAll@5: `{:.4}`\n- PacketRecallAny@{}: `{:.4}`\n- PacketRecallAll@{}: `{:.4}`\n- answer_bearing_ref_selected: `{:.4}`\n- answer_bearing_ref_rendered: `{:.4}`\n- answer_bearing_ref_in_context: `{:.4}`\n- answer_value_visible: `{:.4}`\n- gold_session_plus_neighbor_in_context: `{:.4}`\n- same_session_dedupe_suppression_count: `{}`\n- packet_tokens_mean: `{:.2}`\n- op_plan_counts: `{}`\n- retrieval_only: `{}`\n- official_eval: `{}`\n", + "# klbr memory pipeline benchmark\n\n- suite: `{}`\n- profile: `{}`\n- evaluated: `{}`\n- answerable: `{}`\n- sample: `{}`\n- RecallAny@5: `{:.4}`\n- RecallAll@5: `{:.4}`\n- CandidateSessionRecallAny@{}: `{:.4}`\n- CandidateSessionRecallAll@{}: `{:.4}`\n- PacketRecallAny@{}: `{:.4}`\n- PacketRecallAll@{}: `{:.4}`\n- answer_bearing_ref_selected: `{:.4}`\n- answer_bearing_ref_rendered: `{:.4}`\n- answer_bearing_ref_in_context: `{:.4}`\n- answer_value_visible: `{:.4}`\n- gold_session_plus_neighbor_in_context: `{:.4}`\n- same_session_dedupe_suppression_count: `{}`\n- packet_tokens_mean: `{:.2}`\n- op_plan_counts: `{}`\n- retrieval_only: `{}`\n- official_eval: `{}`\n", suite, profile, evaluated, @@ -1671,6 +1683,10 @@ sample_summary, if answerable == 0 { 0.0 } else { recall_any_at_5 as f64 / answerable as f64 }, if answerable == 0 { 0.0 } else { recall_all_at_5 as f64 / answerable as f64 }, + top_k, + if answerable == 0 { 0.0 } else { candidate_session_recall_any_at_k as f64 / answerable_denominator }, + top_k, + if answerable == 0 { 0.0 } else { candidate_session_recall_all_at_k as f64 / answerable_denominator }, top_k, if answerable == 0 { 0.0 } else { packet_recall_any_at_k as f64 / answerable_denominator }, top_k, @@ -2078,6 +2094,13 @@ .take(k) .filter_map(|packet| packet.session_id.clone()) .collect::>(); + let top_k_candidate_sessions = retrieval + .stage_one + .session_candidates + .iter() + .take(k) + .filter_map(|candidate| candidate.session_id.clone()) + .collect::>(); let selected_refs = retrieval .packets .iter() @@ -2126,6 +2149,26 @@ } else { 0.0 }; + let candidate_session_recall_any_at_k = if gold_sessions.is_empty() { + 0.0 + } else if gold_sessions + .iter() + .any(|session_id| top_k_candidate_sessions.contains(session_id)) + { + 1.0 + } else { + 0.0 + }; + let candidate_session_recall_all_at_k = if gold_sessions.is_empty() { + 0.0 + } else if gold_sessions + .iter() + .all(|session_id| top_k_candidate_sessions.contains(session_id)) + { + 1.0 + } else { + 0.0 + }; let packet_tokens_total = included_packets .iter() .map(|packet| packet.estimated_tokens) @@ -2138,6 +2181,8 @@ PacketSelectionMetrics { packet_k: k, + candidate_session_recall_any_at_k, + candidate_session_recall_all_at_k, answer_bearing_ref_selected, answer_bearing_ref_rendered, answer_bearing_ref_in_context, @@ -2149,7 +2194,12 @@ same_session_dedupe_suppression_count: retrieval .packet_omissions .iter() - .filter(|omission| omission.reason == "same_session_cap") + .filter(|omission| { + matches!( + omission.reason.as_str(), + "same_session_cap" | "same_session_overlap" + ) + }) .count(), packet_tokens_total, packet_tokens_mean, diff --git a/klbr-core/src/evidence.rs b/klbr-core/src/evidence.rs --- a/klbr-core/src/evidence.rs +++ b/klbr-core/src/evidence.rs @@ -11,6 +11,7 @@ pub enum RetrievalSource { Exact, Fts, + Sparse, Dense, Graph, TurnWindow, @@ -26,6 +27,7 @@ match self { Self::Exact => "exact", Self::Fts => "fts", + Self::Sparse => "sparse", Self::Dense => "dense", Self::Graph => "graph", Self::TurnWindow => "turn_window", @@ -41,6 +43,7 @@ match s { "exact" => Self::Exact, "fts" => Self::Fts, + "sparse" | "bge_m3_sparse" => Self::Sparse, "dense" => Self::Dense, "graph" => Self::Graph, "turn_window" => Self::TurnWindow, @@ -592,7 +595,7 @@ .signals .sources .iter() - .any(|source| matches!(source.as_str(), "exact" | "fts" | "dense")); + .any(|source| matches!(source.as_str(), "exact" | "fts" | "sparse" | "dense")); has_graph_signal && !has_corroborating_signal } @@ -641,6 +644,13 @@ .filter(|sig| sig.source == "fts") .map(|sig| sig.rank) .min(); + let sparse_rank = packet + .signals + .per_source + .iter() + .filter(|sig| sig.source == "sparse") + .map(|sig| sig.rank) + .min(); let dense_rank = packet .signals .per_source @@ -663,8 +673,11 @@ if let Some(r) = dense_rank { rrf_score += 1.2 / (rrf_k + r as f32); } - if let Some(r) = fts_rank { + if let Some(r) = sparse_rank { rrf_score += 1.0 / (rrf_k + r as f32); + } + if let Some(r) = fts_rank { + rrf_score += 0.8 / (rrf_k + r as f32); } if let Some(r) = graph_rank { rrf_score += 0.6 / (rrf_k + r as f32); @@ -702,10 +715,11 @@ fn source_priority(source: &str) -> usize { match source { "exact" => 0, - "fts" => 1, - "dense" => 2, - "graph" => 3, - _ => 4, + "dense" => 1, + "sparse" => 2, + "fts" => 3, + "graph" => 4, + _ => 5, } } diff --git a/klbr-core/src/memory.rs b/klbr-core/src/memory.rs --- a/klbr-core/src/memory.rs +++ b/klbr-core/src/memory.rs @@ -1808,7 +1808,7 @@ body: row.get(4)?, token_count: row.get::<_, i64>(5)? as usize, score: score as f32, - source: "bge_m3_sparse".to_string(), + source: "sparse".to_string(), }) })?; diff --git a/klbr-core/src/pipeline.rs b/klbr-core/src/pipeline.rs --- a/klbr-core/src/pipeline.rs +++ b/klbr-core/src/pipeline.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::time::Duration; use anyhow::Result; @@ -111,6 +111,8 @@ pub route: LaneRoute, #[serde(default)] pub op_plan: OpPlan, + #[serde(default)] + pub stage_one: StageOneTrace, pub exact_refs: Vec, pub candidates: Vec, #[serde(default)] @@ -119,6 +121,23 @@ pub packet_omissions: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub packet_rerank: Option, +} + +#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] +pub struct StageOneTrace { + pub strategy: String, + pub session_candidates: Vec, +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct StageOneSessionCandidate { + pub group_key: String, + pub session_id: Option, + pub anchor_ref: String, + pub anchor_entity_type: EntityType, + pub sources: Vec, + pub seed_refs: Vec, + pub score: f32, } #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -165,6 +184,7 @@ pub lexical: bool, pub dense: bool, pub graph: bool, + pub session_first: bool, } impl PipelineProfile { @@ -175,13 +195,19 @@ let dense_only = lower.contains("dense-only"); let no_graph = lower.contains("no-graph") || lower.ends_with("/semantic") || lower == "semantic"; + let session_first = !raw_turns + && (lower.contains("session-first") + || lower.contains("klbr-full") + || lower.contains("dense-only") + || lower.contains("hybrid-packet")); Self { name: name.to_string(), - write_episode_memories: !raw_turns, + write_episode_memories: !raw_turns && !fts_only, write_episode_notes: !raw_turns, lexical: !dense_only, dense: !raw_turns && !fts_only, graph: !no_graph && !raw_turns, + session_first, } } } @@ -388,10 +414,14 @@ candidates.extend(self.expand_graph(&seed_refs, candidate_limit)?); } - let candidates = rank_seed_candidates( - candidates, - packet_limit.saturating_mul(4).max(candidate_limit), - ); + let ranked_limit = packet_limit.saturating_mul(4).max(candidate_limit); + let (candidates, stage_one) = if self.profile.session_first { + rank_session_first_candidates(candidates, ranked_limit) + } else { + let candidates = rank_seed_candidates(candidates, ranked_limit); + let stage_one = trace_from_ranked_candidates("source_interleave", &candidates); + (candidates, stage_one) + }; let packet_plan = self.evidence_planner().build_packets_with_collect_mode( &query.text, &candidates, @@ -407,6 +437,7 @@ routed_lanes, route, op_plan, + stage_one, exact_refs, candidates, packets, @@ -425,7 +456,7 @@ let mut remaining_tokens = budget.max_tokens; let mut packets = String::new(); packets.push_str("\n"); - let mut fallback_packets; + let fallback_packets; let evidence_packets = if retrieved.packets.is_empty() && !retrieved.candidates.is_empty() && retrieved.packet_omissions.is_empty() @@ -515,6 +546,10 @@ routed_lanes: lanes, route, op_plan, + stage_one: StageOneTrace { + strategy: "complete_stored".to_string(), + session_candidates: vec![], + }, exact_refs: vec![], candidates, packets, @@ -810,6 +845,7 @@ let before = out.len(); for source in [ RetrievalSource::Fts, + RetrievalSource::Sparse, RetrievalSource::Dense, RetrievalSource::Graph, ] { @@ -824,6 +860,7 @@ } for source in [ RetrievalSource::Fts, + RetrievalSource::Sparse, RetrievalSource::Dense, RetrievalSource::Graph, ] { @@ -837,6 +874,245 @@ candidate.rank_in_source = *count; } out +} + +fn rank_session_first_candidates( + mut candidates: Vec, + limit: usize, +) -> (Vec, StageOneTrace) { + assign_source_ranks(&mut candidates); + + let mut exact = Vec::new(); + let mut groups = HashMap::::new(); + for candidate in candidates { + if candidate.source == RetrievalSource::Exact { + exact.push(candidate); + continue; + } + let group_key = candidate_group_key(&candidate); + groups + .entry(group_key.clone()) + .or_insert_with(|| CandidateSessionGroup::new(group_key, candidate.session_id.clone())) + .push(candidate); + } + + let mut groups = groups.into_values().collect::>(); + groups.sort_by(|left, right| { + right + .score() + .partial_cmp(&left.score()) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| left.group_key.cmp(&right.group_key)) + }); + + let mut out = Vec::new(); + let mut seen_refs = HashSet::new(); + for candidate in exact { + if seen_refs.insert(candidate.ref_id.clone()) { + out.push(candidate); + } + if out.len() >= limit { + assign_source_ranks(&mut out); + let stage_one = trace_from_ranked_candidates("session_first", &out); + return (out, stage_one); + } + } + + let mut session_candidates = + trace_from_ranked_candidates("session_first", &out).session_candidates; + let mut traced_groups = session_candidates + .iter() + .map(|candidate| candidate.group_key.clone()) + .collect::>(); + for group in groups { + if out.len() >= limit { + break; + } + let trace = group.trace(); + if traced_groups.insert(trace.group_key.clone()) { + session_candidates.push(trace); + } + let group_take = if group.session_id.is_some() { 3 } else { 1 }; + for candidate in group.ordered_candidates().into_iter().take(group_take) { + if seen_refs.insert(candidate.ref_id.clone()) { + out.push(candidate); + } + if out.len() >= limit { + break; + } + } + } + assign_source_ranks(&mut out); + ( + out, + StageOneTrace { + strategy: "session_first".to_string(), + session_candidates, + }, + ) +} + +#[derive(Debug)] +struct CandidateSessionGroup { + group_key: String, + session_id: Option, + candidates: Vec, +} + +impl CandidateSessionGroup { + fn new(group_key: String, session_id: Option) -> Self { + Self { + group_key, + session_id, + candidates: Vec::new(), + } + } + + fn push(&mut self, candidate: EvidenceAtom) { + self.candidates.push(candidate); + } + + fn score(&self) -> f32 { + let source_score = [ + RetrievalSource::Dense, + RetrievalSource::Sparse, + RetrievalSource::Fts, + RetrievalSource::Graph, + RetrievalSource::Other, + ] + .into_iter() + .filter_map(|source| { + self.candidates + .iter() + .filter(|candidate| candidate.source == source) + .map(candidate_rank_score) + .max_by(|left, right| left.partial_cmp(right).unwrap_or(std::cmp::Ordering::Equal)) + }) + .sum::(); + let session_anchor_bonus = if self.candidates.iter().any(is_session_card_candidate) { + 0.35 + } else { + 0.0 + }; + let source_count_bonus = 0.15 * (1.0 + self.source_names().len() as f32).ln(); + source_score + session_anchor_bonus + source_count_bonus + } + + fn trace(&self) -> StageOneSessionCandidate { + let anchor = self + .candidates + .iter() + .min_by(|left, right| candidate_order(left).cmp(&candidate_order(right))) + .expect("session candidate group should not be empty"); + StageOneSessionCandidate { + group_key: self.group_key.clone(), + session_id: self.session_id.clone(), + anchor_ref: anchor.ref_id.clone(), + anchor_entity_type: anchor.entity_type, + sources: self.source_names(), + seed_refs: self + .candidates + .iter() + .map(|candidate| candidate.ref_id.clone()) + .collect(), + score: self.score(), + } + } + + fn ordered_candidates(mut self) -> Vec { + self.candidates.sort_by(|left, right| { + candidate_order(left) + .cmp(&candidate_order(right)) + .then_with(|| left.ref_id.cmp(&right.ref_id)) + }); + self.candidates + } + + fn source_names(&self) -> Vec { + let mut names = Vec::new(); + for candidate in &self.candidates { + let source = candidate.source.as_str().to_string(); + if !names.contains(&source) { + names.push(source); + } + } + names + } +} + +fn assign_source_ranks(candidates: &mut [EvidenceAtom]) { + let mut counts = HashMap::::new(); + for candidate in candidates { + let count = counts.entry(candidate.source).or_insert(0); + *count += 1; + candidate.rank_in_source = *count; + } +} + +fn candidate_group_key(candidate: &EvidenceAtom) -> String { + candidate + .session_id + .as_ref() + .map(|session_id| format!("session:{session_id}")) + .unwrap_or_else(|| format!("ref:{}", candidate.ref_id)) +} + +fn is_session_card_candidate(candidate: &EvidenceAtom) -> bool { + candidate.lane == MemoryLane::Episodic && candidate.entity_type != EntityType::TurnChunk +} + +fn candidate_rank_score(candidate: &EvidenceAtom) -> f32 { + source_weight(candidate.source) / (20.0 + candidate.rank_in_source.max(1) as f32) +} + +fn source_weight(source: RetrievalSource) -> f32 { + match source { + RetrievalSource::Exact => 3.0, + RetrievalSource::Dense => 1.4, + RetrievalSource::Sparse => 1.0, + RetrievalSource::Fts => 0.8, + RetrievalSource::Graph => 0.35, + _ => 0.2, + } +} + +fn candidate_order(candidate: &EvidenceAtom) -> (usize, usize, usize) { + let card_rank = usize::from(!is_session_card_candidate(candidate)); + let source_rank = match candidate.source { + RetrievalSource::Exact => 0, + RetrievalSource::Dense => 1, + RetrievalSource::Sparse => 2, + RetrievalSource::Fts => 3, + RetrievalSource::Graph => 4, + _ => 5, + }; + (card_rank, source_rank, candidate.rank_in_source) +} + +fn trace_from_ranked_candidates(strategy: &str, candidates: &[EvidenceAtom]) -> StageOneTrace { + let mut seen = HashSet::new(); + let session_candidates = candidates + .iter() + .filter_map(|candidate| { + let group_key = candidate_group_key(candidate); + if !seen.insert(group_key.clone()) { + return None; + } + Some(StageOneSessionCandidate { + group_key, + session_id: candidate.session_id.clone(), + anchor_ref: candidate.ref_id.clone(), + anchor_entity_type: candidate.entity_type, + sources: vec![candidate.source.as_str().to_string()], + seed_refs: vec![candidate.ref_id.clone()], + score: candidate_rank_score(candidate), + }) + }) + .collect(); + StageOneTrace { + strategy: strategy.to_string(), + session_candidates, + } } struct PacketRerankDecision { @@ -1293,6 +1569,98 @@ use crate::evidence::{EvidencePacketKind, EvidencePacketSignals, EvidenceSourceSignal}; use crate::planner::QueryOp; use tempfile::NamedTempFile; + + #[test] + fn sparse_candidates_are_ranked_as_retrieval_sources() { + let ranked = rank_seed_candidates( + vec![seed_atom( + "sparse_ref", + Some("s1"), + EntityType::TurnChunk, + RetrievalSource::Sparse, + 1, + MemoryLane::Episodic, + )], + 5, + ); + + assert_eq!(ranked.len(), 1); + assert_eq!(ranked[0].source, RetrievalSource::Sparse); + assert_eq!(ranked[0].rank_in_source, 1); + } + + #[test] + fn profile_selects_session_first_without_changing_raw_turn_ablations() { + assert!(PipelineProfile::named("klbr-full").session_first); + assert!(PipelineProfile::named("dense-only").session_first); + assert!(!PipelineProfile::named("raw-turns/fts-only").session_first); + assert!(!PipelineProfile::named("fts-only").session_first); + assert!(!PipelineProfile::named("fts-only").write_episode_memories); + assert!(PipelineProfile::named("fts-only").write_episode_notes); + } + + #[test] + fn session_first_promotes_episode_anchor_and_keeps_chunk_enrichment() { + let (ranked, trace) = rank_session_first_candidates( + vec![ + seed_atom( + "s1_chunk_sparse", + Some("s1"), + EntityType::TurnChunk, + RetrievalSource::Sparse, + 1, + MemoryLane::Episodic, + ), + seed_atom( + "s2_episode", + Some("s2"), + EntityType::Episode, + RetrievalSource::Dense, + 1, + MemoryLane::Episodic, + ), + seed_atom( + "s1_episode", + Some("s1"), + EntityType::Episode, + RetrievalSource::Dense, + 2, + MemoryLane::Episodic, + ), + seed_atom( + "s1_chunk_fts", + Some("s1"), + EntityType::TurnChunk, + RetrievalSource::Fts, + 1, + MemoryLane::Episodic, + ), + ], + 8, + ); + + assert_eq!(trace.strategy, "session_first"); + assert_eq!( + trace + .session_candidates + .first() + .and_then(|candidate| { candidate.session_id.as_ref().map(String::as_str) }), + Some("s1") + ); + assert_eq!(trace.session_candidates[0].anchor_ref, "s1_episode"); + assert!(trace.session_candidates[0] + .sources + .contains(&"sparse".to_string())); + assert!(trace.session_candidates[0] + .sources + .contains(&"dense".to_string())); + + assert_eq!(ranked[0].ref_id, "s1_episode"); + assert_eq!(ranked[1].ref_id, "s1_chunk_sparse"); + assert!(ranked + .iter() + .any(|candidate| candidate.ref_id == "s1_chunk_fts")); + } #[tokio::test] async fn turn_window_packet_includes_next_turn_answer() -> Result<()> { @@ -1839,6 +2207,29 @@ fn reader_prompt_treats_packet_bodies_as_local_evidence_windows() { assert!(MEMORY_READER_SYSTEM_PROMPT.contains("adjacent bodies inside the same packet")); assert!(MEMORY_READER_SYSTEM_PROMPT.contains("do not answer a discovery/source channel")); + } + + fn seed_atom( + ref_id: &str, + session_id: Option<&str>, + entity_type: EntityType, + source: RetrievalSource, + rank_in_source: usize, + lane: MemoryLane, + ) -> EvidenceAtom { + EvidenceAtom { + ref_id: ref_id.to_string(), + alias: None, + session_id: session_id.map(str::to_string), + entity_type, + lane, + source, + rank_in_source, + score: 0.0, + token_count: 8, + body: format!("body for {ref_id}"), + anchor_kind: AnchorKind::QueryMatch, + } } fn rerank_test_packet( -- tangled.sh