Skip to main content

semantic_memory/
reinstatement.rs

1//! Contextual reinstatement building blocks.
2//!
3//! Reinstatement scoring measures how well a candidate memory chunk (episode,
4//! entity, namespace) matches an active retrieval context. The score is purely
5//! additive and bounded in [0, 1] — callers use it as a soft boost signal, not
6//! a hard gate.
7
8/// Active context that drives reinstatement-biased retrieval.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ReinstatementContext {
11    /// Episode IDs that are currently active / in-focus.
12    pub episode_ids: Vec<String>,
13    /// Entity IDs that are currently in scope.
14    pub entity_ids: Vec<String>,
15    /// Optional namespace constraint.
16    pub namespace: Option<String>,
17}
18
19/// Per-candidate reinstatement score with per-dimension match flags.
20#[derive(Debug, Clone, PartialEq)]
21pub struct ReinstatementScore {
22    /// Aggregate score in [0, 1].
23    pub score: f32,
24    /// True when at least one episode in the context matched the candidate.
25    pub matched_episode: bool,
26    /// True when at least one entity in the context matched the candidate.
27    pub matched_entity: bool,
28    /// True when the namespace matched.
29    pub matched_namespace: bool,
30}
31
32const EPISODE_WEIGHT: f32 = 0.4;
33const ENTITY_WEIGHT: f32 = 0.4;
34const NAMESPACE_WEIGHT: f32 = 0.2;
35
36/// Compute a reinstatement score for one candidate against the active context.
37///
38/// - Episode match contributes 0.4
39/// - Entity match contributes 0.4
40/// - Namespace match contributes 0.2
41///
42/// The returned `score` is always in [0.0, 1.0].
43pub fn compute_reinstatement_score(
44    ctx: &ReinstatementContext,
45    candidate_episode_ids: &[String],
46    candidate_entity_ids: &[String],
47    candidate_namespace: Option<&str>,
48) -> ReinstatementScore {
49    let matched_episode = !ctx.episode_ids.is_empty()
50        && candidate_episode_ids
51            .iter()
52            .any(|id| ctx.episode_ids.contains(id));
53
54    let matched_entity = !ctx.entity_ids.is_empty()
55        && candidate_entity_ids
56            .iter()
57            .any(|id| ctx.entity_ids.contains(id));
58
59    let matched_namespace = match (ctx.namespace.as_deref(), candidate_namespace) {
60        (Some(a), Some(b)) => a == b,
61        _ => false,
62    };
63
64    let raw = if matched_episode { EPISODE_WEIGHT } else { 0.0 }
65        + if matched_entity { ENTITY_WEIGHT } else { 0.0 }
66        + if matched_namespace {
67            NAMESPACE_WEIGHT
68        } else {
69            0.0
70        };
71
72    ReinstatementScore {
73        score: raw.clamp(0.0, 1.0),
74        matched_episode,
75        matched_entity,
76        matched_namespace,
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    fn s(v: &str) -> String {
85        v.to_string()
86    }
87
88    #[test]
89    fn reinstatement_no_match_is_zero() {
90        let ctx = ReinstatementContext {
91            episode_ids: vec![s("ep-1")],
92            entity_ids: vec![s("ent-1")],
93            namespace: Some(s("ns-a")),
94        };
95        let result = compute_reinstatement_score(&ctx, &[], &[], None);
96        assert_eq!(result.score, 0.0);
97        assert!(!result.matched_episode);
98        assert!(!result.matched_entity);
99        assert!(!result.matched_namespace);
100    }
101
102    #[test]
103    fn reinstatement_full_match_is_one() {
104        let ctx = ReinstatementContext {
105            episode_ids: vec![s("ep-1")],
106            entity_ids: vec![s("ent-1")],
107            namespace: Some(s("ns-a")),
108        };
109        let result = compute_reinstatement_score(&ctx, &[s("ep-1")], &[s("ent-1")], Some("ns-a"));
110        assert!((result.score - 1.0).abs() < f32::EPSILON);
111        assert!(result.matched_episode);
112        assert!(result.matched_entity);
113        assert!(result.matched_namespace);
114    }
115
116    #[test]
117    fn reinstatement_episode_only() {
118        let ctx = ReinstatementContext {
119            episode_ids: vec![s("ep-2")],
120            entity_ids: vec![],
121            namespace: None,
122        };
123        let result = compute_reinstatement_score(&ctx, &[s("ep-2")], &[], None);
124        assert!((result.score - 0.4).abs() < 1e-6);
125        assert!(result.matched_episode);
126        assert!(!result.matched_entity);
127        assert!(!result.matched_namespace);
128    }
129
130    #[test]
131    fn reinstatement_entity_only() {
132        let ctx = ReinstatementContext {
133            episode_ids: vec![],
134            entity_ids: vec![s("ent-x")],
135            namespace: None,
136        };
137        let result = compute_reinstatement_score(&ctx, &[], &[s("ent-x")], None);
138        assert!((result.score - 0.4).abs() < 1e-6);
139        assert!(!result.matched_episode);
140        assert!(result.matched_entity);
141    }
142
143    #[test]
144    fn reinstatement_namespace_only() {
145        let ctx = ReinstatementContext {
146            episode_ids: vec![],
147            entity_ids: vec![],
148            namespace: Some(s("ns-b")),
149        };
150        let result = compute_reinstatement_score(&ctx, &[], &[], Some("ns-b"));
151        assert!((result.score - 0.2).abs() < 1e-6);
152        assert!(result.matched_namespace);
153    }
154
155    #[test]
156    fn reinstatement_score_bounded() {
157        let ctx = ReinstatementContext {
158            episode_ids: vec![s("ep-1"), s("ep-2")],
159            entity_ids: vec![s("ent-1")],
160            namespace: Some(s("ns")),
161        };
162        let result =
163            compute_reinstatement_score(&ctx, &[s("ep-1"), s("ep-2")], &[s("ent-1")], Some("ns"));
164        assert!(result.score >= 0.0 && result.score <= 1.0);
165    }
166}