gemini_memory_rs/retrieval/
snapshot.rs

1//! Prepared memory snapshots — the immutable unit the model consumes.
2//!
3//! A snapshot is frozen when a turn begins and never changes while that turn is
4//! in flight. Without that, a slow retrieval landing mid-response could change
5//! what the model "remembers" halfway through a sentence.
6
7use chrono::{DateTime, Duration, Utc};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10
11use crate::bm25::MemoryOrigin;
12use crate::core::{MemoryId, MemoryKind, PlanId, SnapshotId, TemporalScope, TurnId};
13
14/// How long a prepared snapshot stays usable before it is considered stale.
15pub const SNAPSHOT_TTL_SECONDS: i64 = 120;
16
17/// One memory as it will be shown to the model.
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
19pub struct RetrievedMemory {
20    /// Which record this is, for provenance and explanation.
21    pub memory_id: MemoryId,
22    /// The sentence itself.
23    pub statement: String,
24    /// What sort of memory it is.
25    pub kind: MemoryKind,
26    /// How durable the fact is.
27    pub temporal_scope: TemporalScope,
28    /// Whether the fact is committed or still session-local.
29    pub origin: MemoryOrigin,
30    /// Fused retrieval score, retained for debugging.
31    pub score: f32,
32}
33
34impl RetrievedMemory {
35    /// The statement as it should be phrased to the model.
36    ///
37    /// An uncommitted session fact is hedged — "the user mentioned" rather than
38    /// "the user prefers" — because it has not yet survived reconciliation and
39    /// the model should not assert it as settled.
40    pub fn presented_statement(&self) -> String {
41        match self.origin {
42            MemoryOrigin::Canonical => self.statement.clone(),
43            MemoryOrigin::SessionOverlay => {
44                let trimmed = self.statement.trim_end_matches('.');
45                let lowered = lowercase_first(trimmed);
46                format!("The user mentioned that {lowered}.")
47            }
48        }
49    }
50
51    /// Estimated token cost of the presented statement.
52    pub fn token_cost(&self) -> usize {
53        estimate_tokens(&self.presented_statement())
54    }
55}
56
57fn lowercase_first(text: &str) -> String {
58    let mut chars = text.chars();
59    match chars.next() {
60        Some(first) => first.to_lowercase().collect::<String>() + chars.as_str(),
61        None => String::new(),
62    }
63}
64
65/// Approximate the token cost of a string.
66///
67/// Deliberately an over-estimate: the budget exists to protect latency and
68/// context, and being slightly conservative is cheaper than blowing the cap.
69pub fn estimate_tokens(text: &str) -> usize {
70    let words = text.split_whitespace().count();
71    let by_chars = text.chars().count().div_ceil(4);
72    words.max(by_chars).max(1)
73}
74
75/// An immutable, budgeted set of memories prepared for a turn.
76#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
77pub struct PreparedMemorySnapshot {
78    /// Snapshot identity.
79    pub snapshot_id: SnapshotId,
80    /// Canonical repository revision the snapshot was built from.
81    pub user_memory_revision: u64,
82    /// Session overlay revision the snapshot was built from.
83    pub session_overlay_revision: u64,
84    /// The plan that produced it.
85    pub retrieval_plan_id: Option<PlanId>,
86    /// The turn whose transcript produced it.
87    pub source_turn_id: TurnId,
88    /// The first turn it may be served to.
89    pub eligible_from_turn: TurnId,
90    /// The memories, in presentation order.
91    pub facts: Arc<[RetrievedMemory]>,
92    /// Total estimated tokens.
93    pub token_count: u16,
94    /// Cache key of the plan that produced it.
95    pub cache_key: String,
96    /// When it was prepared.
97    pub created_at: DateTime<Utc>,
98    /// When it stops being usable.
99    pub expires_at: DateTime<Utc>,
100}
101
102impl Default for PreparedMemorySnapshot {
103    fn default() -> Self {
104        let now = Utc::now();
105        Self {
106            snapshot_id: SnapshotId::generate(),
107            user_memory_revision: 0,
108            session_overlay_revision: 0,
109            retrieval_plan_id: None,
110            source_turn_id: TurnId::ZERO,
111            eligible_from_turn: TurnId::ZERO,
112            facts: Arc::from(Vec::new()),
113            token_count: 0,
114            cache_key: String::new(),
115            created_at: now,
116            expires_at: now + Duration::seconds(SNAPSHOT_TTL_SECONDS),
117        }
118    }
119}
120
121impl PreparedMemorySnapshot {
122    /// An empty snapshot for a turn.
123    pub fn empty(turn_id: TurnId) -> Self {
124        Self {
125            source_turn_id: turn_id,
126            eligible_from_turn: turn_id,
127            ..Default::default()
128        }
129    }
130
131    /// Whether the snapshot holds anything.
132    pub fn is_empty(&self) -> bool {
133        self.facts.is_empty()
134    }
135
136    /// Whether the snapshot is still fresh at `now`.
137    pub fn is_fresh(&self, now: DateTime<Utc>) -> bool {
138        now < self.expires_at
139    }
140
141    /// Whether this snapshot answers the query the model actually asked with.
142    ///
143    /// A prepared snapshot is speculative: it was built from the transcript,
144    /// not from the tool arguments. When the model asks something the
145    /// speculation did not anticipate, the caller falls back to a live search
146    /// rather than answering the wrong question quickly.
147    pub fn satisfies(&self, query: &str, now: DateTime<Utc>) -> bool {
148        if !self.is_fresh(now) || self.facts.is_empty() {
149            return false;
150        }
151        let query_terms = crate::bm25::tokenize(query);
152        if query_terms.is_empty() {
153            return true;
154        }
155        let covered: Vec<String> = self
156            .facts
157            .iter()
158            .flat_map(|f| crate::bm25::tokenize(&f.statement))
159            .collect();
160        let overlap = query_terms.iter().filter(|t| covered.contains(t)).count();
161        // At least a third of the asked-for terms should appear in what was
162        // prepared, or the speculation missed.
163        overlap * 3 >= query_terms.len()
164    }
165
166    /// The tool payload handed back to the model (§39.2).
167    pub fn to_tool_payload(&self) -> serde_json::Value {
168        if self.facts.is_empty() {
169            return serde_json::json!({ "status": "not_found", "facts": [] });
170        }
171        serde_json::json!({
172            "status": "found",
173            "facts": self
174                .facts
175                .iter()
176                .map(|f| serde_json::json!({
177                    "statement": f.presented_statement(),
178                    "kind": f.kind,
179                    "temporal_scope": f.temporal_scope,
180                }))
181                .collect::<Vec<_>>(),
182            "token_count": self.token_count,
183        })
184    }
185}
186
187/// Merge a live search with a prepared snapshot, ranked by RRF.
188///
189/// Serving one *or* the other is a choice the engine used to make with
190/// [`PreparedMemorySnapshot::satisfies`], and it made it badly: measured
191/// against snapshots that already held the answer, `satisfies` refused 65 of 93
192/// paraphrased questions, discarding a correct snapshot in favour of a lexical
193/// search that could not find one. Neither ranking is reliably better, so
194/// neither gets to win outright.
195///
196/// The same `1/(60 + rank)` the retriever already fuses lexical rankings with,
197/// so a fact both agree on rises and a fact only one found still gets a place.
198/// A stale prepared snapshot contributes nothing — it is a snapshot of a
199/// conversation that has moved on.
200/// `max_tokens` is a hard cap, not a target. Both inputs were assembled under
201/// it independently, which does *not* make their union compliant: two
202/// single-fact snapshots of 400 tokens each fuse to 800 under a 500 cap. The
203/// count limit alone cannot catch that — it is a limit on a different quantity
204/// — so the budget is reapplied here, the same way [`ContextAssembler`] applies
205/// it, and for the same reason: whatever this returns is what fills the model's
206/// context.
207///
208/// [`ContextAssembler`]: crate::retrieval::ContextAssembler
209pub fn fuse_snapshots(
210    live: &PreparedMemorySnapshot,
211    prepared: &PreparedMemorySnapshot,
212    max_memories: usize,
213    max_tokens: usize,
214    now: DateTime<Utc>,
215) -> PreparedMemorySnapshot {
216    if prepared.is_empty() || !prepared.is_fresh(now) {
217        return live.clone();
218    }
219    if live.is_empty() {
220        return prepared.clone();
221    }
222
223    let k = crate::retrieval::fusion::RRF_K as f64;
224    let mut scores: Vec<(f64, RetrievedMemory)> = Vec::new();
225    for ranking in [&live.facts, &prepared.facts] {
226        for (rank, fact) in ranking.iter().enumerate() {
227            let contribution = 1.0 / (k + rank as f64 + 1.0);
228            match scores
229                .iter_mut()
230                .find(|(_, f)| f.memory_id == fact.memory_id)
231            {
232                Some((score, _)) => *score += contribution,
233                None => scores.push((contribution, fact.clone())),
234            }
235        }
236    }
237    scores.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
238    scores.truncate(max_memories);
239
240    // Fill in fused order until the next fact would break the budget, then
241    // stop. `break` rather than `continue`: skipping an expensive fact to fit a
242    // cheaper one behind it would serve a lower-ranked memory in place of a
243    // higher-ranked one, which is the fusion silently disagreeing with itself.
244    // Running out of budget is a truncation, not a re-ranking.
245    let mut facts: Vec<RetrievedMemory> = Vec::with_capacity(scores.len());
246    let mut tokens = 0usize;
247    for (_, fact) in scores {
248        let cost = fact.token_cost();
249        if tokens + cost > max_tokens {
250            break;
251        }
252        tokens += cost;
253        facts.push(fact);
254    }
255    let token_count = tokens.min(u16::MAX as usize) as u16;
256    PreparedMemorySnapshot {
257        facts: Arc::from(facts),
258        token_count,
259        // Provenance follows the live search: it is the one that answered the
260        // question the model actually asked.
261        ..live.clone()
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    fn fact(statement: &str, origin: MemoryOrigin) -> RetrievedMemory {
270        RetrievedMemory {
271            memory_id: MemoryId::new("mem_1"),
272            statement: statement.to_string(),
273            kind: MemoryKind::Preference,
274            temporal_scope: TemporalScope::Persistent,
275            origin,
276            score: 3.0,
277        }
278    }
279
280    fn snapshot(facts: Vec<RetrievedMemory>) -> PreparedMemorySnapshot {
281        let token_count = facts
282            .iter()
283            .map(super::RetrievedMemory::token_cost)
284            .sum::<usize>() as u16;
285        PreparedMemorySnapshot {
286            facts: Arc::from(facts),
287            token_count,
288            ..Default::default()
289        }
290    }
291
292    #[test]
293    fn canonical_facts_are_asserted_and_overlay_facts_are_hedged() {
294        assert_eq!(
295            fact("The user is pescatarian.", MemoryOrigin::Canonical).presented_statement(),
296            "The user is pescatarian."
297        );
298        assert_eq!(
299            fact("The user is pescatarian.", MemoryOrigin::SessionOverlay).presented_statement(),
300            "The user mentioned that the user is pescatarian."
301        );
302    }
303
304    /// The cap `ContextAssembler` enforces has to survive fusion.
305    ///
306    /// Both inputs are individually legal — one fact each, comfortably under a
307    /// generous cap — and their union is not. Before the budget was reapplied
308    /// here, fusion truncated by *count* and then summed whatever it kept, so
309    /// two compliant snapshots produced a non-compliant one and the guarantee
310    /// the assembler makes stopped holding the moment speculation was involved.
311    #[test]
312    fn fusing_two_compliant_snapshots_does_not_exceed_the_token_cap() {
313        let long = "The user has a standing preference that every dinner \
314                    reservation be somewhere quiet enough to hold a conversation \
315                    without raising a voice, ideally with outdoor seating.";
316        let mut live_fact = fact(long, MemoryOrigin::Canonical);
317        live_fact.memory_id = MemoryId::new("mem_live");
318        let mut prepared_fact = fact(long, MemoryOrigin::Canonical);
319        prepared_fact.memory_id = MemoryId::new("mem_prepared");
320
321        let each = live_fact.token_cost();
322        assert!(each > 0, "the fixture fact must cost something");
323        // A cap that either fact fits inside and both together do not.
324        let cap = each + each / 2;
325
326        let fused = fuse_snapshots(
327            &snapshot(vec![live_fact]),
328            &snapshot(vec![prepared_fact]),
329            10,
330            cap,
331            Utc::now(),
332        );
333        assert!(
334            usize::from(fused.token_count) <= cap,
335            "fused snapshot is {} tokens against a {cap}-token cap",
336            fused.token_count
337        );
338        assert_eq!(
339            fused.facts.len(),
340            1,
341            "the budget should have stopped at the first fact"
342        );
343        // The reported count has to match what is actually carried, or the cap
344        // is enforced against a number nobody is serving.
345        assert_eq!(
346            usize::from(fused.token_count),
347            fused
348                .facts
349                .iter()
350                .map(RetrievedMemory::token_cost)
351                .sum::<usize>(),
352        );
353    }
354
355    /// A cap generous enough for both must not drop either — otherwise the fix
356    /// above would "pass" by always truncating.
357    #[test]
358    fn a_cap_that_fits_both_keeps_both() {
359        let mut live_fact = fact("The user drinks cortados.", MemoryOrigin::Canonical);
360        live_fact.memory_id = MemoryId::new("mem_live");
361        let mut prepared_fact = fact("Rhea prefers quiet places.", MemoryOrigin::Canonical);
362        prepared_fact.memory_id = MemoryId::new("mem_prepared");
363
364        let fused = fuse_snapshots(
365            &snapshot(vec![live_fact]),
366            &snapshot(vec![prepared_fact]),
367            10,
368            10_000,
369            Utc::now(),
370        );
371        assert_eq!(fused.facts.len(), 2);
372    }
373
374    #[test]
375    fn an_empty_snapshot_reports_not_found() {
376        let payload = PreparedMemorySnapshot::empty(TurnId(1)).to_tool_payload();
377        assert_eq!(payload["status"], "not_found");
378        assert_eq!(payload["facts"].as_array().unwrap().len(), 0);
379    }
380
381    #[test]
382    fn the_payload_carries_statements_kinds_and_a_token_count() {
383        let payload = snapshot(vec![fact(
384            "The user is pescatarian.",
385            MemoryOrigin::Canonical,
386        )])
387        .to_tool_payload();
388        assert_eq!(payload["status"], "found");
389        assert_eq!(payload["facts"][0]["statement"], "The user is pescatarian.");
390        assert_eq!(payload["facts"][0]["kind"], "preference");
391        assert_eq!(payload["facts"][0]["temporal_scope"], "persistent");
392        assert!(payload["token_count"].as_u64().unwrap() > 0);
393    }
394
395    #[test]
396    fn a_snapshot_satisfies_a_query_it_actually_covers() {
397        let snap = snapshot(vec![fact(
398            "Rhea prefers quiet restaurants.",
399            MemoryOrigin::Canonical,
400        )]);
401        let now = Utc::now();
402        assert!(snap.satisfies("quiet restaurants for Rhea", now));
403        assert!(!snap.satisfies("what medication does the user take", now));
404    }
405
406    #[test]
407    fn a_stale_or_empty_snapshot_satisfies_nothing() {
408        let now = Utc::now();
409        let mut stale = snapshot(vec![fact("Rhea prefers quiet.", MemoryOrigin::Canonical)]);
410        stale.expires_at = now - Duration::seconds(1);
411        assert!(!stale.satisfies("quiet", now));
412        assert!(!PreparedMemorySnapshot::empty(TurnId(1)).satisfies("anything", now));
413    }
414
415    #[test]
416    fn token_estimates_are_conservative_and_never_zero() {
417        assert!(estimate_tokens("The user is pescatarian.") >= 5);
418        assert_eq!(estimate_tokens(""), 1);
419    }
420}
gemini_memory_rs/retrieval/snapshot.rs

gemini_memory_rs/retrieval/
snapshot.rs
