car-memgine 0.48.0

//! Utility-aware memory retrieval — learned utility posteriors + explore/exploit
//! selection (U-Mem, arXiv 2602.22406).
//!
//! Applies *Towards Autonomous Memory Agents* to CAR — see
//! `docs/proposals/autonomous-memory-agents.md`. U-Mem's retrieval picks
//! memories by **Semantic-Aware Thompson Sampling**: it blends semantic
//! similarity with a *learned utility distribution* per memory, so proven
//! memories are exploited, unproven ones are explored, and cold-start bias is
//! mitigated. The utility distribution is a Beta posterior updated by outcomes —
//! exactly what CAR's skill `success_count`/`fail_count` already approximate
//! (see [`crate::engine::DomainStats`]); this generalizes that into a principled
//! posterior and a retrieval-time explore/exploit score.
//!
//! Two selection modes:
//! - [`ucb_score`] — a **deterministic** upper-confidence blend (mean utility +
//!   exploration·uncertainty). Cold-start memories have high uncertainty → a
//!   bonus → they get explored, *without* randomness. This fits CAR's
//!   deterministic-runtime thesis: reproducible retrieval, no RNG in the hot path.
//! - [`thompson_score`] — the paper's stochastic draw, with the Beta sampler
//!   **injected** (like other CAR randomness/inference seams) so the core stays
//!   pure and tests stay deterministic.
//!
//! Pure: no inference/embedding deps. The caller supplies semantic `relevance`
//! (e.g. from the graph's spreading activation) and the per-memory counts.

use serde::{Deserialize, Serialize};

/// A Beta(`alpha`, `beta`) utility posterior for a memory/skill, learned from
/// outcomes. `alpha-1` is the success count and `beta-1` the failure count under
/// a Laplace (uniform) prior, so a brand-new memory is `Beta(1,1)` — maximally
/// uncertain, which is what drives cold-start exploration.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct UtilityPosterior {
    pub alpha: f64,
    pub beta: f64,
}

impl Default for UtilityPosterior {
    fn default() -> Self {
        Self::uniform()
    }
}

impl UtilityPosterior {
    /// The cold-start prior `Beta(1,1)` — uniform, maximally uncertain.
    pub fn uniform() -> Self {
        Self {
            alpha: 1.0,
            beta: 1.0,
        }
    }

    /// Posterior from observed outcome counts under a Laplace prior:
    /// `Beta(success+1, fail+1)`. Matches `DomainStats { success_count,
    /// fail_count }`, so existing skill stats become utility posteriors directly.
    pub fn from_counts(success: u64, fail: u64) -> Self {
        Self {
            alpha: success as f64 + 1.0,
            beta: fail as f64 + 1.0,
        }
    }

    pub fn record_success(&mut self) {
        self.alpha += 1.0;
    }

    pub fn record_failure(&mut self) {
        self.beta += 1.0;
    }

    /// Posterior mean `alpha / (alpha+beta)` — the expected utility.
    pub fn mean(&self) -> f64 {
        self.alpha / (self.alpha + self.beta)
    }

    /// Posterior standard deviation — how unsure we are. High for cold-start,
    /// shrinking as evidence accumulates. This is the exploration signal.
    pub fn uncertainty(&self) -> f64 {
        let s = self.alpha + self.beta;
        // Var(Beta) = ab / ((a+b)^2 (a+b+1))
        ((self.alpha * self.beta) / (s * s * (s + 1.0))).sqrt()
    }

    /// Deterministic upper-confidence utility in `[0,1]`:
    /// `clamp(mean + exploration·uncertainty)`. `exploration = 0` is pure
    /// exploitation (mean); higher values weight exploring uncertain memories.
    pub fn ucb(&self, exploration: f64) -> f64 {
        (self.mean() + exploration * self.uncertainty()).clamp(0.0, 1.0)
    }

    /// A Thompson draw of utility, with the Beta sampler injected as
    /// `draw(alpha, beta) -> f64` (kept out of this pure crate). The result is
    /// clamped to `[0,1]`.
    pub fn thompson(&self, draw: impl FnOnce(f64, f64) -> f64) -> f64 {
        draw(self.alpha, self.beta).clamp(0.0, 1.0)
    }
}

/// A retrieval candidate: an id, its semantic `relevance` in `[0,1]` (e.g. from
/// spreading activation), and its learned utility posterior.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Candidate {
    pub id: String,
    pub relevance: f64,
    #[serde(default)]
    pub posterior: UtilityPosterior,
}

/// A scored, ranked candidate.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScoredCandidate {
    pub id: String,
    pub score: f64,
    pub relevance: f64,
    /// The utility value used in the blend (UCB or Thompson draw).
    pub utility: f64,
}

/// Blend semantic `relevance` with a utility value by `utility_weight` in
/// `[0,1]`: `relevance·(1-w) + utility·w`. `w = 0` reproduces pure relevance
/// ranking (so enabling utility never changes behavior unless opted in); `w = 1`
/// is pure utility.
fn blend(relevance: f64, utility: f64, utility_weight: f64) -> f64 {
    let w = utility_weight.clamp(0.0, 1.0);
    relevance * (1.0 - w) + utility * w
}

