Skip to main content

kimetsu_brain/
tune.rs

1//! v1.5 / S2: Self-Tuning Brain sweep engine.
2//!
3//! Pure functions for objective scoring, holdout splitting, and history I/O.
4//! The sweep itself is driven from the CLI (kimetsu-cli) which calls
5//! `evaluate_combo` in-process with injected embedder + optional reranker.
6//!
7//! Sweep space (config-addressable only):
8//!   - min_lexical_coverage ∈ {0.3, 0.4, 0.5, 0.6}
9//!   - min_semantic_score   ∈ {-1.0(auto), 0.0, 0.25, 0.35, 0.45}
10//!   - reranker id ∈ {off, ms-marco-tinybert-l-2-v2, jina-reranker-v1-tiny-en,
11//!     ms-marco-minilm-l-4-v2}
12//!
13//! NOT swept (compile-time or complex-deploy):
14//!   - RERANK_POOL (compile-time const in the daemon) — deferred.
15//!
16//! Objective (S2.3):
17//!   quality - cost_weight * mean_final_serialized_UTF8_bound
18//! Default cost_weight = 0.05 / 6000; historical regret is diagnostic only.
19//!
20//! S2.1 Re-tune triggers:
21//!   - Corpus milestone: ≥50 memories added since last tune.
22//!   - Drift: insights hit-rate decline OR regret-rate rise beyond thresholds.
23//!
24//! S2.2 Model re-selection advisor:
25//!   Recommends re-running the embedder×reranker grid at corpus milestones.
26//!   Reports download+reindex cost. Never auto-switches.
27
28use serde::{Deserialize, Serialize};
29
30// ─── S2.1: Re-tune trigger constants ─────────────────────────────────────────
31
32/// Corpus milestone: propose a re-tune when ≥ this many memories have been
33/// added since the last tune run.
34pub const RETUNE_CORPUS_MILESTONE: u64 = 50;
35
36/// Drift threshold: propose a re-tune when the regret rate (regrets / served
37/// events in the last 24h window) rises above this fraction.
38pub const RETUNE_REGRET_RATE_THRESHOLD: f64 = 0.10;
39
40/// S2.2: approximate token cost to reindex 1 000 memories when switching
41/// the embedder model (conservative estimate based on batch embed overhead).
42/// Used to report the cost of a full embedder switch in the advisor output.
43pub const REINDEX_TOKENS_PER_1K_MEMORIES: u64 = 2_000;
44
45/// Explicit policy, not a fitted optimum: a full delivery budget costs 0.05 quality units.
46pub const DEFAULT_COST_LAMBDA: f64 = 0.05;
47pub const DEFAULT_COST_WEIGHT: f64 = DEFAULT_COST_LAMBDA / 6000.0;
48
49pub const LEXICAL_FLOORS: &[f32] = &[0.3, 0.4, 0.5, 0.6];
50pub const SEMANTIC_FLOORS: &[f32] = &[-1.0, 0.0, 0.25, 0.35, 0.45];
51pub const RERANKER_IDS: &[&str] = &[
52    "off",
53    "ms-marco-tinybert-l-2-v2",
54    "jina-reranker-v1-tiny-en",
55    "ms-marco-minilm-l-4-v2",
56];
57
58/// v2.6: how the lexical and semantic rankings are merged. See
59/// [`crate::fusion`] for why this is a real ranking decision and not plumbing.
60///
61/// It is swept rather than defaulted because BM25 and cosine live on different
62/// scales, and which merge rule wins depends on the corpus — the whole reason
63/// the semantic floor already had to be calibrated per embedder family. The
64/// shipped default stays `linear` until a corpus says otherwise; this is how
65/// you get it to say so.
66pub const FUSION_MODES: &[&str] = &["linear", "rrf"];
67
68// ─── Combo ────────────────────────────────────────────────────────────────────
69
70#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
71pub struct TuneCombo {
72    pub min_lexical_coverage: f32,
73    pub min_semantic_score: f32,
74    pub reranker_id: String,
75    /// v2.6: `"linear"` or `"rrf"`. `#[serde(default)]` keeps tune-history
76    /// files written before v2.6 deserializing cleanly — they predate the
77    /// dimension, so they describe linear runs.
78    #[serde(default = "default_fusion_mode")]
79    pub fusion: String,
80}
81
82fn default_fusion_mode() -> String {
83    "linear".to_string()
84}
85
86impl TuneCombo {
87    pub fn all_combos() -> Vec<TuneCombo> {
88        let mut out = Vec::new();
89        for &lex in LEXICAL_FLOORS {
90            for &sem in SEMANTIC_FLOORS {
91                for &rr in RERANKER_IDS {
92                    for &fusion in FUSION_MODES {
93                        out.push(TuneCombo {
94                            min_lexical_coverage: lex,
95                            min_semantic_score: sem,
96                            reranker_id: rr.to_string(),
97                            fusion: fusion.to_string(),
98                        });
99                    }
100                }
101            }
102        }
103        out
104    }
105}
106
107// ─── Per-combo result ─────────────────────────────────────────────────────────
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct ComboResult {
111    pub combo: TuneCombo,
112    pub mean_mrr: f64,
113    pub mean_tokens: f64,
114    /// quality minus per-unit cost weight times final serialized UTF-8 bound
115    pub objective: f64,
116}
117
118// ─── Tune history entry ───────────────────────────────────────────────────────
119
120#[derive(Debug, Clone, Serialize, Deserialize)]
121pub struct TuneHistoryEntry {
122    pub timestamp: String,
123    pub before: TuneCombo,
124    pub after: TuneCombo,
125    pub train_objective: f64,
126    pub holdout_objective: f64,
127    pub holdout_mrr: f64,
128    pub baseline_holdout_objective: f64,
129    /// S2.1: corpus size (active memory count) at the time of this tune run.
130    /// `None` for history entries written before S2 (backward compat).
131    #[serde(default, skip_serializing_if = "Option::is_none")]
132    pub memory_count_at_tune: Option<u64>,
133    /// None marks historical objectives whose units/policy were not recorded.
134    #[serde(default, skip_serializing_if = "Option::is_none")]
135    pub measurement: Option<serde_json::Value>,
136}
137
138// ─── S2.1: Re-tune trigger state ─────────────────────────────────────────────
139
140/// Trigger state for S2.1 re-tune proposals.  Computed cheaply (no sweep).
141#[derive(Debug, Clone, Serialize, Deserialize)]
142pub struct RetuneTriggerState {
143    /// Active memory count right now.
144    pub current_memory_count: u64,
145    /// Active memory count at the last tune, or 0 if never tuned.
146    pub memory_count_at_last_tune: u64,
147    /// Memories added since the last tune.
148    pub memories_added_since_tune: u64,
149    /// Whether the corpus milestone threshold has been crossed.
150    pub corpus_milestone_triggered: bool,
151    /// Regrets in the last 24 h window.
152    pub recent_regret_count: u64,
153    /// Context-served events in the last 24 h window.
154    pub recent_served_count: u64,
155    /// Regret rate = recent_regret_count / recent_served_count (0.0 when served=0).
156    pub regret_rate: f64,
157    /// Whether the drift threshold has been crossed.
158    pub drift_triggered: bool,
159    /// True when either trigger is active.
160    pub should_retune: bool,
161    /// Timestamp of the last tune, or `None` if never tuned.
162    pub last_tuned_at: Option<String>,
163}
164
165/// Compute re-tune trigger state from the DB without running any sweep.
166///
167/// Reads:
168/// - Active memory count (current and at last tune from `tune-history.json`).
169/// - `retrieval.regret` event count in the last 24 h.
170/// - `context.served` event count in the last 24 h.
171pub fn compute_retune_trigger(
172    conn: &rusqlite::Connection,
173    kimetsu_dir: &std::path::Path,
174) -> kimetsu_core::KimetsuResult<RetuneTriggerState> {
175    // Current active memory count.
176    let current_memory_count: u64 = conn.query_row(
177        "SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL",
178        [],
179        |r| r.get(0),
180    )?;
181
182    // Last tune entry (if any).
183    let last_entry = latest_tune_history(kimetsu_dir)?;
184    let memory_count_at_last_tune = last_entry
185        .as_ref()
186        .and_then(|e| e.memory_count_at_tune)
187        .unwrap_or(0);
188    let last_tuned_at = last_entry.as_ref().map(|e| e.timestamp.clone());
189
190    let memories_added_since_tune = current_memory_count.saturating_sub(memory_count_at_last_tune);
191    let corpus_milestone_triggered = memories_added_since_tune >= RETUNE_CORPUS_MILESTONE;
192
193    // Regret / served counts in the last 24 h.
194    let cutoff_secs = std::time::SystemTime::now()
195        .duration_since(std::time::UNIX_EPOCH)
196        .map(|d| d.as_secs())
197        .unwrap_or(0)
198        .saturating_sub(86_400);
199    // Convert unix-secs cutoff to an approximate ISO string for the SQL comparison.
200    let cutoff_iso = {
201        let dt = time::OffsetDateTime::from_unix_timestamp(cutoff_secs as i64)
202            .unwrap_or(time::OffsetDateTime::UNIX_EPOCH);
203        dt.format(&time::format_description::well_known::Rfc3339)
204            .unwrap_or_default()
205    };
206
207    let recent_regret_count: u64 = conn.query_row(
208        "SELECT COUNT(*) FROM events WHERE kind = 'retrieval.regret' AND ts >= ?1",
209        rusqlite::params![cutoff_iso],
210        |r| r.get(0),
211    )?;
212
213    let recent_served_count: u64 = conn.query_row(
214        "SELECT COUNT(*) FROM events WHERE ts>=?1 AND kind=CASE WHEN EXISTS(SELECT 1 FROM events WHERE kind='context.injected' AND ts>=?1) THEN 'context.injected' ELSE 'context.served' END",
215        rusqlite::params![cutoff_iso],
216        |r| r.get(0),
217    )?;
218
219    let regret_rate = if recent_served_count > 0 {
220        recent_regret_count as f64 / recent_served_count as f64
221    } else {
222        0.0
223    };
224    let drift_triggered = regret_rate >= RETUNE_REGRET_RATE_THRESHOLD;
225    let should_retune = corpus_milestone_triggered || drift_triggered;
226
227    Ok(RetuneTriggerState {
228        current_memory_count,
229        memory_count_at_last_tune,
230        memories_added_since_tune,
231        corpus_milestone_triggered,
232        recent_regret_count,
233        recent_served_count,
234        regret_rate,
235        drift_triggered,
236        should_retune,
237        last_tuned_at,
238    })
239}
240
241// ─── S2.2: Model re-selection advisor ────────────────────────────────────────
242
243/// Known embedder models and their approximate on-disk download sizes (MiB).
244///
245/// These are estimates for the advisor report; actual sizes vary by format.
246pub const KNOWN_EMBEDDER_MODELS: &[(&str, &str, u32)] = &[
247    // (model_id, description, approx_download_mib)
248    (
249        "jina-embeddings-v2-base-code",
250        "Jina v2 Code (768d, default)",
251        280,
252    ),
253    ("bge-small-en-v1.5", "BGE-small (384d, lightweight)", 130),
254    ("nomic-embed-text-v1.5", "Nomic Embed v1.5 (768d)", 270),
255    ("all-minilm-l6-v2", "MiniLM L6 (384d, fast)", 90),
256];
257
258/// Model re-selection advisor recommendation.
259#[derive(Debug, Clone, Serialize, Deserialize)]
260pub struct ModelAdvisorReport {
261    /// Whether the advisor recommends re-running the grid now.
262    pub recommend_grid_run: bool,
263    /// Reason for the recommendation.
264    pub reason: String,
265    /// Currently active embedder model id.
266    pub current_embedder: String,
267    /// Approximate number of memories to re-embed if the model changes.
268    pub memories_to_reindex: u64,
269    /// Estimated token cost to reindex (conservative lower-bound).
270    pub estimated_reindex_tokens: u64,
271    /// Estimated approximate download size for all candidate models (MiB).
272    pub candidate_models: Vec<ModelCandidate>,
273}
274
275/// A candidate embedder model for the grid sweep.
276#[derive(Debug, Clone, Serialize, Deserialize)]
277pub struct ModelCandidate {
278    pub model_id: String,
279    pub description: String,
280    pub approx_download_mib: u32,
281}
282
283/// Compute the model re-selection advisor report.
284///
285/// Does NOT run the sweep — only computes the metadata needed for the
286/// advisor recommendation.  The actual grid run is a separate `brain tune`
287/// invocation (reuses the existing sweep machinery).
288///
289/// `trigger` must be pre-computed via [`compute_retune_trigger`].
290pub fn compute_model_advisor(
291    current_embedder: &str,
292    trigger: &RetuneTriggerState,
293) -> ModelAdvisorReport {
294    let recommend_grid_run = trigger.corpus_milestone_triggered;
295    let reason = if trigger.corpus_milestone_triggered {
296        format!(
297            "Corpus grew by {} memories since last tune (≥{} threshold). \
298             Re-running the embedder×reranker grid is recommended to verify \
299             the current model remains optimal.",
300            trigger.memories_added_since_tune, RETUNE_CORPUS_MILESTONE,
301        )
302    } else {
303        format!(
304            "No corpus milestone triggered ({} memories added, threshold {}). \
305             Grid re-run is optional.",
306            trigger.memories_added_since_tune, RETUNE_CORPUS_MILESTONE,
307        )
308    };
309
310    let memories_to_reindex = trigger.current_memory_count;
311    let estimated_reindex_tokens =
312        (memories_to_reindex.max(1) / 1_000 + 1).saturating_mul(REINDEX_TOKENS_PER_1K_MEMORIES);
313
314    let candidate_models = KNOWN_EMBEDDER_MODELS
315        .iter()
316        .map(|(id, desc, mib)| ModelCandidate {
317            model_id: id.to_string(),
318            description: desc.to_string(),
319            approx_download_mib: *mib,
320        })
321        .collect();
322
323    ModelAdvisorReport {
324        recommend_grid_run,
325        reason,
326        current_embedder: current_embedder.to_string(),
327        memories_to_reindex,
328        estimated_reindex_tokens,
329        candidate_models,
330    }
331}
332
333// ─── Pure functions ───────────────────────────────────────────────────────────
334
335/// Compute the tuning objective for a combo result.
336///
337/// `objective = quality - cost_weight * mean_final_bound`; legacy callers retain
338/// their explicit per-unit coefficient. CLI uses DEFAULT_COST_WEIGHT.
339pub fn compute_objective(mean_mrr: f64, mean_tokens: f64, cost_weight: f64) -> f64 {
340    mean_mrr - cost_weight * mean_tokens
341}
342
343/// Compatibility wrapper. Historical regret is diagnostic only: it is not
344/// candidate-specific evidence and cannot affect candidate scores.
345pub fn compute_objective_with_regret(
346    quality: f64,
347    mean_bound: f64,
348    cost_weight: f64,
349    _regret_rate: f64,
350) -> f64 {
351    compute_objective(quality, mean_bound, cost_weight)
352}
353
354/// Count `retrieval.regret` events in `conn` within an optional ISO-8601
355/// timestamp window `[since, until]`.
356///
357/// Historical diagnostic and retune trigger only; never a candidate objective term.
358pub fn count_regret_events(
359    conn: &rusqlite::Connection,
360    since: Option<&str>,
361    until: Option<&str>,
362) -> kimetsu_core::KimetsuResult<u64> {
363    let count: u64 = match (since, until) {
364        (Some(lo), Some(hi)) => conn.query_row(
365            "SELECT COUNT(*) FROM events \
366             WHERE kind = 'retrieval.regret' AND ts >= ?1 AND ts <= ?2",
367            rusqlite::params![lo, hi],
368            |r| r.get(0),
369        )?,
370        (Some(lo), None) => conn.query_row(
371            "SELECT COUNT(*) FROM events \
372             WHERE kind = 'retrieval.regret' AND ts >= ?1",
373            rusqlite::params![lo],
374            |r| r.get(0),
375        )?,
376        (None, Some(hi)) => conn.query_row(
377            "SELECT COUNT(*) FROM events \
378             WHERE kind = 'retrieval.regret' AND ts <= ?1",
379            rusqlite::params![hi],
380            |r| r.get(0),
381        )?,
382        (None, None) => conn.query_row(
383            "SELECT COUNT(*) FROM events WHERE kind = 'retrieval.regret'",
384            [],
385            |r| r.get(0),
386        )?,
387    };
388    Ok(count)
389}
390
391/// Split cases into (train, holdout) using a deterministic seed derived from
392/// `case_count`. 80 % train, 20 % holdout. Indices into `cases` are returned.
393///
394/// Legacy index-only API. Evaluation uses grouped_train_holdout_split, which
395/// accepts case identities and prevents family leakage.
396pub fn train_holdout_split(case_count: usize) -> (Vec<usize>, Vec<usize>) {
397    if case_count < 2 {
398        return ((0..case_count).collect(), Vec::new());
399    }
400    let holdout_size = (case_count / 5).max(1); // ≥1 holdout
401    // Deterministic: pick every 5th index as holdout.
402    let holdout: Vec<usize> = (0..case_count).filter(|i| i % 5 == 0).collect();
403    let train: Vec<usize> = (0..case_count).filter(|i| i % 5 != 0).collect();
404    let _ = holdout_size; // used via filter logic above
405    (train, holdout)
406}
407
408/// Connected components join aliases sharing task family, normalized query,
409/// or any relevant/stale ID. Sorting component identities makes the 80/20 split
410/// deterministic under input permutations. One component has no holdout.
411#[derive(Debug)]
412pub struct FamilySplit {
413    pub train: Vec<usize>,
414    pub holdout: Vec<usize>,
415    pub family_count: usize,
416}
417pub fn grouped_train_holdout_split(cases: &[crate::eval::EvalCase]) -> FamilySplit {
418    use std::collections::{BTreeMap, BTreeSet};
419    let mut parents: Vec<usize> = (0..cases.len()).collect();
420    fn root(parents: &[usize], mut i: usize) -> usize {
421        while parents[i] != i {
422            i = parents[i]
423        }
424        i
425    }
426    let mut owners = BTreeMap::<String, usize>::new();
427    let keys: Vec<BTreeSet<String>> = cases
428        .iter()
429        .map(|c| {
430            let mut keys = BTreeSet::new();
431            keys.insert(format!(
432                "q:{}",
433                c.query
434                    .split_whitespace()
435                    .collect::<Vec<_>>()
436                    .join(" ")
437                    .to_lowercase()
438            ));
439            if !c.family.trim().is_empty() {
440                keys.insert(format!("f:{}", c.family.trim()));
441            }
442            for id in c.relevant.iter().chain(&c.stale) {
443                keys.insert(format!("m:{id}"));
444            }
445            keys
446        })
447        .collect();
448    for (i, case_keys) in keys.iter().enumerate() {
449        for key in case_keys {
450            if let Some(&other) = owners.get(key) {
451                let a = root(&parents, i);
452                let b = root(&parents, other);
453                parents[a] = b;
454            } else {
455                owners.insert(key.clone(), i);
456            }
457        }
458    }
459    let mut components = BTreeMap::<usize, (BTreeSet<String>, Vec<usize>)>::new();
460    for (i, case_keys) in keys.into_iter().enumerate() {
461        let entry = components.entry(root(&parents, i)).or_default();
462        entry
463            .0
464            .extend(case_keys.into_iter().filter(|k| !k.starts_with("m:")));
465        entry.1.push(i);
466    }
467    let mut groups: Vec<_> = components.into_values().collect();
468    groups.sort_by(|a, b| a.0.cmp(&b.0));
469    let count = groups.len();
470    let mut split = FamilySplit {
471        train: Vec::new(),
472        holdout: Vec::new(),
473        family_count: count,
474    };
475    for (i, (_, ids)) in groups.into_iter().enumerate() {
476        if count > 1 && i % 5 == 0 {
477            split.holdout.extend(ids)
478        } else {
479            split.train.extend(ids)
480        }
481    }
482    split.train.sort_unstable();
483    split.holdout.sort_unstable();
484    split
485}
486
487/// Select the best combo from a slice of `ComboResult` by objective score.
488/// Returns `None` when the slice is empty.
489pub fn select_winner(results: &[ComboResult]) -> Option<&ComboResult> {
490    results.iter().max_by(|a, b| {
491        a.objective
492            .partial_cmp(&b.objective)
493            .unwrap_or(std::cmp::Ordering::Equal)
494    })
495}
496
497// ─── Tune history I/O ─────────────────────────────────────────────────────────
498
499/// Append a `TuneHistoryEntry` to `.kimetsu/tune-history.json`.
500pub fn append_tune_history(
501    kimetsu_dir: &std::path::Path,
502    entry: TuneHistoryEntry,
503) -> kimetsu_core::KimetsuResult<()> {
504    let path = kimetsu_dir.join("tune-history.json");
505    let mut entries: Vec<TuneHistoryEntry> = if path.exists() {
506        let text = std::fs::read_to_string(&path)?;
507        serde_json::from_str(&text).unwrap_or_default()
508    } else {
509        Vec::new()
510    };
511    entries.push(entry);
512    let json = serde_json::to_string_pretty(&entries)?;
513    std::fs::write(&path, json)?;
514    Ok(())
515}
516
517/// Read the latest entry from `.kimetsu/tune-history.json`, if any.
518pub fn latest_tune_history(
519    kimetsu_dir: &std::path::Path,
520) -> kimetsu_core::KimetsuResult<Option<TuneHistoryEntry>> {
521    let path = kimetsu_dir.join("tune-history.json");
522    if !path.exists() {
523        return Ok(None);
524    }
525    let text = std::fs::read_to_string(&path)?;
526    let entries: Vec<TuneHistoryEntry> = serde_json::from_str(&text).unwrap_or_default();
527    Ok(entries.into_iter().last())
528}
529
530// ─── Unit tests ───────────────────────────────────────────────────────────────
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535    use ulid::Ulid;
536
537    /// The sweep space is the product of every swept dimension. Asserting the
538    /// arithmetic (rather than a literal) keeps this honest when a dimension is
539    /// added: v2.6 added `fusion`, which doubled the grid from 80 to 160.
540    #[test]
541    fn all_combos_covers_the_full_grid() {
542        let combos = TuneCombo::all_combos();
543        let expected =
544            LEXICAL_FLOORS.len() * SEMANTIC_FLOORS.len() * RERANKER_IDS.len() * FUSION_MODES.len();
545        assert_eq!(
546            combos.len(),
547            expected,
548            "expected {}×{}×{}×{}={expected} combos, got {}",
549            LEXICAL_FLOORS.len(),
550            SEMANTIC_FLOORS.len(),
551            RERANKER_IDS.len(),
552            FUSION_MODES.len(),
553            combos.len()
554        );
555
556        // Every combo must be distinct: a duplicated point wastes a full
557        // evaluation pass over the corpus.
558        let mut keys: Vec<String> = combos
559            .iter()
560            .map(|c| {
561                format!(
562                    "{}|{}|{}|{}",
563                    c.min_lexical_coverage, c.min_semantic_score, c.reranker_id, c.fusion
564                )
565            })
566            .collect();
567        keys.sort();
568        let before = keys.len();
569        keys.dedup();
570        assert_eq!(before, keys.len(), "sweep grid contains duplicate combos");
571    }
572
573    #[test]
574    fn compute_objective_formula() {
575        let obj = compute_objective(0.75, 1000.0, 0.005);
576        // 0.75 - 0.005 * 1000 = 0.75 - 5.0 = -4.25
577        assert!((obj - (-4.25)).abs() < 1e-9, "objective: {obj}");
578    }
579
580    #[test]
581    fn compute_objective_zero_cost_weight_is_just_mrr() {
582        let obj = compute_objective(0.85, 500.0, 0.0);
583        assert!((obj - 0.85).abs() < 1e-9, "objective with 0 cost: {obj}");
584    }
585
586    #[test]
587    fn train_holdout_split_80_20() {
588        let (train, holdout) = train_holdout_split(10);
589        // Indices 0..10, every 5th (0,5) → holdout, rest → train.
590        assert_eq!(holdout, vec![0, 5]);
591        assert_eq!(train, vec![1, 2, 3, 4, 6, 7, 8, 9]);
592        assert_eq!(train.len() + holdout.len(), 10);
593    }
594
595    #[test]
596    fn train_holdout_split_empty() {
597        let (train, holdout) = train_holdout_split(0);
598        assert!(train.is_empty());
599        assert!(holdout.is_empty());
600    }
601
602    #[test]
603    fn single_family_cannot_supply_an_independent_holdout() {
604        let (train, holdout) = train_holdout_split(1);
605        assert_eq!(train, vec![0]);
606        assert!(holdout.is_empty());
607    }
608
609    #[test]
610    fn split_membership_is_independent_of_input_order() {
611        let queries = ["alpha", "bravo", "charlie", "delta", "echo"];
612        let mut reversed = queries;
613        reversed.reverse();
614        let held = |q: &[&str]| {
615            let cases: Vec<_> = q
616                .iter()
617                .map(|query| {
618                    serde_json::from_value(serde_json::json!({"query":query,"relevant":[]}))
619                        .unwrap()
620                })
621                .collect();
622            let indexes = grouped_train_holdout_split(&cases).holdout;
623            indexes
624                .into_iter()
625                .map(|i| q[i].to_string())
626                .collect::<std::collections::BTreeSet<_>>()
627        };
628        assert_eq!(held(&queries), held(&reversed));
629    }
630
631    #[test]
632    fn overlapping_aliases_and_task_families_never_leak_into_holdout() {
633        let cases: Vec<crate::eval::EvalCase> = serde_json::from_value(serde_json::json!([
634            {"query":"a","relevant":["one"]},
635            {"query":"b","relevant":["two"],"stale":["one"]},
636            {"query":"c","relevant":["two"],"family":"task"},
637            {"query":"d","relevant":[],"family":"task"},
638            {"query":"e","relevant":[]}
639        ]))
640        .unwrap();
641        let split = grouped_train_holdout_split(&cases);
642        assert_eq!(split.family_count, 2);
643        let in_holdout = split.holdout.contains(&0);
644        for i in 1..4 {
645            assert_eq!(split.holdout.contains(&i), in_holdout);
646        }
647        let one = grouped_train_holdout_split(&cases[..4]);
648        assert!(one.holdout.is_empty());
649        assert_eq!(one.train.len(), 4);
650    }
651
652    #[test]
653    fn explicit_default_cost_policy_uses_budget_fraction_units() {
654        let score = compute_objective(0.75, 512.0, DEFAULT_COST_WEIGHT);
655        assert!((score - 0.7457333333333333).abs() < 1e-12);
656        assert_eq!(compute_objective(1.0, 6000.0, DEFAULT_COST_WEIGHT), 0.95);
657    }
658
659    #[test]
660    fn select_winner_picks_highest_objective() {
661        let combos = vec![
662            ComboResult {
663                combo: TuneCombo {
664                    min_lexical_coverage: 0.3,
665                    min_semantic_score: 0.0,
666                    reranker_id: "off".to_string(),
667                    fusion: "linear".to_string(),
668                },
669                mean_mrr: 0.7,
670                mean_tokens: 100.0,
671                objective: 0.2,
672            },
673            ComboResult {
674                combo: TuneCombo {
675                    min_lexical_coverage: 0.4,
676                    min_semantic_score: 0.25,
677                    reranker_id: "off".to_string(),
678                    fusion: "linear".to_string(),
679                },
680                mean_mrr: 0.9,
681                mean_tokens: 80.0,
682                objective: 0.5,
683            },
684        ];
685        let winner = select_winner(&combos).expect("winner");
686        assert!((winner.objective - 0.5).abs() < 1e-9);
687    }
688
689    #[test]
690    fn tune_history_roundtrip() {
691        let tmp = std::env::temp_dir().join(format!("kimetsu-tune-hist-{}", Ulid::new()));
692        std::fs::create_dir_all(&tmp).unwrap();
693
694        let entry = TuneHistoryEntry {
695            timestamp: "2026-06-11T00:00:00Z".to_string(),
696            before: TuneCombo {
697                min_lexical_coverage: 0.5,
698                min_semantic_score: -1.0,
699                reranker_id: "off".to_string(),
700                fusion: "linear".to_string(),
701            },
702            after: TuneCombo {
703                min_lexical_coverage: 0.4,
704                min_semantic_score: 0.25,
705                reranker_id: "ms-marco-tinybert-l-2-v2".to_string(),
706                fusion: "linear".to_string(),
707            },
708            train_objective: 0.55,
709            holdout_objective: 0.50,
710            holdout_mrr: 0.70,
711            baseline_holdout_objective: 0.45,
712            memory_count_at_tune: None,
713            measurement: None,
714        };
715
716        append_tune_history(&tmp, entry.clone()).unwrap();
717        let latest = latest_tune_history(&tmp).unwrap().unwrap();
718        assert!((latest.holdout_objective - 0.50).abs() < 1e-9);
719        assert_eq!(latest.after.reranker_id, "ms-marco-tinybert-l-2-v2");
720
721        std::fs::remove_dir_all(&tmp).ok();
722    }
723
724    #[test]
725    fn tune_history_empty_when_no_file() {
726        let tmp = std::env::temp_dir().join(format!("kimetsu-tune-empty-{}", Ulid::new()));
727        std::fs::create_dir_all(&tmp).unwrap();
728        let latest = latest_tune_history(&tmp).unwrap();
729        assert!(latest.is_none(), "no history file → None");
730        std::fs::remove_dir_all(&tmp).ok();
731    }
732
733    // ─── S2.3: regret-penalised objective ────────────────────────────────────
734
735    #[test]
736    fn compute_objective_with_regret_zero_rate_matches_base() {
737        let base = compute_objective(0.75, 500.0, 0.005);
738        let with_regret = compute_objective_with_regret(0.75, 500.0, 0.005, 0.0);
739        assert!(
740            (base - with_regret).abs() < 1e-9,
741            "zero regret_rate must give same result as base objective"
742        );
743    }
744
745    #[test]
746    fn historical_regret_cannot_change_candidate_objective() {
747        let base = compute_objective(0.75, 500.0, 0.005);
748        let with_regret = compute_objective_with_regret(0.75, 500.0, 0.005, 0.10);
749        assert_eq!(with_regret, base);
750    }
751
752    #[test]
753    fn even_full_historical_regret_is_diagnostic_only() {
754        let base = compute_objective(0.8, 0.0, 0.0);
755        let with_full = compute_objective_with_regret(0.8, 0.0, 0.0, 1.0);
756        assert_eq!(with_full, base);
757    }
758
759    // ─── S2.1: RetuneTriggerState ─────────────────────────────────────────────
760
761    use crate::{
762        project::{init_project, load_project},
763        projector,
764        user_brain::with_user_brain_disabled,
765    };
766    use kimetsu_core::{event::Event, ids::RunId};
767
768    fn trigger_test_root(label: &str) -> std::path::PathBuf {
769        let root =
770            std::env::temp_dir().join(format!("kimetsu-tune-trigger-{label}-{}", Ulid::new()));
771        kimetsu_core::paths::git_init_boundary(&root);
772        root
773    }
774
775    #[test]
776    fn retune_trigger_no_history_no_events() {
777        with_user_brain_disabled(|| {
778            let root = trigger_test_root("empty");
779            std::fs::create_dir_all(&root).expect("mkdir");
780            init_project(&root, false).expect("init");
781            let paths = kimetsu_core::paths::ProjectPaths::discover(&root).expect("paths");
782            let (_, _, conn) = load_project(&root).expect("load");
783            let state = compute_retune_trigger(&conn, &paths.kimetsu_dir).expect("trigger");
784            assert_eq!(state.current_memory_count, 0);
785            assert_eq!(state.memories_added_since_tune, 0);
786            assert!(!state.corpus_milestone_triggered);
787            assert!(!state.drift_triggered);
788            assert!(!state.should_retune);
789            assert!(state.last_tuned_at.is_none());
790            std::fs::remove_dir_all(&root).ok();
791        });
792    }
793
794    #[test]
795    fn retune_trigger_corpus_milestone_when_enough_memories() {
796        with_user_brain_disabled(|| {
797            let root = trigger_test_root("milestone");
798            std::fs::create_dir_all(&root).expect("mkdir");
799            init_project(&root, false).expect("init");
800            let paths = kimetsu_core::paths::ProjectPaths::discover(&root).expect("paths");
801
802            // Seed a fake tune-history entry with memory_count_at_tune = 0.
803            let entry = TuneHistoryEntry {
804                timestamp: "2026-01-01T00:00:00Z".to_string(),
805                before: TuneCombo {
806                    min_lexical_coverage: 0.4,
807                    min_semantic_score: 0.0,
808                    reranker_id: "off".to_string(),
809                    fusion: "linear".to_string(),
810                },
811                after: TuneCombo {
812                    min_lexical_coverage: 0.4,
813                    min_semantic_score: 0.0,
814                    reranker_id: "off".to_string(),
815                    fusion: "linear".to_string(),
816                },
817                train_objective: 0.5,
818                holdout_objective: 0.5,
819                holdout_mrr: 0.7,
820                baseline_holdout_objective: 0.45,
821                memory_count_at_tune: Some(0),
822                measurement: None,
823            };
824            append_tune_history(&paths.kimetsu_dir, entry).expect("append");
825
826            // Add RETUNE_CORPUS_MILESTONE memories via the add_memory API.
827            for i in 0..RETUNE_CORPUS_MILESTONE {
828                crate::project::add_memory(
829                    &root,
830                    kimetsu_core::memory::MemoryScope::Project,
831                    kimetsu_core::memory::MemoryKind::Fact,
832                    &format!("milestone memory {i}"),
833                )
834                .expect("add memory");
835            }
836
837            let (_, _, conn) = load_project(&root).expect("load");
838            let state = compute_retune_trigger(&conn, &paths.kimetsu_dir).expect("trigger");
839            assert!(
840                state.corpus_milestone_triggered,
841                "milestone must trigger at ≥{RETUNE_CORPUS_MILESTONE} memories added"
842            );
843            assert!(state.should_retune);
844            std::fs::remove_dir_all(&root).ok();
845        });
846    }
847
848    #[test]
849    fn retune_trigger_drift_when_regret_rate_high() {
850        with_user_brain_disabled(|| {
851            let root = trigger_test_root("drift");
852            std::fs::create_dir_all(&root).expect("mkdir");
853            init_project(&root, false).expect("init");
854            let paths = kimetsu_core::paths::ProjectPaths::discover(&root).expect("paths");
855            let (_, _, conn) = load_project(&root).expect("load");
856
857            // Seed 1 served event + 1 regret event (rate = 100% >> threshold).
858            let run_id = RunId::new();
859            let served_ev = Event::new(
860                run_id,
861                "context.injected",
862                serde_json::json!({"query_hash":"abc","capsule_count":1,"skipped":false}),
863            );
864            projector::apply_events(&conn, &[served_ev]).expect("seed served");
865            let regret_ev = Event::new(
866                run_id,
867                "retrieval.regret",
868                serde_json::json!({"memory_id":"m1","dropped_at":0,"cited_at":1}),
869            );
870            projector::apply_events(&conn, &[regret_ev]).expect("seed regret");
871
872            let state = compute_retune_trigger(&conn, &paths.kimetsu_dir).expect("trigger");
873            assert!(
874                state.drift_triggered,
875                "regret_rate ({:.2}) must exceed threshold ({RETUNE_REGRET_RATE_THRESHOLD})",
876                state.regret_rate
877            );
878            assert!(state.should_retune);
879            std::fs::remove_dir_all(&root).ok();
880        });
881    }
882
883    // ─── S2.2: ModelAdvisorReport ─────────────────────────────────────────────
884
885    #[test]
886    fn model_advisor_recommends_at_milestone() {
887        let trigger = RetuneTriggerState {
888            current_memory_count: 100,
889            memory_count_at_last_tune: 10,
890            memories_added_since_tune: 90,
891            corpus_milestone_triggered: true,
892            recent_regret_count: 0,
893            recent_served_count: 20,
894            regret_rate: 0.0,
895            drift_triggered: false,
896            should_retune: true,
897            last_tuned_at: Some("2026-01-01T00:00:00Z".to_string()),
898        };
899        let report = compute_model_advisor("jina-embeddings-v2-base-code", &trigger);
900        assert!(report.recommend_grid_run, "must recommend at milestone");
901        assert!(report.estimated_reindex_tokens > 0, "cost must be stated");
902        assert!(!report.candidate_models.is_empty());
903    }
904
905    #[test]
906    fn model_advisor_no_recommendation_below_milestone() {
907        let trigger = RetuneTriggerState {
908            current_memory_count: 30,
909            memory_count_at_last_tune: 25,
910            memories_added_since_tune: 5,
911            corpus_milestone_triggered: false,
912            recent_regret_count: 0,
913            recent_served_count: 10,
914            regret_rate: 0.0,
915            drift_triggered: false,
916            should_retune: false,
917            last_tuned_at: None,
918        };
919        let report = compute_model_advisor("jina-embeddings-v2-base-code", &trigger);
920        assert!(
921            !report.recommend_grid_run,
922            "must NOT recommend below milestone"
923        );
924    }
925
926    // ─── S2.3: count_regret_events ────────────────────────────────────────────
927
928    #[test]
929    fn count_regret_events_zero_in_empty_db() {
930        with_user_brain_disabled(|| {
931            let root = trigger_test_root("regret-count");
932            std::fs::create_dir_all(&root).expect("mkdir");
933            init_project(&root, false).expect("init");
934            let (_, _, conn) = load_project(&root).expect("load");
935            let count = count_regret_events(&conn, None, None).expect("count");
936            assert_eq!(count, 0);
937            std::fs::remove_dir_all(&root).ok();
938        });
939    }
940
941    #[test]
942    fn tune_history_entry_memory_count_roundtrip() {
943        let tmp = std::env::temp_dir().join(format!("kimetsu-tune-memcount-{}", Ulid::new()));
944        std::fs::create_dir_all(&tmp).unwrap();
945
946        let entry = TuneHistoryEntry {
947            timestamp: "2026-06-11T00:00:00Z".to_string(),
948            before: TuneCombo {
949                min_lexical_coverage: 0.5,
950                min_semantic_score: -1.0,
951                reranker_id: "off".to_string(),
952                fusion: "linear".to_string(),
953            },
954            after: TuneCombo {
955                min_lexical_coverage: 0.4,
956                min_semantic_score: 0.25,
957                reranker_id: "off".to_string(),
958                fusion: "linear".to_string(),
959            },
960            train_objective: 0.55,
961            holdout_objective: 0.50,
962            holdout_mrr: 0.70,
963            baseline_holdout_objective: 0.45,
964            memory_count_at_tune: Some(123),
965            measurement: None,
966        };
967
968        append_tune_history(&tmp, entry).unwrap();
969        let latest = latest_tune_history(&tmp).unwrap().unwrap();
970        assert_eq!(
971            latest.memory_count_at_tune,
972            Some(123),
973            "memory_count_at_tune must round-trip"
974        );
975
976        std::fs::remove_dir_all(&tmp).ok();
977    }
978}