mentedb_cognitive/
trajectory.rs

1use mentedb_core::types::Timestamp;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::io;
5use std::path::Path;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub enum DecisionState {
9    Investigating,
10    NarrowedTo(String),
11    Decided(String),
12    Interrupted,
13    Completed,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct TrajectoryNode {
18    pub turn_id: u64,
19    pub topic_embedding: Vec<f32>,
20    pub topic_summary: String,
21    pub decision_state: DecisionState,
22    pub open_questions: Vec<String>,
23    pub timestamp: Timestamp,
24}
25
26const MAX_TURNS_DEFAULT: usize = 100;
27const REINFORCEMENT_BONUS: u32 = 2;
28
29/// Basic topic normalization: lowercase, collapse whitespace, trim.
30/// This handles the easy cases (casing, extra spaces) without attempting
31/// semantic canonicalization (tracked in #22).
32fn normalize_topic(raw: &str) -> String {
33    raw.split_whitespace()
34        .map(|w| w.to_lowercase())
35        .collect::<Vec<_>>()
36        .join(" ")
37}
38
39/// Tracks topic transitions as a Markov chain. Maps
40/// from_topic -> (to_topic -> frequency_count).
41#[derive(Debug, Clone, Default, Serialize, Deserialize)]
42pub struct TransitionMap {
43    transitions: HashMap<String, HashMap<String, u32>>,
44}
45
46#[derive(Serialize, Deserialize)]
47struct TransitionSnapshot {
48    version: u32,
49    transitions: HashMap<String, HashMap<String, u32>>,
50}
51
52const TRANSITION_SNAPSHOT_VERSION: u32 = 1;
53
54impl TransitionMap {
55    pub fn record(&mut self, from: &str, to: &str) {
56        let from = normalize_topic(from);
57        let to = normalize_topic(to);
58        *self
59            .transitions
60            .entry(from)
61            .or_default()
62            .entry(to)
63            .or_insert(0) += 1;
64    }
65
66    pub fn reinforce(&mut self, from: &str, to: &str) {
67        let from = normalize_topic(from);
68        let to = normalize_topic(to);
69        *self
70            .transitions
71            .entry(from)
72            .or_default()
73            .entry(to)
74            .or_insert(0) += REINFORCEMENT_BONUS;
75    }
76
77    pub fn decay(&mut self, from: &str, to: &str) {
78        let from = normalize_topic(from);
79        let to = normalize_topic(to);
80        if let Some(targets) = self.transitions.get_mut(&from) {
81            if let Some(count) = targets.get_mut(&to) {
82                *count = count.saturating_sub(1);
83                if *count == 0 {
84                    targets.remove(&to);
85                }
86            }
87            if targets.is_empty() {
88                self.transitions.remove(&from);
89            }
90        }
91    }
92
93    /// Returns the top N predicted topics from a given topic,
94    /// sorted by frequency descending.
95    pub fn predict_from(&self, topic: &str, limit: usize) -> Vec<(String, u32)> {
96        let topic = normalize_topic(topic);
97        let Some(targets) = self.transitions.get(&topic) else {
98            return Vec::new();
99        };
100        let mut ranked: Vec<(String, u32)> = targets.iter().map(|(t, &c)| (t.clone(), c)).collect();
101        ranked.sort_by(|a, b| b.1.cmp(&a.1));
102        ranked.truncate(limit);
103        ranked
104    }
105
106    pub fn is_empty(&self) -> bool {
107        self.transitions.is_empty()
108    }
109
110    pub fn total_transitions(&self) -> usize {
111        self.transitions.values().map(|t| t.len()).sum()
112    }
113
114    /// Save the transition map to a JSON file. Prunes transitions with
115    /// count below `min_count` to keep the file from growing unbounded.
116    /// Uses atomic write (temp file + rename) to avoid corruption.
117    pub fn save(&self, path: &Path, min_count: u32) -> io::Result<()> {
118        let pruned: HashMap<String, HashMap<String, u32>> = self
119            .transitions
120            .iter()
121            .filter_map(|(from, targets)| {
122                let kept: HashMap<String, u32> = targets
123                    .iter()
124                    .filter(|(_, c)| **c >= min_count)
125                    .map(|(t, &c)| (t.clone(), c))
126                    .collect();
127                if kept.is_empty() {
128                    None
129                } else {
130                    Some((from.clone(), kept))
131                }
132            })
133            .collect();
134        let snapshot = TransitionSnapshot {
135            version: TRANSITION_SNAPSHOT_VERSION,
136            transitions: pruned,
137        };
138        let json = serde_json::to_string(&snapshot)
139            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
140        let tmp = path.with_extension("tmp");
141        std::fs::write(&tmp, json)?;
142        std::fs::rename(&tmp, path)
143    }
144
145    /// Load a transition map from a JSON file, merging counts into the
146    /// current map so that patterns accumulate across sessions.
147    pub fn load(&mut self, path: &Path) -> io::Result<()> {
148        let json = std::fs::read_to_string(path)?;
149        let snapshot: TransitionSnapshot = serde_json::from_str(&json)
150            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
151
152        if snapshot.version != TRANSITION_SNAPSHOT_VERSION {
153            return Err(io::Error::new(
154                io::ErrorKind::InvalidData,
155                format!(
156                    "unsupported transition snapshot version: {} (expected {})",
157                    snapshot.version, TRANSITION_SNAPSHOT_VERSION
158                ),
159            ));
160        }
161
162        for (from, targets) in snapshot.transitions {
163            let entry = self.transitions.entry(from).or_default();
164            for (to, count) in targets {
165                *entry.entry(to).or_insert(0) += count;
166            }
167        }
168        Ok(())
169    }
170}
171
172pub struct TrajectoryTracker {
173    trajectory: Vec<TrajectoryNode>,
174    max_turns: usize,
175    pub transitions: TransitionMap,
176}
177
178impl TrajectoryTracker {
179    pub fn new(max_turns: usize) -> Self {
180        Self {
181            trajectory: Vec::new(),
182            max_turns,
183            transitions: TransitionMap::default(),
184        }
185    }
186
187    pub fn record_turn(&mut self, turn: TrajectoryNode) {
188        if let Some(prev) = self.trajectory.last() {
189            self.transitions
190                .record(&prev.topic_summary, &turn.topic_summary);
191        }
192
193        if self.trajectory.len() >= self.max_turns {
194            self.trajectory.remove(0);
195        }
196        self.trajectory.push(turn);
197    }
198
199    pub fn get_trajectory(&self) -> &[TrajectoryNode] {
200        &self.trajectory
201    }
202
203    pub fn get_resume_context(&self) -> Option<String> {
204        if self.trajectory.is_empty() {
205            return None;
206        }
207
208        let mut parts = Vec::new();
209
210        // Find the last non-completed topic
211        if let Some(last) = self.trajectory.last() {
212            parts.push(format!("You were working on: {}", last.topic_summary));
213
214            match &last.decision_state {
215                DecisionState::Investigating => {
216                    parts.push("Status: Still investigating.".to_string());
217                }
218                DecisionState::NarrowedTo(choice) => {
219                    parts.push(format!("You narrowed down to: {}", choice));
220                }
221                DecisionState::Decided(decision) => {
222                    parts.push(format!("You decided on: {}", decision));
223                }
224                DecisionState::Interrupted => {
225                    parts.push("Status: Was interrupted before completion.".to_string());
226                }
227                DecisionState::Completed => {
228                    parts.push("Status: Completed.".to_string());
229                }
230            }
231
232            if !last.open_questions.is_empty() {
233                let qs: Vec<String> = last
234                    .open_questions
235                    .iter()
236                    .map(|q| format!("- {}", q))
237                    .collect();
238                parts.push(format!("Open questions:\n{}", qs.join("\n")));
239            }
240        }
241
242        // Add recent trajectory summary
243        if self.trajectory.len() > 1 {
244            let recent: Vec<String> = self
245                .trajectory
246                .iter()
247                .rev()
248                .skip(1)
249                .take(3)
250                .rev()
251                .map(|t| t.topic_summary.clone())
252                .collect();
253            parts.push(format!("Recent trajectory: {}", recent.join(" → ")));
254        }
255
256        Some(parts.join(" "))
257    }
258
259    pub fn predict_next_topics(&self) -> Vec<String> {
260        let mut predictions = Vec::new();
261        let mut seen = ahash::AHashSet::new();
262
263        let Some(last) = self.trajectory.last() else {
264            return predictions;
265        };
266
267        // Learned transitions are the strongest signal
268        let learned = self.transitions.predict_from(&last.topic_summary, 3);
269        for (topic, _count) in &learned {
270            if seen.insert(topic.clone()) {
271                predictions.push(topic.clone());
272            }
273        }
274
275        // Open questions fill remaining slots
276        for q in &last.open_questions {
277            if predictions.len() >= 3 {
278                break;
279            }
280            if seen.insert(q.clone()) {
281                predictions.push(q.clone());
282            }
283        }
284
285        // Continuation of current topic
286        if predictions.len() < 3 {
287            let cont = format!("{} (continued)", last.topic_summary);
288            if seen.insert(cont.clone()) {
289                predictions.push(cont);
290            }
291        }
292
293        // Revisit previous topic
294        if predictions.len() < 3 && self.trajectory.len() >= 2 {
295            let prev = &self.trajectory[self.trajectory.len() - 2];
296            let rev = format!("{} (revisit)", prev.topic_summary);
297            if seen.insert(rev.clone()) {
298                predictions.push(rev);
299            }
300        }
301
302        predictions.truncate(3);
303        predictions
304    }
305
306    /// Called when the speculative cache gets a hit. Reinforces the
307    /// transition from the previous topic to the hit topic.
308    pub fn reinforce_transition(&mut self, hit_topic: &str) {
309        if let Some(last) = self.trajectory.last() {
310            self.transitions.reinforce(&last.topic_summary, hit_topic);
311        }
312    }
313
314    /// Called when the speculative cache misses. Slightly decays the
315    /// transition from the previous topic to the predicted topic.
316    pub fn decay_transition(&mut self, predicted_topic: &str) {
317        if let Some(last) = self.trajectory.last() {
318            self.transitions.decay(&last.topic_summary, predicted_topic);
319        }
320    }
321}
322
323impl Default for TrajectoryTracker {
324    fn default() -> Self {
325        Self::new(MAX_TURNS_DEFAULT)
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332
333    fn make_turn(
334        id: u64,
335        summary: &str,
336        state: DecisionState,
337        questions: Vec<&str>,
338    ) -> TrajectoryNode {
339        TrajectoryNode {
340            turn_id: id,
341            topic_embedding: vec![0.0; 4],
342            topic_summary: summary.to_string(),
343            decision_state: state,
344            open_questions: questions.into_iter().map(String::from).collect(),
345            timestamp: id * 1000,
346        }
347    }
348
349    #[test]
350    fn test_record_and_resume() {
351        let mut tracker = TrajectoryTracker::default();
352        tracker.record_turn(make_turn(
353            1,
354            "JWT auth design",
355            DecisionState::Investigating,
356            vec![],
357        ));
358        tracker.record_turn(make_turn(
359            2,
360            "Token refresh strategy",
361            DecisionState::Decided("short-lived access tokens (15min)".into()),
362            vec!["Where to store refresh tokens?"],
363        ));
364
365        let ctx = tracker.get_resume_context().unwrap();
366        assert!(ctx.contains("Token refresh strategy"));
367        assert!(ctx.contains("short-lived access tokens"));
368        assert!(ctx.contains("refresh tokens"));
369    }
370
371    #[test]
372    fn test_predict_topics() {
373        let mut tracker = TrajectoryTracker::default();
374        tracker.record_turn(make_turn(
375            1,
376            "Database schema",
377            DecisionState::Decided("normalized".into()),
378            vec!["How to handle migrations?", "Index strategy?"],
379        ));
380
381        let preds = tracker.predict_next_topics();
382        assert!(!preds.is_empty());
383        assert!(preds.iter().any(|p| p.contains("migrations")));
384    }
385
386    #[test]
387    fn test_fifo_eviction() {
388        let mut tracker = TrajectoryTracker::default();
389        for i in 0..105 {
390            tracker.record_turn(make_turn(
391                i,
392                &format!("turn {}", i),
393                DecisionState::Investigating,
394                vec![],
395            ));
396        }
397        assert_eq!(tracker.get_trajectory().len(), MAX_TURNS_DEFAULT);
398        assert_eq!(tracker.get_trajectory()[0].turn_id, 5);
399    }
400
401    #[test]
402    fn test_transition_recording() {
403        let mut tracker = TrajectoryTracker::default();
404        tracker.record_turn(make_turn(1, "auth", DecisionState::Investigating, vec![]));
405        tracker.record_turn(make_turn(
406            2,
407            "database",
408            DecisionState::Investigating,
409            vec![],
410        ));
411        tracker.record_turn(make_turn(3, "auth", DecisionState::Investigating, vec![]));
412        tracker.record_turn(make_turn(
413            4,
414            "database",
415            DecisionState::Investigating,
416            vec![],
417        ));
418        tracker.record_turn(make_turn(5, "auth", DecisionState::Investigating, vec![]));
419        tracker.record_turn(make_turn(
420            6,
421            "deployment",
422            DecisionState::Investigating,
423            vec![],
424        ));
425
426        // auth -> database happened twice, auth -> deployment once
427        let preds = tracker.transitions.predict_from("auth", 5);
428        assert_eq!(preds.len(), 2);
429        assert_eq!(preds[0].0, "database");
430        assert_eq!(preds[0].1, 2);
431        assert_eq!(preds[1].0, "deployment");
432        assert_eq!(preds[1].1, 1);
433    }
434
435    #[test]
436    fn test_learned_predictions_take_priority() {
437        let mut tracker = TrajectoryTracker::default();
438
439        // Build a pattern: auth -> database (3 times)
440        for _ in 0..3 {
441            tracker.record_turn(make_turn(0, "auth", DecisionState::Investigating, vec![]));
442            tracker.record_turn(make_turn(
443                0,
444                "database",
445                DecisionState::Investigating,
446                vec![],
447            ));
448        }
449
450        // Now land on auth with an open question
451        tracker.record_turn(make_turn(
452            0,
453            "auth",
454            DecisionState::Investigating,
455            vec!["how to handle JWT expiry?"],
456        ));
457
458        let preds = tracker.predict_next_topics();
459        // Learned transition "database" should come first
460        assert_eq!(preds[0], "database");
461    }
462
463    #[test]
464    fn test_reinforce_and_decay() {
465        let mut map = TransitionMap::default();
466        map.record("auth", "database");
467        map.record("auth", "database");
468        assert_eq!(map.predict_from("auth", 1)[0].1, 2);
469
470        // Reinforce adds bonus
471        map.reinforce("auth", "database");
472        assert_eq!(map.predict_from("auth", 1)[0].1, 4);
473
474        // Decay subtracts 1
475        map.decay("auth", "database");
476        assert_eq!(map.predict_from("auth", 1)[0].1, 3);
477    }
478
479    #[test]
480    fn test_decay_removes_zero_entries() {
481        let mut map = TransitionMap::default();
482        map.record("auth", "database");
483        assert_eq!(map.total_transitions(), 1);
484
485        map.decay("auth", "database");
486        assert!(map.is_empty());
487    }
488
489    #[test]
490    fn test_reinforce_via_tracker() {
491        let mut tracker = TrajectoryTracker::default();
492        tracker.record_turn(make_turn(1, "auth", DecisionState::Investigating, vec![]));
493        tracker.record_turn(make_turn(
494            2,
495            "database",
496            DecisionState::Investigating,
497            vec![],
498        ));
499
500        // One natural transition recorded
501        assert_eq!(tracker.transitions.predict_from("auth", 1)[0].1, 1);
502
503        // Simulate cache hit reinforcement
504        tracker.reinforce_transition("database");
505        assert_eq!(
506            tracker.transitions.predict_from("database", 1)[0].1,
507            REINFORCEMENT_BONUS
508        );
509    }
510
511    #[test]
512    fn test_no_duplicate_predictions() {
513        let mut tracker = TrajectoryTracker::default();
514
515        // Build pattern: auth -> database
516        tracker.record_turn(make_turn(1, "auth", DecisionState::Investigating, vec![]));
517        tracker.record_turn(make_turn(
518            2,
519            "database",
520            DecisionState::Investigating,
521            vec![],
522        ));
523
524        // Land on auth with "database" as an open question too
525        tracker.record_turn(make_turn(
526            3,
527            "auth",
528            DecisionState::Investigating,
529            vec!["database"],
530        ));
531
532        let preds = tracker.predict_next_topics();
533        let unique: ahash::AHashSet<&String> = preds.iter().collect();
534        assert_eq!(preds.len(), unique.len(), "predictions should be unique");
535    }
536
537    #[test]
538    fn test_normalization_collapses_variants() {
539        let mut map = TransitionMap::default();
540        map.record("Auth Setup", "database");
541        map.record("auth setup", "DATABASE");
542        map.record("  auth   setup  ", "  database  ");
543
544        // All three should collapse into one transition with count 3
545        let preds = map.predict_from("AUTH SETUP", 1);
546        assert_eq!(preds.len(), 1);
547        assert_eq!(preds[0].0, "database");
548        assert_eq!(preds[0].1, 3);
549    }
550
551    #[test]
552    fn test_transition_map_save_and_load() {
553        let dir = tempfile::tempdir().unwrap();
554        let path = dir.path().join("transitions.json");
555
556        let mut map = TransitionMap::default();
557        map.record("auth", "database");
558        map.record("auth", "database");
559        map.record("auth", "deploy");
560        map.save(&path, 1).unwrap();
561
562        // Load into a fresh map — counts should carry over
563        let mut loaded = TransitionMap::default();
564        loaded.load(&path).unwrap();
565        let preds = loaded.predict_from("auth", 5);
566        assert_eq!(preds[0].0, "database");
567        assert_eq!(preds[0].1, 2);
568        // deploy has count 1, should be saved with min_count=1
569        assert_eq!(preds[1].0, "deploy");
570        assert_eq!(preds[1].1, 1);
571    }
572
573    #[test]
574    fn test_transition_map_save_prunes_low_counts() {
575        let dir = tempfile::tempdir().unwrap();
576        let path = dir.path().join("transitions.json");
577
578        let mut map = TransitionMap::default();
579        map.record("auth", "database");
580        map.record("auth", "database");
581        map.record("auth", "deploy"); // count 1
582        map.save(&path, 2).unwrap(); // only keep count >= 2
583
584        let mut loaded = TransitionMap::default();
585        loaded.load(&path).unwrap();
586        let preds = loaded.predict_from("auth", 5);
587        assert_eq!(preds.len(), 1);
588        assert_eq!(preds[0].0, "database");
589        assert_eq!(preds[0].1, 2);
590    }
591
592    #[test]
593    fn test_transition_map_load_merges() {
594        let dir = tempfile::tempdir().unwrap();
595        let path = dir.path().join("transitions.json");
596
597        let mut map = TransitionMap::default();
598        map.record("auth", "database");
599        map.save(&path, 1).unwrap();
600
601        // Load into a map that already has data — counts should add
602        let mut existing = TransitionMap::default();
603        existing.record("auth", "database");
604        existing.record("auth", "testing");
605        existing.load(&path).unwrap();
606
607        let preds = existing.predict_from("auth", 5);
608        // database: 1 existing + 1 loaded = 2
609        assert_eq!(preds[0].0, "database");
610        assert_eq!(preds[0].1, 2);
611        // testing: 1 existing only
612        assert_eq!(preds[1].0, "testing");
613        assert_eq!(preds[1].1, 1);
614    }
615}
mentedb_cognitive/trajectory.rs

mentedb_cognitive/
trajectory.rs
