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    /// The current cache entries (predicted topic, pre-assembled context,
179    /// source memories, hit count, age), for introspection and display.
180    pub fn entries(&self) -> &[CacheEntry] {
181        &self.entries
182    }
183
184    fn evict_lru(&mut self) {
185        if self.entries.is_empty() {
186            return;
187        }
188        let lru_idx = self
189            .entries
190            .iter()
191            .enumerate()
192            .min_by_key(|(_, e)| e.last_accessed)
193            .map(|(i, _)| i)
194            .unwrap();
195        self.entries.remove(lru_idx);
196        self.stats.evictions += 1;
197    }
198}
199
200impl Default for SpeculativeCache {
201    fn default() -> Self {
202        Self::new(10, 0.5, 0.4)
203    }
204}
205
206/// On-disk format for the speculative cache. Only entries that have been
207/// hit at least once are worth persisting.
208#[derive(Serialize, Deserialize)]
209struct CacheSnapshot {
210    version: u32,
211    entries: Vec<CacheEntry>,
212    stats: CacheStats,
213}
214
215const CACHE_SNAPSHOT_VERSION: u32 = 1;
216
217impl SpeculativeCache {
218    /// Save cache entries with at least `min_hits` to a JSON file.
219    /// Entries below the threshold are considered stale and dropped.
220    /// Uses atomic write (temp file + rename) to avoid corruption.
221    pub fn save(&self, path: &Path, min_hits: u32) -> io::Result<()> {
222        let entries: Vec<CacheEntry> = self
223            .entries
224            .iter()
225            .filter(|e| e.hit_count >= min_hits)
226            .cloned()
227            .collect();
228        let snapshot = CacheSnapshot {
229            version: CACHE_SNAPSHOT_VERSION,
230            entries,
231            stats: self.stats.clone(),
232        };
233        let json = serde_json::to_string(&snapshot)
234            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
235        let tmp = path.with_extension("tmp");
236        std::fs::write(&tmp, json)?;
237        std::fs::rename(&tmp, path)
238    }
239
240    /// Load cache entries from a JSON file, merging into the current cache.
241    /// Preserves the configured max_size and thresholds. Resets last_accessed
242    /// timestamps so LRU eviction works correctly after reload.
243    pub fn load(&mut self, path: &Path) -> io::Result<()> {
244        let json = std::fs::read_to_string(path)?;
245        let snapshot: CacheSnapshot = serde_json::from_str(&json)
246            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
247
248        if snapshot.version != CACHE_SNAPSHOT_VERSION {
249            return Err(io::Error::new(
250                io::ErrorKind::InvalidData,
251                format!(
252                    "unsupported cache snapshot version: {} (expected {})",
253                    snapshot.version, CACHE_SNAPSHOT_VERSION
254                ),
255            ));
256        }
257
258        let now = std::time::SystemTime::now()
259            .duration_since(std::time::UNIX_EPOCH)
260            .unwrap_or_default()
261            .as_micros() as u64;
262
263        for mut entry in snapshot.entries {
264            if self.entries.len() >= self.max_size {
265                self.evict_lru();
266            }
267            if !self.entries.iter().any(|e| e.topic == entry.topic) {
268                entry.last_accessed = now;
269                self.entries.push(entry);
270            }
271        }
272        self.stats.hits += snapshot.stats.hits;
273        self.stats.misses += snapshot.stats.misses;
274        self.stats.evictions += snapshot.stats.evictions;
275        Ok(())
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    #[test]
284    fn test_pre_assemble_and_hit() {
285        let mut cache = SpeculativeCache::default();
286        cache.pre_assemble(
287            vec![
288                "database schema design".to_string(),
289                "API authentication".to_string(),
290            ],
291            |topic| {
292                Some((
293                    format!("Context for {}", topic),
294                    vec![MemoryId::new()],
295                    None,
296                ))
297            },
298        );
299
300        assert_eq!(cache.stats().cache_size, 2);
301
302        let hit = cache.try_hit("database schema", None);
303        assert!(hit.is_some());
304        assert!(hit.unwrap().context_text.contains("database schema design"));
305    }
306
307    #[test]
308    fn test_cache_miss() {
309        let mut cache = SpeculativeCache::default();
310        cache.pre_assemble(vec!["database schema".to_string()], |topic| {
311            Some((format!("Context for {}", topic), vec![], None))
312        });
313
314        let hit = cache.try_hit("cooking recipes", None);
315        assert!(hit.is_none());
316        assert_eq!(cache.stats().misses, 1);
317    }
318
319    #[test]
320    fn test_embedding_hit() {
321        let mut cache = SpeculativeCache::default();
322        cache.pre_assemble(vec!["rust ownership".to_string()], |_topic| {
323            Some((
324                "Context about ownership".to_string(),
325                vec![],
326                Some(vec![1.0, 0.0, 0.0, 0.0]),
327            ))
328        });
329
330        // Query with similar embedding but different words
331        let query_emb = vec![0.95, 0.1, 0.0, 0.0];
332        let hit = cache.try_hit("memory safety borrow checker", Some(&query_emb));
333        assert!(hit.is_some());
334    }
335
336    #[test]
337    fn test_embedding_miss() {
338        let mut cache = SpeculativeCache::default();
339        cache.pre_assemble(vec!["rust ownership".to_string()], |_topic| {
340            Some((
341                "Context about ownership".to_string(),
342                vec![],
343                Some(vec![1.0, 0.0, 0.0, 0.0]),
344            ))
345        });
346
347        // Orthogonal embedding should miss
348        let query_emb = vec![0.0, 0.0, 0.0, 1.0];
349        let hit = cache.try_hit("cooking recipes", Some(&query_emb));
350        assert!(hit.is_none());
351    }
352
353    #[test]
354    fn test_lru_eviction() {
355        let mut cache = SpeculativeCache::new(10, 0.5, 0.4);
356        for i in 0..12 {
357            cache.pre_assemble(vec![format!("topic {}", i)], |topic| {
358                Some((format!("Context for {}", topic), vec![], None))
359            });
360        }
361        assert!(cache.stats().cache_size <= 10);
362        assert!(cache.stats().evictions > 0);
363    }
364
365    #[test]
366    fn test_evict_stale() {
367        let mut cache = SpeculativeCache::default();
368        cache.pre_assemble(vec!["old topic".to_string()], |topic| {
369            Some((format!("Context for {}", topic), vec![], None))
370        });
371        cache.evict_stale(0, u64::MAX);
372        assert_eq!(cache.stats().cache_size, 0);
373    }
374
375    #[test]
376    fn test_save_and_load() {
377        let dir = tempfile::tempdir().unwrap();
378        let path = dir.path().join("cache.json");
379
380        let mut cache = SpeculativeCache::default();
381        cache.pre_assemble(
382            vec!["hot topic".to_string(), "cold topic".to_string()],
383            |topic| {
384                Some((
385                    format!("Context for {}", topic),
386                    vec![MemoryId::new()],
387                    None,
388                ))
389            },
390        );
391        // Hit the hot topic so it has hit_count >= 1
392        cache.try_hit("hot topic", None);
393
394        // Save with min_hits=1 — only "hot topic" should persist
395        cache.save(&path, 1).unwrap();
396
397        let mut loaded = SpeculativeCache::default();
398        loaded.load(&path).unwrap();
399        assert_eq!(loaded.stats().cache_size, 1);
400
401        let hit = loaded.try_hit("hot topic", None);
402        assert!(hit.is_some());
403        let miss = loaded.try_hit("cold topic", None);
404        assert!(miss.is_none());
405    }
406
407    #[test]
408    fn test_save_empty_cache() {
409        let dir = tempfile::tempdir().unwrap();
410        let path = dir.path().join("empty.json");
411
412        let cache = SpeculativeCache::default();
413        cache.save(&path, 1).unwrap();
414
415        let mut loaded = SpeculativeCache::default();
416        loaded.load(&path).unwrap();
417        assert_eq!(loaded.stats().cache_size, 0);
418    }
419
420    #[test]
421    fn test_load_missing_file() {
422        let mut cache = SpeculativeCache::default();
423        let result = cache.load(Path::new("/nonexistent/cache.json"));
424        assert!(result.is_err());
425    }
426}
mentedb_cognitive/speculative.rs

mentedb_cognitive/
speculative.rs
