gemini_memory_rs/retrieval/
assembler.rs

1//! Context assembly — turning a ranked candidate list into a bounded snapshot.
2//!
3//! The assembler receives ten to twenty candidates and emits three to five
4//! facts. Everything it does is subtraction: drop weak matches, drop
5//! near-duplicates of the same predicate, and stop at the token budget. A
6//! memory block that grows without bound is a latency and a distraction
7//! problem, not a helpfulness one.
8
9use chrono::{DateTime, Duration, Utc};
10use std::collections::HashMap;
11use std::sync::Arc;
12
13use super::fusion::FusedCandidate;
14use super::plan::{RetrievalIntent, RetrievalPlan};
15use super::snapshot::{PreparedMemorySnapshot, RetrievedMemory, SNAPSHOT_TTL_SECONDS};
16use crate::bm25::{MemoryIndex, MemoryOrigin};
17use crate::core::{RetrievalConfig, SnapshotId, TurnId};
18
19/// Assembles snapshots under a budget.
20#[derive(Debug, Clone)]
21pub struct ContextAssembler {
22    config: RetrievalConfig,
23}
24
25impl ContextAssembler {
26    /// An assembler using the given budgets.
27    pub fn new(config: RetrievalConfig) -> Self {
28        Self { config }
29    }
30
31    /// The budgets in force.
32    pub fn config(&self) -> &RetrievalConfig {
33        &self.config
34    }
35
36    /// Build a snapshot from fused candidates.
37    ///
38    /// `index` supplies the predicate and temporal metadata that ranking does
39    /// not carry; candidates missing from it are dropped rather than guessed at.
40    pub fn assemble(
41        &self,
42        plan: &RetrievalPlan,
43        candidates: &[FusedCandidate],
44        index: &MemoryIndex,
45        overlay: Option<&MemoryIndex>,
46        revisions: (u64, u64),
47        now: DateTime<Utc>,
48    ) -> PreparedMemorySnapshot {
49        let mut facts: Vec<RetrievedMemory> = Vec::new();
50        let mut tokens = 0usize;
51        let mut per_predicate: HashMap<String, usize> = HashMap::new();
52
53        // An explicit recall request is the user asking to hear what is stored,
54        // so the diversity cap that stops ordinary turns being flooded with
55        // near-identical preferences is relaxed.
56        let per_predicate_cap = if plan.intent == RetrievalIntent::ExplicitRecall {
57            self.config.max_memories
58        } else {
59            self.config.max_per_predicate
60        };
61
62        for candidate in candidates {
63            if facts.len() >= self.config.max_memories {
64                break;
65            }
66            if candidate.hit.score < self.config.minimum_candidate_score {
67                continue;
68            }
69            let Some(doc) = overlay
70                .and_then(|o| o.get(&candidate.hit.id))
71                .or_else(|| index.get(&candidate.hit.id))
72            else {
73                continue;
74            };
75
76            let predicate = doc.predicate.to_string();
77            let used = per_predicate.entry(predicate).or_insert(0);
78            if *used >= per_predicate_cap {
79                continue;
80            }
81
82            let fact = RetrievedMemory {
83                memory_id: doc.id.clone(),
84                statement: doc.statement.clone(),
85                kind: doc.kind,
86                temporal_scope: doc.temporal_scope,
87                origin: doc.origin,
88                score: candidate.hit.score,
89            };
90            let cost = fact.token_cost();
91            if tokens + cost > self.config.max_tokens {
92                // The budget is a hard cap: stop rather than truncate a
93                // statement into something that reads as a different fact.
94                break;
95            }
96            *used += 1;
97            tokens += cost;
98            facts.push(fact);
99
100            if tokens >= self.config.target_tokens {
101                break;
102            }
103        }
104
105        // Committed facts first: the model should see settled memory before
106        // anything still provisional.
107        facts.sort_by_key(|f| match f.origin {
108            MemoryOrigin::Canonical => 0,
109            MemoryOrigin::SessionOverlay => 1,
110        });
111
112        PreparedMemorySnapshot {
113            snapshot_id: SnapshotId::generate(),
114            user_memory_revision: revisions.0,
115            session_overlay_revision: revisions.1,
116            retrieval_plan_id: Some(plan.plan_id.clone()),
117            source_turn_id: plan.turn_id,
118            eligible_from_turn: TurnId(plan.turn_id.0 + 1),
119            token_count: tokens.min(u16::MAX as usize) as u16,
120            facts: Arc::from(facts),
121            cache_key: plan.cache_key(),
122            created_at: now,
123            expires_at: now + Duration::seconds(SNAPSHOT_TTL_SECONDS),
124        }
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use crate::bm25::{IndexedMemory, SearchExplanation, SearchHit};
132    use crate::core::{
133        CanonicalMemory, CanonicalPredicate, EntityRef, EvidenceCounters, Explicitness, MemoryId,
134        MemoryKind, MemorySource, MemoryStatus, MemoryValue, PlanId, PrivacyMetadata,
135        RetrievalMetadata, SessionId, TemporalMetadata, TemporalScope, UserId,
136    };
137
138    fn record(id: &str, predicate: &str, statement: &str) -> CanonicalMemory {
139        CanonicalMemory {
140            id: MemoryId::new(id),
141            owner: UserId::new("usr_1"),
142            kind: MemoryKind::Preference,
143            predicate: CanonicalPredicate::new(predicate),
144            status: MemoryStatus::Active,
145            confidence: 0.9,
146            subject: EntityRef::user(),
147            value: MemoryValue::Text(statement.into()),
148            statement: statement.into(),
149            evidence_summary: "stated".into(),
150            source: MemorySource::from_explicitness(
151                Explicitness::ExplicitStatement,
152                SessionId::new("ses_1"),
153                TurnId(1),
154            ),
155            temporal: TemporalMetadata::created_at(Utc::now()),
156            retrieval: RetrievalMetadata {
157                subject: "user".into(),
158                ..Default::default()
159            },
160            evidence: EvidenceCounters::first(),
161            privacy: PrivacyMetadata::default(),
162            temporal_scope: TemporalScope::Persistent,
163            supersedes: Vec::new(),
164            superseded_by: None,
165            qualifier: None,
166        }
167    }
168
169    fn candidate(id: &str, score: f32) -> FusedCandidate {
170        FusedCandidate {
171            hit: SearchHit {
172                id: MemoryId::new(id),
173                score,
174                statement: String::new(),
175                kind: MemoryKind::Preference,
176                origin: MemoryOrigin::Canonical,
177                explanation: SearchExplanation {
178                    memory_id: MemoryId::new(id),
179                    components: Vec::new(),
180                    boosts: Vec::new(),
181                    lexical_score: score,
182                    final_score: score,
183                },
184            },
185            rrf_score: score,
186            appearances: 1,
187            best_rank: 0,
188        }
189    }
190
191    fn plan() -> RetrievalPlan {
192        RetrievalPlan {
193            plan_id: PlanId::new("pln_1"),
194            turn_id: TurnId(4),
195            generation: 4,
196            requires_memory: true,
197            confidence: 0.9,
198            intent: RetrievalIntent::PersonalRecommendation,
199            entities: Vec::new(),
200            topics: vec!["food".into()],
201            predicates: Vec::new(),
202            lexical_queries: vec!["food".into()],
203            scopes: Vec::new(),
204            kind_filter: Vec::new(),
205            subject_hint: None,
206            predicate_hint: None,
207            temporal: None,
208            source_transcript_hash: "h".into(),
209        }
210    }
211
212    #[test]
213    fn caps_the_number_of_memories_returned() {
214        let records: Vec<_> = (0..10)
215            .map(|i| record(&format!("mem_{i}"), &format!("pred_{i}"), "A preference."))
216            .collect();
217        let index = MemoryIndex::build(records.iter().map(IndexedMemory::from_canonical));
218        let candidates: Vec<_> = (0..10)
219            .map(|i| candidate(&format!("mem_{i}"), 5.0))
220            .collect();
221
222        let snapshot = ContextAssembler::new(RetrievalConfig::default()).assemble(
223            &plan(),
224            &candidates,
225            &index,
226            None,
227            (1, 0),
228            Utc::now(),
229        );
230        assert!(snapshot.facts.len() <= RetrievalConfig::default().max_memories);
231    }
232
233    #[test]
234    fn never_exceeds_the_hard_token_cap() {
235        let long = "The user has an extremely detailed preference ".repeat(20);
236        let records: Vec<_> = (0..10)
237            .map(|i| record(&format!("mem_{i}"), &format!("pred_{i}"), &long))
238            .collect();
239        let index = MemoryIndex::build(records.iter().map(IndexedMemory::from_canonical));
240        let candidates: Vec<_> = (0..10)
241            .map(|i| candidate(&format!("mem_{i}"), 5.0))
242            .collect();
243
244        let config = RetrievalConfig::default();
245        let snapshot = ContextAssembler::new(config.clone()).assemble(
246            &plan(),
247            &candidates,
248            &index,
249            None,
250            (1, 0),
251            Utc::now(),
252        );
253        assert!(
254            usize::from(snapshot.token_count) <= config.max_tokens,
255            "token count {} exceeded cap",
256            snapshot.token_count
257        );
258    }
259
260    #[test]
261    fn limits_near_duplicates_of_the_same_predicate() {
262        let records: Vec<_> = (0..5)
263            .map(|i| {
264                record(
265                    &format!("mem_{i}"),
266                    "coffee_order",
267                    "The user drinks flat whites.",
268                )
269            })
270            .collect();
271        let index = MemoryIndex::build(records.iter().map(IndexedMemory::from_canonical));
272        let candidates: Vec<_> = (0..5)
273            .map(|i| candidate(&format!("mem_{i}"), 5.0))
274            .collect();
275
276        let snapshot = ContextAssembler::new(RetrievalConfig::default()).assemble(
277            &plan(),
278            &candidates,
279            &index,
280            None,
281            (1, 0),
282            Utc::now(),
283        );
284        assert_eq!(
285            snapshot.facts.len(),
286            RetrievalConfig::default().max_per_predicate
287        );
288    }
289
290    #[test]
291    fn an_explicit_recall_request_relaxes_the_diversity_cap() {
292        let records: Vec<_> = (0..5)
293            .map(|i| {
294                record(
295                    &format!("mem_{i}"),
296                    "coffee_order",
297                    "The user drinks flat whites.",
298                )
299            })
300            .collect();
301        let index = MemoryIndex::build(records.iter().map(IndexedMemory::from_canonical));
302        let candidates: Vec<_> = (0..5)
303            .map(|i| candidate(&format!("mem_{i}"), 5.0))
304            .collect();
305
306        let mut plan = plan();
307        plan.intent = RetrievalIntent::ExplicitRecall;
308        let snapshot = ContextAssembler::new(RetrievalConfig::default()).assemble(
309            &plan,
310            &candidates,
311            &index,
312            None,
313            (1, 0),
314            Utc::now(),
315        );
316        assert!(snapshot.facts.len() > RetrievalConfig::default().max_per_predicate);
317    }
318
319    #[test]
320    fn weak_candidates_are_dropped() {
321        let index = MemoryIndex::build([IndexedMemory::from_canonical(&record(
322            "mem_weak",
323            "pred",
324            "Barely relevant.",
325        ))]);
326        let snapshot = ContextAssembler::new(RetrievalConfig::default()).assemble(
327            &plan(),
328            &[candidate("mem_weak", 0.01)],
329            &index,
330            None,
331            (1, 0),
332            Utc::now(),
333        );
334        assert!(snapshot.is_empty());
335        assert_eq!(snapshot.to_tool_payload()["status"], "not_found");
336    }
337
338    #[test]
339    fn committed_facts_are_presented_before_provisional_ones() {
340        let canonical = IndexedMemory::from_canonical(&record("mem_c", "p1", "Committed fact."));
341        let overlay_doc =
342            IndexedMemory::from_canonical(&record("mem_o", "p2", "Provisional fact."))
343                .as_session_overlay();
344        let index = MemoryIndex::build([canonical]);
345        let overlay = MemoryIndex::build([overlay_doc]);
346
347        let snapshot = ContextAssembler::new(RetrievalConfig::default()).assemble(
348            &plan(),
349            // Overlay ranks higher, but presentation order still puts the
350            // settled fact first.
351            &[candidate("mem_o", 9.0), candidate("mem_c", 5.0)],
352            &index,
353            Some(&overlay),
354            (1, 3),
355            Utc::now(),
356        );
357        assert_eq!(snapshot.facts[0].memory_id.as_str(), "mem_c");
358        assert_eq!(snapshot.facts[1].origin, MemoryOrigin::SessionOverlay);
359    }
360
361    #[test]
362    fn the_snapshot_records_what_it_was_built_from() {
363        let index = MemoryIndex::build([IndexedMemory::from_canonical(&record(
364            "mem_a", "p", "A fact.",
365        ))]);
366        let snapshot = ContextAssembler::new(RetrievalConfig::default()).assemble(
367            &plan(),
368            &[candidate("mem_a", 5.0)],
369            &index,
370            None,
371            (7, 3),
372            Utc::now(),
373        );
374        assert_eq!(snapshot.user_memory_revision, 7);
375        assert_eq!(snapshot.session_overlay_revision, 3);
376        assert_eq!(snapshot.source_turn_id, TurnId(4));
377        assert_eq!(snapshot.eligible_from_turn, TurnId(5));
378        assert_eq!(
379            snapshot.retrieval_plan_id.as_ref().unwrap().as_str(),
380            "pln_1"
381        );
382    }
383}
gemini_memory_rs/retrieval/assembler.rs

gemini_memory_rs/retrieval/
assembler.rs
