mentedb_cognitive/
trajectory.rs

1use mentedb_core::types::Timestamp;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Serialize, Deserialize)]
5pub enum DecisionState {
6    Investigating,
7    NarrowedTo(String),
8    Decided(String),
9    Interrupted,
10    Completed,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct TrajectoryNode {
15    pub turn_id: u64,
16    pub topic_embedding: Vec<f32>,
17    pub topic_summary: String,
18    pub decision_state: DecisionState,
19    pub open_questions: Vec<String>,
20    pub timestamp: Timestamp,
21}
22
23const MAX_TURNS_DEFAULT: usize = 100;
24
25pub struct TrajectoryTracker {
26    trajectory: Vec<TrajectoryNode>,
27    max_turns: usize,
28}
29
30impl TrajectoryTracker {
31    pub fn new(max_turns: usize) -> Self {
32        Self {
33            trajectory: Vec::new(),
34            max_turns,
35        }
36    }
37
38    pub fn record_turn(&mut self, turn: TrajectoryNode) {
39        if self.trajectory.len() >= self.max_turns {
40            self.trajectory.remove(0);
41        }
42        self.trajectory.push(turn);
43    }
44
45    pub fn get_trajectory(&self) -> &[TrajectoryNode] {
46        &self.trajectory
47    }
48
49    pub fn get_resume_context(&self) -> Option<String> {
50        if self.trajectory.is_empty() {
51            return None;
52        }
53
54        let mut parts = Vec::new();
55
56        // Find the last non-completed topic
57        if let Some(last) = self.trajectory.last() {
58            parts.push(format!("You were working on: {}", last.topic_summary));
59
60            match &last.decision_state {
61                DecisionState::Investigating => {
62                    parts.push("Status: Still investigating.".to_string());
63                }
64                DecisionState::NarrowedTo(choice) => {
65                    parts.push(format!("You narrowed down to: {}", choice));
66                }
67                DecisionState::Decided(decision) => {
68                    parts.push(format!("You decided on: {}", decision));
69                }
70                DecisionState::Interrupted => {
71                    parts.push("Status: Was interrupted before completion.".to_string());
72                }
73                DecisionState::Completed => {
74                    parts.push("Status: Completed.".to_string());
75                }
76            }
77
78            if !last.open_questions.is_empty() {
79                let qs: Vec<String> = last
80                    .open_questions
81                    .iter()
82                    .map(|q| format!("- {}", q))
83                    .collect();
84                parts.push(format!("Open questions:\n{}", qs.join("\n")));
85            }
86        }
87
88        // Add recent trajectory summary
89        if self.trajectory.len() > 1 {
90            let recent: Vec<String> = self
91                .trajectory
92                .iter()
93                .rev()
94                .skip(1)
95                .take(3)
96                .rev()
97                .map(|t| t.topic_summary.clone())
98                .collect();
99            parts.push(format!("Recent trajectory: {}", recent.join(" → ")));
100        }
101
102        Some(parts.join(" "))
103    }
104
105    pub fn predict_next_topics(&self) -> Vec<String> {
106        let mut predictions = Vec::new();
107
108        if let Some(last) = self.trajectory.last() {
109            // Open questions are the best predictors
110            for q in &last.open_questions {
111                predictions.push(q.clone());
112                if predictions.len() >= 3 {
113                    return predictions;
114                }
115            }
116
117            // Last topic as continuation
118            if predictions.len() < 3 {
119                predictions.push(format!("{} (continued)", last.topic_summary));
120            }
121
122            // Look at trajectory pattern for related topics
123            if predictions.len() < 3 && self.trajectory.len() >= 2 {
124                let prev = &self.trajectory[self.trajectory.len() - 2];
125                predictions.push(format!("{} (revisit)", prev.topic_summary));
126            }
127        }
128
129        predictions.truncate(3);
130        predictions
131    }
132}
133
134impl Default for TrajectoryTracker {
135    fn default() -> Self {
136        Self::new(MAX_TURNS_DEFAULT)
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    fn make_turn(
145        id: u64,
146        summary: &str,
147        state: DecisionState,
148        questions: Vec<&str>,
149    ) -> TrajectoryNode {
150        TrajectoryNode {
151            turn_id: id,
152            topic_embedding: vec![0.0; 4],
153            topic_summary: summary.to_string(),
154            decision_state: state,
155            open_questions: questions.into_iter().map(String::from).collect(),
156            timestamp: id * 1000,
157        }
158    }
159
160    #[test]
161    fn test_record_and_resume() {
162        let mut tracker = TrajectoryTracker::default();
163        tracker.record_turn(make_turn(
164            1,
165            "JWT auth design",
166            DecisionState::Investigating,
167            vec![],
168        ));
169        tracker.record_turn(make_turn(
170            2,
171            "Token refresh strategy",
172            DecisionState::Decided("short-lived access tokens (15min)".into()),
173            vec!["Where to store refresh tokens?"],
174        ));
175
176        let ctx = tracker.get_resume_context().unwrap();
177        assert!(ctx.contains("Token refresh strategy"));
178        assert!(ctx.contains("short-lived access tokens"));
179        assert!(ctx.contains("refresh tokens"));
180    }
181
182    #[test]
183    fn test_predict_topics() {
184        let mut tracker = TrajectoryTracker::default();
185        tracker.record_turn(make_turn(
186            1,
187            "Database schema",
188            DecisionState::Decided("normalized".into()),
189            vec!["How to handle migrations?", "Index strategy?"],
190        ));
191
192        let preds = tracker.predict_next_topics();
193        assert!(!preds.is_empty());
194        assert!(preds.iter().any(|p| p.contains("migrations")));
195    }
196
197    #[test]
198    fn test_fifo_eviction() {
199        let mut tracker = TrajectoryTracker::default();
200        for i in 0..105 {
201            tracker.record_turn(make_turn(
202                i,
203                &format!("turn {}", i),
204                DecisionState::Investigating,
205                vec![],
206            ));
207        }
208        assert_eq!(tracker.get_trajectory().len(), MAX_TURNS_DEFAULT);
209        assert_eq!(tracker.get_trajectory()[0].turn_id, 5);
210    }
211}
mentedb_cognitive/trajectory.rs

mentedb_cognitive/
trajectory.rs
