use khive_brain_core::{BalancedRecallState, BetaPosterior, FeedbackEventKind, FeedbackSignal};
use uuid::Uuid;
const FAST_US: i64 = 50_000;
pub fn on_recall_hit(state: &mut BalancedRecallState, target_id: Uuid, latency_us: i64) {
state.total_events += 1;
state.relevance.update_success();
if latency_us <= FAST_US {
state.temporal.update_success();
} else {
state.temporal.update_failure();
}
let posterior = state
.entity_posteriors
.get_or_insert(target_id, || BetaPosterior::new(1.0, 1.0));
posterior.update_success();
}
pub fn on_recall_miss(state: &mut BalancedRecallState) {
state.total_events += 1;
state.relevance.update_failure();
state.temporal.update_failure();
}
pub fn on_explicit_feedback(state: &mut BalancedRecallState, target_id: Uuid, signal: &str) {
if let Some(event_kind) = FeedbackEventKind::from_signal_str(signal) {
let w = event_kind.update_weight();
if event_kind.is_positive() {
state.salience.update_success_weighted(w);
} else {
state.salience.update_failure_weighted(w);
}
if event_kind == FeedbackEventKind::Correction {
state.relevance.update_failure_weighted(w);
}
let posterior = state
.entity_posteriors
.get_or_insert(target_id, || BetaPosterior::new(1.0, 1.0));
if event_kind.is_positive() {
posterior.update_success_weighted(w);
} else {
posterior.update_failure_weighted(w);
}
state.total_events += 1;
} else if let Ok(fb) =
serde_json::from_value::<FeedbackSignal>(serde_json::Value::String(signal.to_owned()))
{
let positive = matches!(fb, FeedbackSignal::Useful);
if positive {
state.salience.update_success();
} else {
state.salience.update_failure();
}
let posterior = state
.entity_posteriors
.get_or_insert(target_id, || BetaPosterior::new(1.0, 1.0));
if positive {
posterior.update_success();
} else {
posterior.update_failure();
}
state.total_events += 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
fn fresh() -> BalancedRecallState {
BalancedRecallState::new(100)
}
#[test]
fn recall_hit_fast_increments_relevance_temporal_and_entity_alpha() {
let mut s = fresh();
let id = Uuid::new_v4();
on_recall_hit(&mut s, id, 10_000);
assert_eq!(s.total_events, 1);
assert!((s.relevance.alpha() - 8.0).abs() < 1e-12);
assert!((s.relevance.beta() - 3.0).abs() < 1e-12);
assert!((s.temporal.alpha() - 2.0).abs() < 1e-12);
assert!((s.temporal.beta() - 9.0).abs() < 1e-12);
let ep = s.entity_posteriors.get(&id).unwrap();
assert!((ep.alpha() - 2.0).abs() < 1e-12);
}
#[test]
fn recall_hit_slow_increments_relevance_alpha_but_temporal_beta() {
let mut s = fresh();
let id = Uuid::new_v4();
on_recall_hit(&mut s, id, 100_000);
assert!((s.relevance.alpha() - 8.0).abs() < 1e-12);
assert!((s.temporal.beta() - 10.0).abs() < 1e-12);
assert!((s.temporal.alpha() - 1.0).abs() < 1e-12);
}
#[test]
fn recall_miss_increments_relevance_beta_and_temporal_beta() {
let mut s = fresh();
on_recall_miss(&mut s);
assert_eq!(s.total_events, 1);
assert!((s.relevance.beta() - 4.0).abs() < 1e-12); assert!((s.temporal.beta() - 10.0).abs() < 1e-12); assert!(s.entity_posteriors.is_empty());
}
#[test]
fn explicit_feedback_useful_increments_salience_alpha() {
let mut s = fresh();
let id = Uuid::new_v4();
on_explicit_feedback(&mut s, id, "useful");
assert_eq!(s.total_events, 1);
assert!((s.salience.alpha() - 3.0).abs() < 1e-12);
assert!((s.salience.beta() - 8.0).abs() < 1e-12);
let ep = s.entity_posteriors.get(&id).unwrap();
assert!((ep.alpha() - 2.0).abs() < 1e-12);
}
#[test]
fn explicit_feedback_not_useful_increments_salience_beta() {
let mut s = fresh();
let id = Uuid::new_v4();
on_explicit_feedback(&mut s, id, "not_useful");
assert!((s.salience.beta() - 9.0).abs() < 1e-12); let ep = s.entity_posteriors.get(&id).unwrap();
assert!((ep.beta() - 2.0).abs() < 1e-12);
}
#[test]
fn explicit_feedback_wrong_increments_salience_beta() {
let mut s = fresh();
let id = Uuid::new_v4();
on_explicit_feedback(&mut s, id, "wrong");
assert!((s.salience.beta() - 9.0).abs() < 1e-12); }
#[test]
fn explicit_feedback_explicit_positive_applies_weight_1_5() {
let mut s = fresh();
let id = Uuid::new_v4();
on_explicit_feedback(&mut s, id, "explicit_positive");
assert!((s.salience.alpha() - 3.5).abs() < 1e-12); assert!((s.salience.beta() - 8.0).abs() < 1e-12);
let ep = s.entity_posteriors.get(&id).unwrap();
assert!((ep.alpha() - 2.5).abs() < 1e-12); }
#[test]
fn explicit_feedback_correction_updates_salience_beta_and_relevance_beta() {
let mut s = fresh();
let id = Uuid::new_v4();
on_explicit_feedback(&mut s, id, "correction");
assert!((s.salience.beta() - 10.0).abs() < 1e-12); assert!((s.relevance.beta() - 5.0).abs() < 1e-12); let ep = s.entity_posteriors.get(&id).unwrap();
assert!((ep.beta() - 3.0).abs() < 1e-12); }
#[test]
fn explicit_feedback_unknown_signal_is_noop() {
let mut s = fresh();
let id = Uuid::new_v4();
let sal_before = (s.salience.alpha(), s.salience.beta());
on_explicit_feedback(&mut s, id, "bad_value");
assert_eq!(s.total_events, 0);
assert!((s.salience.alpha() - sal_before.0).abs() < 1e-12);
assert!((s.salience.beta() - sal_before.1).abs() < 1e-12);
assert!(s.entity_posteriors.is_empty());
}
}