vtcode_core/context/
conversation_memory.rs

1//! Conversation memory for vibe coding support
2//!
3//! Tracks entities mentioned across conversation turns to enable pronoun resolution
4//! and contextual understanding.
5
6use hashbrown::HashMap;
7use serde::{Deserialize, Serialize};
8use std::collections::VecDeque;
9use std::fmt::Write;
10use std::path::PathBuf;
11use vtcode_commons::utils::current_timestamp;
12
13/// Maximum number of conversation turns to remember
14const MAX_MEMORY_TURNS: usize = 50;
15
16/// Maximum entity mentions to track
17const MAX_ENTITY_MENTIONS: usize = 200;
18
19/// Type of entity mention
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
21pub enum MentionType {
22    /// Direct mention (explicit naming)
23    Direct,
24    /// Pronoun reference (it, that, this)
25    Pronoun,
26    /// Implicit (inferred from context)
27    Implicit,
28}
29
30/// History of an entity's mentions
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct MentionHistory {
33    pub entity: String,
34    pub first_mention: u64,
35    pub last_mention: u64,
36    pub mention_count: usize,
37    pub context_snippets: Vec<String>,
38}
39
40impl MentionHistory {
41    /// Create new mention history
42    pub fn new(entity: String) -> Self {
43        let now = current_timestamp();
44        Self {
45            entity,
46            first_mention: now,
47            last_mention: now,
48            mention_count: 1,
49            context_snippets: Vec::new(),
50        }
51    }
52
53    /// Record a mention
54    pub fn record_mention(&mut self, _turn: usize) {
55        self.last_mention = current_timestamp();
56        self.mention_count += 1;
57    }
58}
59
60/// A single entity mention in the timeline
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct EntityMention {
63    pub turn: usize,
64    pub entity: String,
65    pub mention_type: MentionType,
66    pub file_context: Option<PathBuf>,
67}
68
69/// A user message for context
70#[derive(Debug, Clone, Serialize, Deserialize)]
71pub struct UserMessage {
72    pub turn: usize,
73    pub content: String,
74    pub entities: Vec<String>,
75}
76
77/// Pronoun reference waiting to be resolved
78#[derive(Debug, Clone)]
79pub struct PronounReference {
80    pub pronoun: String,
81    pub turn: usize,
82    pub context: String,
83}
84
85/// Conversation memory system
86pub struct ConversationMemory {
87    /// Mentioned entities with history
88    mentioned_entities: HashMap<String, MentionHistory>,
89
90    /// Timeline of entity mentions
91    entity_timeline: VecDeque<EntityMention>,
92
93    /// Recent user messages
94    recent_user_messages: VecDeque<UserMessage>,
95
96    /// Recent file contexts (files mentioned or edited)
97
98    /// Resolved reference mappings
99
100    /// Current turn number
101    current_turn: usize,
102}
103
104impl Default for ConversationMemory {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110impl ConversationMemory {
111    /// Create a new conversation memory
112    pub fn new() -> Self {
113        Self {
114            mentioned_entities: HashMap::new(),
115            entity_timeline: VecDeque::with_capacity(MAX_ENTITY_MENTIONS),
116            recent_user_messages: VecDeque::with_capacity(MAX_MEMORY_TURNS),
117            current_turn: 0,
118        }
119    }
120
121    /// Extract entities from a message
122    pub fn extract_entities(&mut self, message: &str, turn: usize) {
123        self.current_turn = turn;
124
125        let entities = self.extract_nouns_and_identifiers(message);
126        let mut extracted = Vec::new();
127
128        for entity in entities {
129            self.record_entity_mention(&entity, turn, MentionType::Direct);
130            extracted.push(entity);
131        }
132
133        // Store user message
134        self.recent_user_messages.push_back(UserMessage {
135            turn,
136            content: message.to_string(),
137            entities: extracted,
138        });
139
140        // Keep bounded
141        while self.recent_user_messages.len() > MAX_MEMORY_TURNS {
142            self.recent_user_messages.pop_front();
143        }
144    }
145
146    /// Extract nouns and identifiers from text
147    fn extract_nouns_and_identifiers(&self, text: &str) -> Vec<String> {
148        let mut entities = Vec::new();
149
150        const ENTITY_STOPWORDS: &[&str] = &[
151            "update",
152            "fix",
153            "test",
154            "look",
155            "add",
156            "remove",
157            "create",
158            "delete",
159            "refactor",
160            "implement",
161            "check",
162            "review",
163        ];
164
165        // Simple extraction: capitalize words, camelCase, PascalCase
166        for word in text.split_whitespace() {
167            let cleaned = word.trim_matches(|c: char| !c.is_alphanumeric() && c != '.');
168            let candidate = cleaned.split('.').next().unwrap_or(cleaned);
169            let candidate_lower = candidate.to_ascii_lowercase();
170
171            // Skip empty and short words
172            if candidate.len() < 3 {
173                continue;
174            }
175
176            // Skip common action words that are not entities.
177            if ENTITY_STOPWORDS.contains(&candidate_lower.as_str()) {
178                continue;
179            }
180
181            // Capitalized words (likely proper nouns)
182            if candidate.chars().next().map(|c| c.is_uppercase()).unwrap_or(false) {
183                entities.push(candidate.to_string());
184                continue;
185            }
186
187            // camelCase or PascalCase (likely identifiers)
188            let has_mixed_case =
189                candidate.chars().any(|c| c.is_uppercase()) && candidate.chars().any(|c| c.is_lowercase());
190
191            if has_mixed_case {
192                entities.push(candidate.to_string());
193            }
194        }
195
196        entities
197    }
198
199    /// Record an entity mention
200    fn record_entity_mention(&mut self, entity: &str, turn: usize, mention_type: MentionType) {
201        let entity_lower = entity.to_lowercase();
202
203        // Update or create history
204        self.mentioned_entities
205            .entry(entity_lower.clone())
206            .and_modify(|history| history.record_mention(turn))
207            .or_insert_with(|| MentionHistory::new(entity.to_string()));
208
209        // Add to timeline
210        self.entity_timeline.push_back(EntityMention {
211            turn,
212            entity: entity.to_string(),
213            mention_type,
214            file_context: None,
215        });
216
217        // Keep bounded
218        while self.entity_timeline.len() > MAX_ENTITY_MENTIONS {
219            self.entity_timeline.pop_front();
220        }
221    }
222
223    /// Get recent entities (N most recent)
224    pub fn get_recent_entities(&self, count: usize) -> Vec<String> {
225        self.entity_timeline
226            .iter()
227            .rev()
228            .take(count)
229            .map(|mention| mention.entity.clone())
230            .collect()
231    }
232
233    /// Resolve a pronoun to an entity
234    pub fn resolve_pronoun(&self, pronoun: &str, turn: usize) -> Option<String> {
235        let pronoun_lower = pronoun.to_lowercase();
236
237        match pronoun_lower.as_str() {
238            "it" => {
239                // Look back at last mentioned entity
240                self.entity_timeline
241                    .iter()
242                    .rev()
243                    .find(|m| m.turn < turn)
244                    .map(|m| m.entity.clone())
245            }
246            "that" | "this" => {
247                // Look at last direct mention
248                self.entity_timeline
249                    .iter()
250                    .rev()
251                    .filter(|m| m.turn < turn)
252                    .find(|m| matches!(m.mention_type, MentionType::Direct))
253                    .map(|m| m.entity.clone())
254            }
255            "those" | "these" => {
256                // Multiple entities - return most recent direct mention
257                self.entity_timeline
258                    .iter()
259                    .rev()
260                    .filter(|m| m.turn < turn && matches!(m.mention_type, MentionType::Direct))
261                    .take(2)
262                    .map(|m| m.entity.clone())
263                    .next()
264            }
265            _ => None,
266        }
267    }
268
269    /// Get entity mention count
270    pub fn mention_count(&self, entity: &str) -> usize {
271        self.mentioned_entities
272            .get(&entity.to_lowercase())
273            .map(|h| h.mention_count)
274            .unwrap_or(0)
275    }
276
277    /// Get context summary for recent conversation
278    pub fn get_context_summary(&self, turns: usize) -> String {
279        let messages: Vec<_> = self.recent_user_messages.iter().rev().take(turns).collect();
280
281        if messages.is_empty() {
282            return String::from("No recent context available");
283        }
284
285        let mut summary = String::from("Recent conversation:\n");
286        for msg in messages.iter().rev() {
287            let _ = writeln!(summary, "Turn {}: {}", msg.turn, msg.content);
288        }
289
290        summary
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn test_extract_entities() {
300        let mut memory = ConversationMemory::new();
301
302        memory.extract_entities("Update the Sidebar component in App.tsx", 1);
303
304        assert_eq!(memory.mentioned_entities.len(), 2); // Sidebar, App
305        assert!(memory.mentioned_entities.contains_key("sidebar"));
306        assert!(memory.mentioned_entities.contains_key("app"));
307    }
308
309    #[test]
310    fn test_pronoun_resolution_it() {
311        let mut memory = ConversationMemory::new();
312
313        // Turn 1: Mention Sidebar
314        memory.extract_entities("The Sidebar is too wide", 1);
315
316        // Turn 2: Reference with "it"
317        let resolved = memory.resolve_pronoun("it", 2);
318
319        assert!(resolved.is_some());
320        assert_eq!(resolved.unwrap(), "Sidebar");
321    }
322
323    #[test]
324    fn test_pronoun_resolution_that() {
325        let mut memory = ConversationMemory::new();
326
327        memory.extract_entities("Look at the Button component", 1);
328
329        let resolved = memory.resolve_pronoun("that", 2);
330
331        assert!(resolved.is_some());
332        assert_eq!(resolved.unwrap(), "Button");
333    }
334
335    #[test]
336    fn test_recent_entities() {
337        let mut memory = ConversationMemory::new();
338
339        memory.extract_entities("Update Sidebar", 1);
340        memory.extract_entities("Fix Button", 2);
341        memory.extract_entities("Test Form", 3);
342
343        let recent = memory.get_recent_entities(2);
344
345        assert_eq!(recent.len(), 2);
346        assert_eq!(recent[0], "Form");
347        assert_eq!(recent[1], "Button");
348    }
349
350    #[test]
351    fn test_mention_count() {
352        let mut memory = ConversationMemory::new();
353
354        memory.extract_entities("Update Sidebar", 1);
355        memory.extract_entities("The Sidebar is nice", 2);
356        memory.extract_entities("Sidebar needs work", 3);
357
358        assert_eq!(memory.mention_count("sidebar"), 3);
359        assert_eq!(memory.mention_count("Button"), 0);
360    }
361
362    #[test]
363    fn test_context_summary() {
364        let mut memory = ConversationMemory::new();
365
366        memory.extract_entities("First message", 1);
367        memory.extract_entities("Second message", 2);
368
369        let summary = memory.get_context_summary(2);
370
371        assert!(summary.contains("First message"));
372        assert!(summary.contains("Second message"));
373    }
374}
vtcode_core/context/conversation_memory.rs

vtcode_core/context/
conversation_memory.rs