/// Deterministic utility-aware ranking (the runtime default): blend each
/// candidate's relevance with its UCB utility, sort descending. Reproducible —
/// no RNG. `exploration` controls the cold-start/uncertainty bonus;
/// `utility_weight` controls how much utility matters vs. raw relevance.
pub fn rank_ucb(
    candidates: &[Candidate],
    exploration: f64,
    utility_weight: f64,
) -> Vec<ScoredCandidate> {
    let mut scored: Vec<ScoredCandidate> = candidates
        .iter()
        .map(|c| {
            let utility = c.posterior.ucb(exploration);
            ScoredCandidate {
                id: c.id.clone(),
                score: blend(c.relevance, utility, utility_weight),
                relevance: c.relevance,
                utility,
            }
        })
        .collect();
    sort_desc(&mut scored);
    scored
}

/// Stochastic (Thompson) utility-aware ranking — the paper's SA-CTS — with the
/// Beta sampler injected as `draw(alpha, beta) -> f64`. Tests/callers supply a
/// deterministic or seeded draw; the core stays pure.
pub fn rank_thompson(
    candidates: &[Candidate],
    utility_weight: f64,
    mut draw: impl FnMut(f64, f64) -> f64,
) -> Vec<ScoredCandidate> {
    let mut scored: Vec<ScoredCandidate> = candidates
        .iter()
        .map(|c| {
            let utility = c.posterior.thompson(&mut draw);
            ScoredCandidate {
                id: c.id.clone(),
                score: blend(c.relevance, utility, utility_weight),
                relevance: c.relevance,
                utility,
            }
        })
        .collect();
    sort_desc(&mut scored);
    scored
}

fn sort_desc(scored: &mut [ScoredCandidate]) {
    // Score descending; ties broken by id for a deterministic order.
    scored.sort_by(|a, b| {
        b.score
            .partial_cmp(&a.score)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(a.id.cmp(&b.id))
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn uniform_prior_is_uncertain_with_mean_half() {
        let p = UtilityPosterior::uniform();
        assert!((p.mean() - 0.5).abs() < 1e-9);
        assert!(p.uncertainty() > 0.2, "cold start should be very uncertain");
    }

    #[test]
    fn evidence_shrinks_uncertainty_and_moves_mean() {
        let proven = UtilityPosterior::from_counts(20, 0);
        let cold = UtilityPosterior::uniform();
        assert!(proven.mean() > 0.9);
        assert!(
            proven.uncertainty() < cold.uncertainty(),
            "more evidence = less uncertainty"
        );
    }

    #[test]
    fn record_outcomes_updates_posterior() {
        let mut p = UtilityPosterior::uniform();
        p.record_success();
        p.record_success();
        p.record_failure();
        // Beta(3,2)
        assert!((p.mean() - 0.6).abs() < 1e-9);
    }

    #[test]
    fn ucb_gives_coldstart_an_exploration_bonus() {
        // Same mean (0.5), but the uncertain one gets a higher UCB.
        let cold = UtilityPosterior::uniform(); // Beta(1,1), mean .5, high var
        let seen = UtilityPosterior::from_counts(10, 10); // Beta(11,11), mean .5, low var
        assert!((cold.mean() - seen.mean()).abs() < 1e-9);
        assert!(
            cold.ucb(1.0) > seen.ucb(1.0),
            "cold start is explored under UCB"
        );
        // With no exploration, both reduce to the mean.
        assert!((cold.ucb(0.0) - seen.ucb(0.0)).abs() < 1e-9);
    }

    fn cand(id: &str, relevance: f64, success: u64, fail: u64) -> Candidate {
        Candidate {
            id: id.to_string(),
            relevance,
            posterior: UtilityPosterior::from_counts(success, fail),
        }
    }

    #[test]
    fn weight_zero_is_pure_relevance_order() {
        // b has higher relevance but worse utility; at weight 0 b still wins.
        let cands = vec![cand("a", 0.4, 50, 0), cand("b", 0.9, 0, 50)];
        let ranked = rank_ucb(&cands, 0.0, 0.0);
        assert_eq!(ranked[0].id, "b");
    }

    #[test]
    fn utility_can_outrank_relevance_when_weighted() {
        // a: lower relevance, proven; b: higher relevance, proven-bad.
        let cands = vec![cand("a", 0.5, 50, 0), cand("b", 0.7, 0, 50)];
        let ranked = rank_ucb(&cands, 0.0, 0.8); // heavy utility weight
        assert_eq!(
            ranked[0].id, "a",
            "proven-useful memory outranks more-relevant useless one"
        );
    }

    #[test]
    fn thompson_uses_injected_draw_deterministically() {
        // Inject a draw that returns the posterior mean -> deterministic test.
        let cands = vec![cand("a", 0.5, 9, 0), cand("b", 0.5, 0, 9)];
        let ranked = rank_thompson(&cands, 1.0, |a, b| a / (a + b)); // = mean
                                                                     // a (mean 10/11) beats b (mean 1/11) under pure-utility weight.
        assert_eq!(ranked[0].id, "a");
    }
}