gemini_memory_rs/bm25/
schema.rs

1//! The lexical surface of a memory: which text is indexed, in which field, and
2//! how much each field is worth.
3//!
4//! Fields are weighted rather than concatenated because *where* a term matched
5//! carries most of the signal. "Rhea" appearing as the subject of a record is a
6//! far stronger indication of relevance than "Rhea" appearing in the middle of
7//! a sentence.
8
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11
12use crate::core::{
13    CanonicalMemory, CanonicalPredicate, MemoryId, MemoryKind, MemoryStatus, TemporalScope,
14    normalize_token,
15};
16
17/// The indexed fields, in weight order.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum Field {
21    /// The memory's subject surface form.
22    Subject,
23    /// Other entities the memory mentions.
24    Entities,
25    /// Paraphrases the fact may be asked for by.
26    Aliases,
27    /// The canonical predicate.
28    Predicate,
29    /// Topical tags.
30    Tags,
31    /// A place the fact is scoped to.
32    Location,
33    /// The natural-language statement.
34    Statement,
35}
36
37impl Field {
38    /// Every field, in a stable order.
39    pub const ALL: [Field; 7] = [
40        Field::Subject,
41        Field::Entities,
42        Field::Aliases,
43        Field::Predicate,
44        Field::Tags,
45        Field::Location,
46        Field::Statement,
47    ];
48
49    /// Index into per-field arrays.
50    pub fn slot(self) -> usize {
51        match self {
52            Field::Subject => 0,
53            Field::Entities => 1,
54            Field::Aliases => 2,
55            Field::Predicate => 3,
56            Field::Tags => 4,
57            Field::Location => 5,
58            Field::Statement => 6,
59        }
60    }
61
62    /// The field's contribution multiplier (§13.3).
63    pub fn weight(self) -> f32 {
64        match self {
65            Field::Subject => 3.0,
66            Field::Entities => 3.0,
67            Field::Aliases => 2.5,
68            Field::Predicate => 2.2,
69            Field::Tags => 2.0,
70            Field::Location => 1.5,
71            Field::Statement => 1.0,
72        }
73    }
74
75    /// A short label for search explanations.
76    pub fn label(self) -> &'static str {
77        match self {
78            Field::Subject => "subject",
79            Field::Entities => "entities",
80            Field::Aliases => "aliases",
81            Field::Predicate => "predicate",
82            Field::Tags => "tags",
83            Field::Location => "location",
84            Field::Statement => "statement",
85        }
86    }
87}
88
89/// Where an indexed memory came from.
90///
91/// Overlay facts are things the user said moments ago and the engine has not
92/// yet committed. They are retrievable immediately and ranked above canonical
93/// memory, but presented more cautiously.
94#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
95#[serde(rename_all = "snake_case")]
96pub enum MemoryOrigin {
97    /// Reconciled and durable.
98    Canonical,
99    /// Learned in the current session and not yet committed.
100    SessionOverlay,
101}
102
103/// A memory reduced to what the index needs.
104#[derive(Debug, Clone)]
105pub struct IndexedMemory {
106    /// Record identity.
107    pub id: MemoryId,
108    /// Memory kind, for scope filtering.
109    pub kind: MemoryKind,
110    /// Lifecycle state.
111    pub status: MemoryStatus,
112    /// Canonical predicate, for per-predicate diversity limits.
113    pub predicate: CanonicalPredicate,
114    /// Aggregated confidence.
115    pub confidence: f32,
116    /// Whether the record rests on something the user said outright.
117    pub explicit: bool,
118    /// Where the record came from.
119    pub origin: MemoryOrigin,
120    /// Expected persistence.
121    pub temporal_scope: TemporalScope,
122    /// When the fact started holding.
123    pub valid_from: DateTime<Utc>,
124    /// When it stops being retrievable, if ever.
125    pub expires_at: Option<DateTime<Utc>>,
126    /// The sentence handed to the model.
127    pub statement: String,
128    /// The value side of the fact, for filling governed state slots.
129    ///
130    /// A slot wants "pescatarian", not "The user is pescatarian." — the
131    /// statement is for the model, the value is for the application.
132    pub value: String,
133    /// Normalized subject surface form, for exact-entity boosting.
134    pub subject_form: String,
135    /// Normalized entity surface forms, for exact-entity boosting.
136    pub entity_forms: Vec<String>,
137    /// Tokenized field contents.
138    pub fields: [Vec<String>; 7],
139}
140
141impl IndexedMemory {
142    /// Project a canonical record into the index.
143    pub fn from_canonical(memory: &CanonicalMemory) -> Self {
144        let mut fields: [Vec<String>; 7] = Default::default();
145        fields[Field::Subject.slot()] = tokenize(&memory.retrieval.subject);
146        fields[Field::Entities.slot()] = memory
147            .retrieval
148            .entities
149            .iter()
150            .flat_map(|e| tokenize(e))
151            .collect();
152        fields[Field::Aliases.slot()] = memory
153            .retrieval
154            .aliases
155            .iter()
156            .chain(memory.subject.aliases.iter())
157            .flat_map(|a| tokenize(a))
158            .collect();
159        fields[Field::Predicate.slot()] = tokenize(memory.predicate.as_str());
160        fields[Field::Tags.slot()] = memory
161            .retrieval
162            .tags
163            .iter()
164            .flat_map(|t| tokenize(t))
165            .collect();
166        fields[Field::Location.slot()] = memory
167            .retrieval
168            .location
169            .as_deref()
170            .map(tokenize)
171            .unwrap_or_default();
172        fields[Field::Statement.slot()] = tokenize(&memory.statement);
173
174        let mut entity_forms: Vec<String> = memory
175            .retrieval
176            .entities
177            .iter()
178            .map(|e| normalize_token(e))
179            .collect();
180        entity_forms.extend(memory.subject.surface_forms());
181        entity_forms.retain(|f| !f.is_empty());
182        entity_forms.sort();
183        entity_forms.dedup();
184
185        Self {
186            id: memory.id.clone(),
187            kind: memory.kind,
188            status: memory.status,
189            predicate: memory.predicate.clone(),
190            confidence: memory.confidence,
191            explicit: memory.source.is_explicit(),
192            origin: MemoryOrigin::Canonical,
193            temporal_scope: memory.temporal_scope,
194            valid_from: memory.temporal.valid_from,
195            expires_at: memory.temporal.expires_at,
196            statement: memory.statement.clone(),
197            value: memory.value.display(),
198            subject_form: normalize_token(&memory.subject.display),
199            entity_forms,
200            fields,
201        }
202    }
203
204    /// Mark this document as an uncommitted session fact.
205    pub fn as_session_overlay(mut self) -> Self {
206        self.origin = MemoryOrigin::SessionOverlay;
207        self
208    }
209
210    /// Token count in a field.
211    pub fn field_len(&self, field: Field) -> usize {
212        self.fields[field.slot()].len()
213    }
214
215    /// Whether the document may be returned at `now`.
216    pub fn is_retrievable(&self, now: DateTime<Utc>) -> bool {
217        let status_ok = match self.origin {
218            // Overlay facts are staged by definition; they are still usable.
219            MemoryOrigin::SessionOverlay => self.status != MemoryStatus::Deleted,
220            MemoryOrigin::Canonical => self.status == MemoryStatus::Active,
221        };
222        status_ok && self.expires_at.is_none_or(|e| e > now)
223    }
224}
225
226/// Split text into normalized, indexable terms.
227///
228/// Lowercase, split on anything non-alphanumeric, fold regular plurals, and
229/// drop stop words. Deliberately not a stemmer: a personal corpus is full of
230/// names, and aggressive stemming conflates them ("Rhea" and "rhe") for very
231/// little recall. Plural folding is the one exception, because "restaurants"
232/// and "restaurant" are the same word to every user who says either.
233///
234/// Paraphrases with no shared term at all ("diet" against a record indexed
235/// under "dietary") are deliberately out of scope here — that is what record
236/// aliases and the semantic fallback exist for.
237pub fn tokenize(text: &str) -> Vec<String> {
238    text.split(|c: char| !c.is_alphanumeric())
239        .filter(|t| !t.is_empty())
240        .map(|t| singularize(&t.to_lowercase()))
241        .filter(|t| !is_stop_word(t))
242        .collect()
243}
244
245/// Fold regular English plurals onto their singular form.
246fn singularize(token: &str) -> String {
247    if token.len() <= 3 {
248        return token.to_string();
249    }
250    if let Some(stem) = token.strip_suffix("ies") {
251        return format!("{stem}y");
252    }
253    for suffix in ["ches", "shes", "sses", "xes", "zes"] {
254        if let Some(stem) = token.strip_suffix("es")
255            && token.ends_with(suffix)
256        {
257            return stem.to_string();
258        }
259    }
260    if token.ends_with('s')
261        && !token.ends_with("ss")
262        && !token.ends_with("us")
263        && !token.ends_with("is")
264        // "does" and "goes" are verbs, not plurals; folding them to "doe" and
265        // "goe" invents words that then collide with real ones.
266        && !token.ends_with("oes")
267    {
268        return token[..token.len() - 1].to_string();
269    }
270    token.to_string()
271}
272
273/// Words carrying no retrieval signal in a personal-memory corpus.
274fn is_stop_word(token: &str) -> bool {
275    const STOP: &[&str] = &[
276        "a", "an", "and", "are", "as", "at", "be", "but", "by", "do", "does", "for", "from", "had",
277        "has", "have", "in", "is", "it", "its", "of", "on", "or", "that", "the", "to", "was",
278        "were", "will", "with",
279    ];
280    STOP.contains(&token)
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286    use crate::core::{
287        EntityRef, EvidenceCounters, Explicitness, MemorySource, MemoryValue, PrivacyMetadata,
288        RetrievalMetadata, SessionId, TemporalMetadata, TurnId, UserId,
289    };
290
291    fn canonical() -> CanonicalMemory {
292        CanonicalMemory {
293            id: MemoryId::new("mem_1"),
294            owner: UserId::new("usr_1"),
295            kind: MemoryKind::RelationshipPreference,
296            predicate: CanonicalPredicate::new("venue_preference"),
297            status: MemoryStatus::Active,
298            confidence: 0.9,
299            subject: EntityRef::named("Rhea").with_alias("my wife"),
300            value: MemoryValue::Text("quiet restaurants".into()),
301            statement: "Rhea prefers quiet restaurants.".into(),
302            evidence_summary: "stated".into(),
303            source: MemorySource::from_explicitness(
304                Explicitness::ExplicitStatement,
305                SessionId::new("ses_1"),
306                TurnId(3),
307            ),
308            temporal: TemporalMetadata::created_at(Utc::now()),
309            retrieval: RetrievalMetadata {
310                subject: "rhea".into(),
311                tags: vec!["restaurant".into(), "noise".into()],
312                aliases: vec!["wife".into()],
313                entities: vec!["Rhea".into()],
314                location: Some("Bandra".into()),
315            },
316            evidence: EvidenceCounters::first(),
317            privacy: PrivacyMetadata::default(),
318            temporal_scope: TemporalScope::Persistent,
319            supersedes: Vec::new(),
320            superseded_by: None,
321            qualifier: None,
322        }
323    }
324
325    #[test]
326    fn tokenizer_normalizes_and_drops_stop_words() {
327        assert_eq!(
328            tokenize("The user is a Pescatarian!"),
329            vec!["user", "pescatarian"]
330        );
331        assert!(tokenize("   ").is_empty());
332    }
333
334    #[test]
335    fn regular_plurals_fold_onto_their_singular() {
336        assert_eq!(tokenize("restaurants"), tokenize("restaurant"));
337        assert_eq!(tokenize("preferences"), tokenize("preference"));
338        assert_eq!(tokenize("allergies"), vec!["allergy"]);
339        assert_eq!(tokenize("lunches"), vec!["lunch"]);
340    }
341
342    #[test]
343    fn plural_folding_leaves_verbs_that_merely_end_in_s() {
344        assert_eq!(tokenize("does"), vec!["doe"; 0], "does is a stop word");
345        assert_eq!(super::singularize("does"), "does");
346        assert_eq!(super::singularize("goes"), "goes");
347    }
348
349    #[test]
350    fn plural_folding_leaves_names_and_short_words_alone() {
351        assert_eq!(tokenize("Rhea"), vec!["rhea"]);
352        assert_eq!(tokenize("Kushal"), vec!["kushal"]);
353        assert_eq!(tokenize("gas"), vec!["gas"]);
354        assert_eq!(tokenize("this"), vec!["this"]);
355    }
356
357    #[test]
358    fn every_field_is_populated_from_the_record() {
359        let doc = IndexedMemory::from_canonical(&canonical());
360        assert_eq!(doc.fields[Field::Subject.slot()], vec!["rhea"]);
361        assert_eq!(doc.fields[Field::Tags.slot()], vec!["restaurant", "noise"]);
362        assert!(doc.fields[Field::Aliases.slot()].contains(&"wife".to_string()));
363        assert_eq!(doc.fields[Field::Location.slot()], vec!["bandra"]);
364        assert!(doc.fields[Field::Statement.slot()].contains(&"quiet".to_string()));
365        assert!(doc.entity_forms.contains(&"rhea".to_string()));
366        assert!(doc.entity_forms.contains(&"my wife".to_string()));
367    }
368
369    #[test]
370    fn the_value_is_kept_separately_from_the_sentence() {
371        let doc = IndexedMemory::from_canonical(&canonical());
372        assert_eq!(doc.value, "quiet restaurants");
373        assert_eq!(doc.statement, "Rhea prefers quiet restaurants.");
374    }
375
376    #[test]
377    fn subject_and_entity_fields_outweigh_the_statement() {
378        assert!(Field::Subject.weight() > Field::Statement.weight());
379        assert!(Field::Aliases.weight() > Field::Tags.weight());
380    }
381
382    #[test]
383    fn expired_and_superseded_documents_are_not_retrievable() {
384        let now = Utc::now();
385        let mut doc = IndexedMemory::from_canonical(&canonical());
386        assert!(doc.is_retrievable(now));
387
388        doc.expires_at = Some(now - chrono::Duration::hours(1));
389        assert!(!doc.is_retrievable(now));
390
391        let mut superseded = IndexedMemory::from_canonical(&canonical());
392        superseded.status = MemoryStatus::Superseded;
393        assert!(!superseded.is_retrievable(now));
394    }
395
396    #[test]
397    fn overlay_facts_stay_retrievable_while_staged() {
398        let mut doc = IndexedMemory::from_canonical(&canonical()).as_session_overlay();
399        doc.status = MemoryStatus::Staged;
400        assert!(doc.is_retrievable(Utc::now()));
401    }
402}
gemini_memory_rs/bm25/schema.rs

gemini_memory_rs/bm25/
schema.rs
