gemini_memory_rs/retrieval/
vocabulary.rs

1//! An inventory of what a user's memory contains, for the model to read.
2//!
3//! # Why this exists
4//!
5//! [`recall_context`](crate::runtime::tools) lets the model narrow a search by
6//! `about` and `attribute`. Those are only useful if it can name the values,
7//! and it cannot guess them: the corpus files a haircut under `barber` and a
8//! coffee order under `beverage_preference`, and a model asked to invent those
9//! from an utterance produces something reasonable and wrong.
10//!
11//! Measured over 93 questions (`tests/memory_map_probe.rs`), asking
12//! `gemini-2.5-flash-lite` to fill the filter fields:
13//!
14//! | condition | `about` | `attribute` | `about` + `attribute` |
15//! |---|---|---|---|
16//! | with no map | 49% | 2% | 2% |
17//! | **with this map** | 67% | **69%** | **48%** |
18//!
19//! `attribute` goes from 2% to 69%. That is the difference between a filter
20//! that costs more than it earns and one that pays: with the filter applied as
21//! a *soft* ranking, break-even is 8% accuracy, so 2% is a net loss and 48% is
22//! worth about five questions on top-5.
23//!
24//! # It is a fixed cost
25//!
26//! The map is bounded by the user's vocabulary rather than by how much they
27//! have accumulated — people acquire more facts about the same handful of
28//! people and properties, not endlessly more kinds of thing:
29//!
30//! | records | subjects | predicates | map tokens |
31//! |---|---|---|---|
32//! | 250 | 38 | 16 | 242 |
33//! | 1,000 | 42 | 16 | 262 |
34//! | 16,000 | 42 | 16 | **282** |
35//!
36//! Flat from a thousand records up.
37//!
38//! # Where to put it
39//!
40//! In the **system instruction**, not the tool description. Live sessions fix
41//! tool declarations at connect time and the corpus grows while the session
42//! runs, so a map in the schema goes stale and cannot be refreshed. Instructions
43//! can be updated mid-session. It is also set once and cached rather than
44//! resent per call, which is what makes 282 tokens the whole price.
45
46use std::collections::BTreeMap;
47
48use crate::core::{CanonicalMemory, MemoryStatus};
49
50/// How many values of each field the map names before summarising the rest.
51///
52/// Chosen so the map stays inside a few hundred tokens even for a user whose
53/// vocabulary is unusually wide. Values are listed most-frequent first, so a
54/// truncated tail is the long tail.
55pub const DEFAULT_LIMIT: usize = 40;
56
57/// Render the inventory a model needs to write filters.
58///
59/// Only active records are counted — a model should not narrow a search to a
60/// value that exists solely on superseded facts. Counts are included because
61/// they tell the model which values are worth filtering by and which are
62/// one-offs, and they cost almost nothing.
63///
64/// Returns an empty string when there is nothing to describe, so it can be
65/// concatenated into an instruction unconditionally.
66pub fn memory_map(records: &[CanonicalMemory]) -> String {
67    memory_map_with_limit(records, DEFAULT_LIMIT)
68}
69
70/// [`memory_map`] with an explicit cap on values named per field.
71pub fn memory_map_with_limit(records: &[CanonicalMemory], limit: usize) -> String {
72    let active: Vec<&CanonicalMemory> = records
73        .iter()
74        .filter(|m| m.status == MemoryStatus::Active)
75        .collect();
76    if active.is_empty() {
77        return String::new();
78    }
79
80    let mut subjects: BTreeMap<&str, usize> = BTreeMap::new();
81    let mut predicates: BTreeMap<&str, usize> = BTreeMap::new();
82    for memory in &active {
83        *subjects
84            .entry(memory.retrieval.subject.as_str())
85            .or_default() += 1;
86        *predicates.entry(memory.predicate.as_str()).or_default() += 1;
87    }
88    assemble(subjects, predicates, limit)
89}
90
91/// The same map, built from the live BM25 index rather than a record list.
92///
93/// This is the form the runtime uses. The canonical index is already resident,
94/// so the map costs one pass over memory rather than a repository read — which
95/// matters because it has to be rebuilt whenever the corpus changes, and the
96/// corpus changes mid-conversation.
97pub fn memory_map_from_index(index: &crate::bm25::MemoryIndex, limit: usize) -> String {
98    let now = chrono::Utc::now();
99    let mut subjects: BTreeMap<&str, usize> = BTreeMap::new();
100    let mut predicates: BTreeMap<&str, usize> = BTreeMap::new();
101    for doc in index.documents() {
102        // `is_retrievable` rather than a status check, because it also drops
103        // records that have expired — exactly as unavailable as superseded
104        // ones, and exactly as wrong to advertise as filterable values.
105        if !doc.is_retrievable(now) {
106            continue;
107        }
108        *subjects.entry(doc.subject_form.as_str()).or_default() += 1;
109        *predicates.entry(doc.predicate.as_str()).or_default() += 1;
110    }
111    assemble(subjects, predicates, limit)
112}
113
114/// Render the two vocabularies into the block the model reads.
115fn assemble(
116    subjects: BTreeMap<&str, usize>,
117    predicates: BTreeMap<&str, usize>,
118    limit: usize,
119) -> String {
120    if subjects.is_empty() && predicates.is_empty() {
121        return String::new();
122    }
123    let mut map = String::from(
124        "The values that exist in this user's memory. When you call \
125         recall_context, `about` and `attribute` must come from these lists — \
126         omit either if none fits.\n",
127    );
128    // Subjects are listed in their *normalised* form ("rhea", not "Rhea"),
129    // because that is the form the hint matcher compares against. Showing the
130    // display name would read better and would silently fail to match wherever
131    // the two differ — a subject displayed as "Rhea Kapoor" is stored as
132    // "rhea", and a model echoing the pretty version would narrow to nothing.
133    map.push_str(&render("about", subjects, limit));
134    map.push_str(&render("attribute", predicates, limit));
135    map
136}
137
138/// One line: `label: value (count), value (count), and N more`.
139fn render(label: &str, counts: BTreeMap<&str, usize>, limit: usize) -> String {
140    let mut sorted: Vec<(&str, usize)> = counts.into_iter().collect();
141    // Frequency first so truncation drops the long tail, then name so the
142    // output is stable for a given corpus rather than varying run to run.
143    sorted.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(b.0)));
144    let shown: Vec<String> = sorted
145        .iter()
146        .take(limit)
147        .map(|(name, count)| format!("{name} ({count})"))
148        .collect();
149    let tail = sorted.len().saturating_sub(limit);
150    let suffix = if tail > 0 {
151        format!(", and {tail} more")
152    } else {
153        String::new()
154    };
155    format!("{label}: {}{suffix}\n", shown.join(", "))
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use chrono::Utc;
162
163    use crate::core::{
164        CanonicalPredicate, EntityRef, Explicitness, MemoryId, MemoryKind, MemorySource,
165        MemoryValue, RetrievalMetadata, SessionId, TemporalMetadata, TemporalScope, TurnId, UserId,
166    };
167
168    fn record(subject: &str, predicate: &str, statement: &str) -> CanonicalMemory {
169        CanonicalMemory {
170            id: MemoryId::new(format!("mem_{predicate}_{subject}")),
171            owner: UserId::new("usr_test"),
172            kind: MemoryKind::Preference,
173            predicate: CanonicalPredicate::new(predicate),
174            status: MemoryStatus::Active,
175            confidence: 0.9,
176            subject: EntityRef::named(subject),
177            value: MemoryValue::Text(statement.into()),
178            statement: statement.into(),
179            evidence_summary: "stated".into(),
180            source: MemorySource::from_explicitness(
181                Explicitness::ExplicitStatement,
182                SessionId::new("ses_1"),
183                TurnId(1),
184            ),
185            temporal: TemporalMetadata::created_at(Utc::now()),
186            retrieval: RetrievalMetadata {
187                subject: crate::core::normalize_token(subject),
188                ..Default::default()
189            },
190            temporal_scope: TemporalScope::Persistent,
191            qualifier: None,
192            evidence: Default::default(),
193            privacy: Default::default(),
194            supersedes: Vec::new(),
195            superseded_by: None,
196        }
197    }
198
199    #[test]
200    fn it_names_the_values_a_model_would_otherwise_have_to_guess() {
201        let records = vec![
202            record("user", "barber", "The user's barber is Deepa."),
203            record("user", "beverage_preference", "The user drinks cortados."),
204            record("Rhea", "barber", "Rhea's barber is Sam."),
205        ];
206        let map = memory_map(&records);
207        assert!(map.contains("user (2)"), "{map}");
208        // Normalised, not the display form — see the note in the renderer.
209        assert!(map.contains("rhea (1)"), "{map}");
210        assert!(map.contains("barber (2)"), "{map}");
211        assert!(map.contains("beverage_preference (1)"), "{map}");
212    }
213
214    #[test]
215    fn superseded_records_do_not_advertise_values_that_no_longer_apply() {
216        let mut stale = record("Priya", "barber", "Priya's barber was Sam.");
217        stale.status = MemoryStatus::Superseded;
218        let records = vec![
219            record("user", "barber", "The user's barber is Deepa."),
220            stale,
221        ];
222        let map = memory_map(&records);
223        assert!(
224            !map.contains("priya"),
225            "a value only present on superseded records should not be offered: {map}"
226        );
227    }
228
229    #[test]
230    fn the_tail_is_summarised_rather_than_dropped_silently() {
231        let records: Vec<CanonicalMemory> = (0..10)
232            .map(|i| record("user", &format!("attribute_{i}"), "x"))
233            .collect();
234        let map = memory_map_with_limit(&records, 3);
235        assert!(map.contains("and 7 more"), "{map}");
236    }
237
238    /// The runtime path: the map built from the live index must say the same
239    /// thing as the map built from records, or the thing measured and the thing
240    /// shipped are different maps.
241    #[test]
242    fn the_index_backed_map_agrees_with_the_record_backed_one() {
243        let records = vec![
244            record("user", "barber", "The user's barber is Deepa."),
245            record("user", "beverage_preference", "The user drinks cortados."),
246            record("Rhea", "barber", "Rhea's barber is Sam."),
247        ];
248        let index = crate::bm25::MemoryIndex::build(
249            records
250                .iter()
251                .map(crate::bm25::IndexedMemory::from_canonical),
252        );
253        assert_eq!(
254            memory_map_from_index(&index, DEFAULT_LIMIT),
255            memory_map(&records),
256            "the two builders must not drift; the experiments measured one and \
257             the runtime serves the other"
258        );
259    }
260
261    /// A superseded record is dropped by the index, so it cannot reach the map.
262    #[test]
263    fn the_index_backed_map_omits_records_that_are_not_retrievable() {
264        let mut stale = record("Priya", "barber", "Priya's barber was Sam.");
265        stale.status = MemoryStatus::Superseded;
266        let records = [
267            record("user", "barber", "The user's barber is Deepa."),
268            stale,
269        ];
270        let index = crate::bm25::MemoryIndex::build(
271            records
272                .iter()
273                .map(crate::bm25::IndexedMemory::from_canonical),
274        );
275        let map = memory_map_from_index(&index, DEFAULT_LIMIT);
276        assert!(
277            !map.contains("priya"),
278            "a superseded record must not be offered as a filter value: {map}"
279        );
280        assert!(map.contains("user"), "{map}");
281    }
282
283    #[test]
284    fn an_empty_index_renders_nothing_to_concatenate() {
285        let index = crate::bm25::MemoryIndex::new();
286        assert!(memory_map_from_index(&index, DEFAULT_LIMIT).is_empty());
287    }
288
289    #[test]
290    fn an_empty_corpus_renders_nothing_to_concatenate() {
291        assert!(memory_map(&[]).is_empty());
292    }
293
294    /// The property the token budget rests on: the map grows with the
295    /// vocabulary, not with how much the user has accumulated.
296    #[test]
297    fn more_records_over_the_same_vocabulary_do_not_grow_the_map() {
298        let few: Vec<CanonicalMemory> = (0..10)
299            .map(|_| record("user", "barber", "The user's barber is Deepa."))
300            .collect();
301        let many: Vec<CanonicalMemory> = (0..10_000)
302            .map(|_| record("user", "barber", "The user's barber is Deepa."))
303            .collect();
304        // Same values, so the same lines — only the counts differ in width.
305        assert!(memory_map(&many).len() < memory_map(&few).len() + 16);
306    }
307}
gemini_memory_rs/retrieval/vocabulary.rs

gemini_memory_rs/retrieval/
vocabulary.rs
