gemini_memory_rs/retrieval/
plan.rs

1//! The retrieval plan — a transient, query-oriented reading of the transcript.
2//!
3//! A plan answers "which parts of existing memory might matter to this turn?"
4//! It never answers "what should be stored?" — that is a separate extraction
5//! with a separate schema, because conflating the two produces a model that
6//! stores what it was asked to recall.
7
8use chrono::{DateTime, Utc};
9use schemars::JsonSchema;
10use serde::{Deserialize, Serialize};
11
12use crate::core::{CanonicalPredicate, MemoryKind, PlanId, TurnId, stable_hash};
13
14/// Caps enforced on every plan, whoever produced it (§12.2).
15///
16/// A model asked for search terms will happily return forty. These bounds are
17/// applied after extraction so an over-eager plan degrades to a good one rather
18/// than a slow one.
19pub mod limits {
20    /// Maximum entities.
21    pub const ENTITIES: usize = 5;
22    /// Maximum topics.
23    pub const TOPICS: usize = 8;
24    /// Maximum predicates.
25    pub const PREDICATES: usize = 5;
26    /// Maximum lexical queries.
27    pub const LEXICAL_QUERIES: usize = 3;
28    /// Maximum terms within a single lexical query.
29    pub const QUERY_TERMS: usize = 16;
30    /// Maximum memory scopes.
31    pub const SCOPES: usize = 5;
32}
33
34/// What the user is trying to do with memory this turn.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
36#[serde(rename_all = "snake_case")]
37pub enum RetrievalIntent {
38    /// No memory is needed — a generic factual or visual question.
39    #[default]
40    None,
41    /// The user asked what B knows.
42    ExplicitRecall,
43    /// The user wants a suggestion informed by their preferences.
44    PersonalRecommendation,
45    /// The user referred to something that happened before.
46    PriorEventReference,
47    /// The user referred to a person in their life.
48    RelationshipReference,
49    /// The user is comparing options against their preferences.
50    Comparison,
51    /// Memory may help but the intent is unclear.
52    Ambient,
53}
54
55impl RetrievalIntent {
56    /// Whether this intent justifies searching memory at all.
57    pub fn requires_memory(self) -> bool {
58        !matches!(self, Self::None)
59    }
60}
61
62/// An entity the plan wants memories about.
63#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
64pub struct RetrievalEntity {
65    /// The surface form as spoken ("my wife", "Rhea").
66    pub surface: String,
67    /// A canonical form when one is known.
68    #[serde(default)]
69    pub canonical: Option<String>,
70}
71
72impl RetrievalEntity {
73    /// An entity known only by how the user said it.
74    pub fn surface(surface: impl Into<String>) -> Self {
75        Self {
76            surface: surface.into(),
77            canonical: None,
78        }
79    }
80
81    /// An entity resolved to a canonical name.
82    pub fn resolved(surface: impl Into<String>, canonical: impl Into<String>) -> Self {
83        Self {
84            surface: surface.into(),
85            canonical: Some(canonical.into()),
86        }
87    }
88
89    /// Every form this entity should be matched by.
90    pub fn forms(&self) -> Vec<String> {
91        let mut forms = vec![self.surface.clone()];
92        if let Some(canonical) = &self.canonical {
93            forms.push(canonical.clone());
94        }
95        forms
96    }
97}
98
99/// A time window the plan is interested in.
100#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
101pub struct TemporalConstraint {
102    /// Earliest instant of interest.
103    #[serde(default)]
104    pub after: Option<DateTime<Utc>>,
105    /// Latest instant of interest.
106    #[serde(default)]
107    pub before: Option<DateTime<Utc>>,
108    /// A human label such as `last_week`, kept for explanation.
109    #[serde(default)]
110    pub label: Option<String>,
111}
112
113/// A transient plan for one turn's retrieval.
114#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
115pub struct RetrievalPlan {
116    /// Plan identity.
117    pub plan_id: PlanId,
118    /// The turn the plan was derived from.
119    pub turn_id: TurnId,
120    /// The generation the plan was derived at.
121    pub generation: u64,
122    /// Whether memory should be consulted at all.
123    pub requires_memory: bool,
124    /// Confidence in that judgement.
125    pub confidence: f32,
126    /// What the user appears to want.
127    pub intent: RetrievalIntent,
128    /// Entities of interest.
129    pub entities: Vec<RetrievalEntity>,
130    /// Topical terms.
131    pub topics: Vec<String>,
132    /// Canonical predicates of interest.
133    pub predicates: Vec<CanonicalPredicate>,
134    /// Independent lexical queries to run and fuse.
135    pub lexical_queries: Vec<String>,
136    /// Memory kinds worth searching — a hint used for explanation, never a
137    /// filter. See `kind_filter` for the restricting form.
138    pub scopes: Vec<MemoryKind>,
139    /// A hard kind restriction, set only when a caller explicitly asked for one.
140    #[serde(default)]
141    pub kind_filter: Vec<MemoryKind>,
142    /// Whose fact the caller believes this is — the subject surface form.
143    ///
144    /// A *hint*, never a restriction: records matching it are boosted in the
145    /// fusion, records that do not are still returned. That asymmetry is the
146    /// whole design, and it is measured. Over 93 questions, a correct
147    /// `about`+`attribute` pair applied softly takes top-5 from 79 to 91; a
148    /// wrong one costs one question. Applied as a hard filter, a correct pair
149    /// reaches the same 91 and a **wrong one returns nothing at all** — 0 of
150    /// 93, with the answer absent from the candidate set every single time.
151    ///
152    /// Break-even follows from that: filtering softly beats not filtering once
153    /// the caller is right 8% of the time, where filtering hard needs 87%. See
154    /// `tests/filter_dsl_probe.rs`.
155    #[serde(default)]
156    pub subject_hint: Option<String>,
157    /// Which attribute of the subject — the canonical predicate.
158    ///
159    /// Same soft semantics as [`subject_hint`](Self::subject_hint), and the
160    /// more valuable of the two: a model shown
161    /// [`memory_map`](crate::retrieval::memory_map) names the right predicate
162    /// 69% of the time against 2% without it.
163    #[serde(default)]
164    pub predicate_hint: Option<String>,
165    /// Time window, if the user named one.
166    #[serde(default)]
167    pub temporal: Option<TemporalConstraint>,
168    /// Hash of the transcript the plan came from, for cache keying.
169    pub source_transcript_hash: String,
170}
171
172impl RetrievalPlan {
173    /// A plan that says "do not search".
174    pub fn skip(turn_id: TurnId, generation: u64, transcript: &str) -> Self {
175        Self {
176            plan_id: PlanId::generate(),
177            turn_id,
178            generation,
179            requires_memory: false,
180            confidence: 1.0,
181            intent: RetrievalIntent::None,
182            entities: Vec::new(),
183            topics: Vec::new(),
184            predicates: Vec::new(),
185            lexical_queries: Vec::new(),
186            scopes: Vec::new(),
187            kind_filter: Vec::new(),
188            subject_hint: None,
189            predicate_hint: None,
190            temporal: None,
191            source_transcript_hash: stable_hash(transcript),
192        }
193    }
194
195    /// Apply the §12.2 caps and drop empties, returning the trimmed plan.
196    ///
197    /// Called on every plan regardless of origin, so a model that returns
198    /// forty search terms produces a fast plan rather than a slow one.
199    pub fn normalized(mut self) -> Self {
200        self.entities.retain(|e| !e.surface.trim().is_empty());
201        self.entities.truncate(limits::ENTITIES);
202
203        self.topics = dedup_non_empty(self.topics);
204        self.topics.truncate(limits::TOPICS);
205
206        self.predicates.dedup();
207        self.predicates.truncate(limits::PREDICATES);
208
209        self.lexical_queries = dedup_non_empty(self.lexical_queries)
210            .into_iter()
211            .map(|q| {
212                q.split_whitespace()
213                    .take(limits::QUERY_TERMS)
214                    .collect::<Vec<_>>()
215                    .join(" ")
216            })
217            .collect();
218        self.lexical_queries.truncate(limits::LEXICAL_QUERIES);
219
220        self.scopes.dedup();
221        self.scopes.truncate(limits::SCOPES);
222
223        self.confidence = self.confidence.clamp(0.0, 1.0);
224
225        // A plan with nothing to search for cannot require memory, whatever it
226        // claims about itself.
227        if self.lexical_queries.is_empty() && self.entities.is_empty() {
228            self.requires_memory = false;
229            self.intent = RetrievalIntent::None;
230        }
231        self
232    }
233
234    /// A cache key covering everything that changes the result.
235    pub fn cache_key(&self) -> String {
236        let mut parts = self.lexical_queries.clone();
237        parts.extend(self.entities.iter().flat_map(RetrievalEntity::forms));
238        parts.extend(self.scopes.iter().map(|s| s.scope_label().to_string()));
239        parts.sort();
240        stable_hash(&parts.join("|"))
241    }
242
243    /// Every entity surface form the index should boost on.
244    pub fn entity_forms(&self) -> Vec<String> {
245        self.entities
246            .iter()
247            .flat_map(RetrievalEntity::forms)
248            .collect()
249    }
250}
251
252fn dedup_non_empty(values: Vec<String>) -> Vec<String> {
253    let mut seen = std::collections::HashSet::new();
254    values
255        .into_iter()
256        .map(|v| v.trim().to_string())
257        .filter(|v| !v.is_empty())
258        .filter(|v| seen.insert(v.to_lowercase()))
259        .collect()
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265
266    fn plan_with(lexical: Vec<&str>) -> RetrievalPlan {
267        RetrievalPlan {
268            plan_id: PlanId::new("pln_1"),
269            turn_id: TurnId(1),
270            generation: 1,
271            requires_memory: true,
272            confidence: 1.2,
273            intent: RetrievalIntent::PersonalRecommendation,
274            entities: Vec::new(),
275            topics: Vec::new(),
276            predicates: Vec::new(),
277            lexical_queries: lexical.into_iter().map(str::to_string).collect(),
278            scopes: Vec::new(),
279            kind_filter: Vec::new(),
280            subject_hint: None,
281            predicate_hint: None,
282            temporal: None,
283            source_transcript_hash: "hash".into(),
284        }
285    }
286
287    #[test]
288    fn normalization_applies_every_cap() {
289        let mut plan = plan_with(vec!["a", "b", "c", "d", "e"]);
290        plan.entities = (0..9)
291            .map(|i| RetrievalEntity::surface(format!("e{i}")))
292            .collect();
293        plan.topics = (0..20).map(|i| format!("t{i}")).collect();
294        plan.scopes = vec![MemoryKind::Preference; 9];
295
296        let normalized = plan.normalized();
297        assert_eq!(normalized.entities.len(), limits::ENTITIES);
298        assert_eq!(normalized.topics.len(), limits::TOPICS);
299        assert_eq!(normalized.lexical_queries.len(), limits::LEXICAL_QUERIES);
300        assert_eq!(normalized.scopes.len(), 1, "duplicate scopes collapse");
301        assert!(normalized.confidence <= 1.0);
302    }
303
304    #[test]
305    fn overlong_queries_are_trimmed_to_their_leading_terms() {
306        let long = (0..40)
307            .map(|i| format!("w{i}"))
308            .collect::<Vec<_>>()
309            .join(" ");
310        let normalized = plan_with(vec![&long]).normalized();
311        assert_eq!(
312            normalized.lexical_queries[0].split_whitespace().count(),
313            limits::QUERY_TERMS
314        );
315    }
316
317    #[test]
318    fn blank_and_duplicate_terms_are_dropped() {
319        let mut plan = plan_with(vec!["quiet restaurant", "  ", "Quiet Restaurant"]);
320        plan.topics = vec!["food".into(), "".into(), "FOOD".into()];
321        let normalized = plan.normalized();
322        assert_eq!(normalized.lexical_queries.len(), 1);
323        assert_eq!(normalized.topics, vec!["food"]);
324    }
325
326    #[test]
327    fn a_plan_with_nothing_to_search_for_cannot_require_memory() {
328        let normalized = plan_with(vec!["", "   "]).normalized();
329        assert!(!normalized.requires_memory);
330        assert_eq!(normalized.intent, RetrievalIntent::None);
331    }
332
333    #[test]
334    fn a_skip_plan_requires_nothing() {
335        let plan = RetrievalPlan::skip(TurnId(3), 7, "what is the capital of France");
336        assert!(!plan.requires_memory);
337        assert!(!plan.intent.requires_memory());
338    }
339
340    #[test]
341    fn the_cache_key_ignores_query_ordering_but_not_content() {
342        let a = plan_with(vec!["quiet restaurant", "wife preference"]).normalized();
343        let b = plan_with(vec!["wife preference", "quiet restaurant"]).normalized();
344        assert_eq!(a.cache_key(), b.cache_key());
345
346        let c = plan_with(vec!["loud restaurant"]).normalized();
347        assert_ne!(a.cache_key(), c.cache_key());
348    }
349}
gemini_memory_rs/retrieval/plan.rs

gemini_memory_rs/retrieval/
plan.rs
