mentedb_cognitive/
speculative.rs

1use mentedb_core::types::{MemoryId, Timestamp};
2use serde::{Deserialize, Serialize};
3use std::io;
4use std::path::Path;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct CacheEntry {
8    pub topic: String,
9    pub topic_embedding: Option<Vec<f32>>,
10    pub context_text: String,
11    pub memory_ids: Vec<MemoryId>,
12    pub created_at: Timestamp,
13    pub hit_count: u32,
14    last_accessed: Timestamp,
15}
16
17#[derive(Debug, Clone, Default, Serialize, Deserialize)]
18pub struct CacheStats {
19    pub hits: u64,
20    pub misses: u64,
21    pub evictions: u64,
22    pub cache_size: usize,
23}
24
25pub struct SpeculativeCache {
26    entries: Vec<CacheEntry>,
27    stats: CacheStats,
28    max_size: usize,
29    keyword_threshold: f32,
30    embedding_threshold: f32,
31}
32
33fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
34    if a.len() != b.len() || a.is_empty() {
35        return 0.0;
36    }
37    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
38    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
39    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
40    if norm_a == 0.0 || norm_b == 0.0 {
41        return 0.0;
42    }
43    dot / (norm_a * norm_b)
44}
45
46fn keyword_overlap_score(query: &str, topic: &str) -> f32 {
47    let query_words: ahash::AHashSet<String> = query
48        .to_lowercase()
49        .split_whitespace()
50        .filter(|w| w.len() >= 2)
51        .map(String::from)
52        .collect();
53    let topic_words: ahash::AHashSet<String> = topic
54        .to_lowercase()
55        .split_whitespace()
56        .filter(|w| w.len() >= 2)
57        .map(String::from)
58        .collect();
59
60    if query_words.is_empty() || topic_words.is_empty() {
61        return 0.0;
62    }
63
64    let intersection = query_words.intersection(&topic_words).count();
65    let union = query_words.union(&topic_words).count();
66    if union == 0 {
67        0.0
68    } else {
69        intersection as f32 / union as f32
70    }
71}
72
73impl SpeculativeCache {
74    pub fn new(max_size: usize, keyword_threshold: f32, embedding_threshold: f32) -> Self {
75        Self {
76            entries: Vec::new(),
77            stats: CacheStats::default(),
78            max_size,
79            keyword_threshold,
80            embedding_threshold,
81        }
82    }
83
84    pub fn pre_assemble(
85        &mut self,
86        predictions: Vec<String>,
87        builder: impl Fn(&str) -> Option<(String, Vec<MemoryId>, Option<Vec<f32>>)>,
88    ) {
89        let now = std::time::SystemTime::now()
90            .duration_since(std::time::UNIX_EPOCH)
91            .unwrap_or_default()
92            .as_micros() as u64;
93
94        for topic in predictions {
95            if self.entries.iter().any(|e| e.topic == topic) {
96                continue;
97            }
98
99            if let Some((context_text, memory_ids, embedding)) = builder(&topic) {
100                if self.entries.len() >= self.max_size {
101                    self.evict_lru();
102                }
103
104                self.entries.push(CacheEntry {
105                    topic,
106                    topic_embedding: embedding,
107                    context_text,
108                    memory_ids,
109                    created_at: now,
110                    hit_count: 0,
111                    last_accessed: now,
112                });
113            }
114        }
115    }
116
117    /// Try to find a cached context for this query. Uses cosine similarity on
118    /// embeddings when available, falls back to keyword overlap only when
119    /// no embeddings exist.
120    pub fn try_hit(&mut self, query: &str, query_embedding: Option<&[f32]>) -> Option<CacheEntry> {
121        let now = std::time::SystemTime::now()
122            .duration_since(std::time::UNIX_EPOCH)
123            .unwrap_or_default()
124            .as_micros() as u64;
125
126        let mut best_idx = None;
127        let mut best_score = 0.0f32;
128        let mut used_embeddings = false;
129
130        for (i, entry) in self.entries.iter().enumerate() {
131            let score = match (query_embedding, &entry.topic_embedding) {
132                (Some(qe), Some(te)) => {
133                    used_embeddings = true;
134                    cosine_similarity(qe, te)
135                }
136                _ => keyword_overlap_score(query, &entry.topic),
137            };
138
139            if score > best_score {
140                best_score = score;
141                best_idx = Some(i);
142            }
143        }
144
145        let threshold = if used_embeddings {
146            self.embedding_threshold
147        } else {
148            self.keyword_threshold
149        };
150
151        if best_score > threshold
152            && let Some(idx) = best_idx
153        {
154            self.entries[idx].hit_count += 1;
155            self.entries[idx].last_accessed = now;
156            self.stats.hits += 1;
157            return Some(self.entries[idx].clone());
158        }
159
160        self.stats.misses += 1;
161        None
162    }
163
164    pub fn evict_stale(&mut self, max_age_us: u64, now: Timestamp) {
165        let before = self.entries.len();
166        self.entries.retain(|e| now - e.created_at <= max_age_us);
167        let evicted = before - self.entries.len();
168        self.stats.evictions += evicted as u64;
169    }
170
171    pub fn stats(&self) -> CacheStats {
172        CacheStats {
173            cache_size: self.entries.len(),
174            ..self.stats.clone()
175        }
176    }
177
178    fn evict_lru(&mut self) {
179        if self.entries.is_empty() {
180            return;
181        }
182        let lru_idx = self
183            .entries
184            .iter()
185            .enumerate()
186            .min_by_key(|(_, e)| e.last_accessed)
187            .map(|(i, _)| i)
188            .unwrap();
189        self.entries.remove(lru_idx);
190        self.stats.evictions += 1;
191    }
192}
193
194impl Default for SpeculativeCache {
195    fn default() -> Self {
196        Self::new(10, 0.5, 0.4)
197    }
198}
199
200/// On-disk format for the speculative cache. Only entries that have been
201/// hit at least once are worth persisting.
202#[derive(Serialize, Deserialize)]
203struct CacheSnapshot {
204    version: u32,
205    entries: Vec<CacheEntry>,
206    stats: CacheStats,
207}
208
209const CACHE_SNAPSHOT_VERSION: u32 = 1;
210
211impl SpeculativeCache {
212    /// Save cache entries with at least `min_hits` to a JSON file.
213    /// Entries below the threshold are considered stale and dropped.
214    /// Uses atomic write (temp file + rename) to avoid corruption.
215    pub fn save(&self, path: &Path, min_hits: u32) -> io::Result<()> {
216        let entries: Vec<CacheEntry> = self
217            .entries
218            .iter()
219            .filter(|e| e.hit_count >= min_hits)
220            .cloned()
221            .collect();
222        let snapshot = CacheSnapshot {
223            version: CACHE_SNAPSHOT_VERSION,
224            entries,
225            stats: self.stats.clone(),
226        };
227        let json = serde_json::to_string(&snapshot)
228            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
229        let tmp = path.with_extension("tmp");
230        std::fs::write(&tmp, json)?;
231        std::fs::rename(&tmp, path)
232    }
233
234    /// Load cache entries from a JSON file, merging into the current cache.
235    /// Preserves the configured max_size and thresholds. Resets last_accessed
236    /// timestamps so LRU eviction works correctly after reload.
237    pub fn load(&mut self, path: &Path) -> io::Result<()> {
238        let json = std::fs::read_to_string(path)?;
239        let snapshot: CacheSnapshot = serde_json::from_str(&json)
240            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
241
242        if snapshot.version != CACHE_SNAPSHOT_VERSION {
243            return Err(io::Error::new(
244                io::ErrorKind::InvalidData,
245                format!(
246                    "unsupported cache snapshot version: {} (expected {})",
247                    snapshot.version, CACHE_SNAPSHOT_VERSION
248                ),
249            ));
250        }
251
252        let now = std::time::SystemTime::now()
253            .duration_since(std::time::UNIX_EPOCH)
254            .unwrap_or_default()
255            .as_micros() as u64;
256
257        for mut entry in snapshot.entries {
258            if self.entries.len() >= self.max_size {
259                self.evict_lru();
260            }
261            if !self.entries.iter().any(|e| e.topic == entry.topic) {
262                entry.last_accessed = now;
263                self.entries.push(entry);
264            }
265        }
266        self.stats.hits += snapshot.stats.hits;
267        self.stats.misses += snapshot.stats.misses;
268        self.stats.evictions += snapshot.stats.evictions;
269        Ok(())
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn test_pre_assemble_and_hit() {
279        let mut cache = SpeculativeCache::default();
280        cache.pre_assemble(
281            vec![
282                "database schema design".to_string(),
283                "API authentication".to_string(),
284            ],
285            |topic| {
286                Some((
287                    format!("Context for {}", topic),
288                    vec![MemoryId::new()],
289                    None,
290                ))
291            },
292        );
293
294        assert_eq!(cache.stats().cache_size, 2);
295
296        let hit = cache.try_hit("database schema", None);
297        assert!(hit.is_some());
298        assert!(hit.unwrap().context_text.contains("database schema design"));
299    }
300
301    #[test]
302    fn test_cache_miss() {
303        let mut cache = SpeculativeCache::default();
304        cache.pre_assemble(vec!["database schema".to_string()], |topic| {
305            Some((format!("Context for {}", topic), vec![], None))
306        });
307
308        let hit = cache.try_hit("cooking recipes", None);
309        assert!(hit.is_none());
310        assert_eq!(cache.stats().misses, 1);
311    }
312
313    #[test]
314    fn test_embedding_hit() {
315        let mut cache = SpeculativeCache::default();
316        cache.pre_assemble(vec!["rust ownership".to_string()], |_topic| {
317            Some((
318                "Context about ownership".to_string(),
319                vec![],
320                Some(vec![1.0, 0.0, 0.0, 0.0]),
321            ))
322        });
323
324        // Query with similar embedding but different words
325        let query_emb = vec![0.95, 0.1, 0.0, 0.0];
326        let hit = cache.try_hit("memory safety borrow checker", Some(&query_emb));
327        assert!(hit.is_some());
328    }
329
330    #[test]
331    fn test_embedding_miss() {
332        let mut cache = SpeculativeCache::default();
333        cache.pre_assemble(vec!["rust ownership".to_string()], |_topic| {
334            Some((
335                "Context about ownership".to_string(),
336                vec![],
337                Some(vec![1.0, 0.0, 0.0, 0.0]),
338            ))
339        });
340
341        // Orthogonal embedding should miss
342        let query_emb = vec![0.0, 0.0, 0.0, 1.0];
343        let hit = cache.try_hit("cooking recipes", Some(&query_emb));
344        assert!(hit.is_none());
345    }
346
347    #[test]
348    fn test_lru_eviction() {
349        let mut cache = SpeculativeCache::new(10, 0.5, 0.4);
350        for i in 0..12 {
351            cache.pre_assemble(vec![format!("topic {}", i)], |topic| {
352                Some((format!("Context for {}", topic), vec![], None))
353            });
354        }
355        assert!(cache.stats().cache_size <= 10);
356        assert!(cache.stats().evictions > 0);
357    }
358
359    #[test]
360    fn test_evict_stale() {
361        let mut cache = SpeculativeCache::default();
362        cache.pre_assemble(vec!["old topic".to_string()], |topic| {
363            Some((format!("Context for {}", topic), vec![], None))
364        });
365        cache.evict_stale(0, u64::MAX);
366        assert_eq!(cache.stats().cache_size, 0);
367    }
368
369    #[test]
370    fn test_save_and_load() {
371        let dir = tempfile::tempdir().unwrap();
372        let path = dir.path().join("cache.json");
373
374        let mut cache = SpeculativeCache::default();
375        cache.pre_assemble(
376            vec!["hot topic".to_string(), "cold topic".to_string()],
377            |topic| {
378                Some((
379                    format!("Context for {}", topic),
380                    vec![MemoryId::new()],
381                    None,
382                ))
383            },
384        );
385        // Hit the hot topic so it has hit_count >= 1
386        cache.try_hit("hot topic", None);
387
388        // Save with min_hits=1 — only "hot topic" should persist
389        cache.save(&path, 1).unwrap();
390
391        let mut loaded = SpeculativeCache::default();
392        loaded.load(&path).unwrap();
393        assert_eq!(loaded.stats().cache_size, 1);
394
395        let hit = loaded.try_hit("hot topic", None);
396        assert!(hit.is_some());
397        let miss = loaded.try_hit("cold topic", None);
398        assert!(miss.is_none());
399    }
400
401    #[test]
402    fn test_save_empty_cache() {
403        let dir = tempfile::tempdir().unwrap();
404        let path = dir.path().join("empty.json");
405
406        let cache = SpeculativeCache::default();
407        cache.save(&path, 1).unwrap();
408
409        let mut loaded = SpeculativeCache::default();
410        loaded.load(&path).unwrap();
411        assert_eq!(loaded.stats().cache_size, 0);
412    }
413
414    #[test]
415    fn test_load_missing_file() {
416        let mut cache = SpeculativeCache::default();
417        let result = cache.load(Path::new("/nonexistent/cache.json"));
418        assert!(result.is_err());
419    }
420}
mentedb_cognitive/speculative.rs

mentedb_cognitive/
speculative.rs
