vtcode_core/context/
entity_resolver.rs

1//! Entity resolution for vibe coding support
2//!
3//! This module provides fuzzy entity matching to resolve vague terms like
4//! "the sidebar" or "that button" to actual workspace entities (files, components, etc.)
5
6use hashbrown::HashMap;
7use serde::{Deserialize, Serialize};
8use std::collections::VecDeque;
9use std::path::PathBuf;
10use vtcode_commons::utils::current_timestamp;
11
12/// Maximum number of entity matches to return
13const MAX_ENTITY_MATCHES: usize = 5;
14
15/// Maximum number of recent edits to track
16const MAX_RECENT_EDITS: usize = 50;
17
18/// Location of an entity within a file
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct FileLocation {
21    pub path: PathBuf,
22    pub line_start: usize,
23    pub line_end: usize,
24    pub content_preview: String,
25}
26
27/// A matched entity with confidence score
28#[derive(Debug, Clone)]
29pub struct EntityMatch {
30    pub entity: String,
31    pub locations: Vec<FileLocation>,
32    pub confidence: f32,
33    pub recency_score: f32,
34    pub mention_score: f32,
35    pub proximity_score: f32,
36}
37
38impl EntityMatch {
39    /// Calculate total score for ranking
40    pub fn total_score(&self) -> f32 {
41        self.confidence + self.recency_score + self.mention_score + self.proximity_score
42    }
43}
44
45/// Reference to an entity that was recently edited
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct EntityReference {
48    pub entity: String,
49    pub file: PathBuf,
50    pub timestamp: u64,
51}
52
53/// Value found in a style file (CSS, SCSS, etc.)
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct StyleValue {
56    pub property: String,
57    pub value: String,
58    pub file: PathBuf,
59    pub line: usize,
60}
61
62/// Index of entities in the workspace
63#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct EntityIndex {
65    /// UI components (Sidebar, Button, etc.)
66    pub ui_components: HashMap<String, Vec<FileLocation>>,
67
68    /// Functions and methods
69    pub functions: HashMap<String, Vec<FileLocation>>,
70
71    /// Classes and structs
72    pub classes: HashMap<String, Vec<FileLocation>>,
73
74    /// Style properties (padding, color, etc.)
75    pub style_properties: HashMap<String, Vec<StyleValue>>,
76
77    /// Config keys
78    pub config_keys: HashMap<String, Vec<FileLocation>>,
79
80    /// Recent edits for recency ranking
81    pub recent_edits: VecDeque<EntityReference>,
82
83    /// Last mentioned entities with timestamps
84    pub last_mentioned: HashMap<String, u64>,
85
86    /// Last update timestamp
87    pub last_updated: u64,
88}
89
90impl Default for EntityIndex {
91    fn default() -> Self {
92        Self {
93            ui_components: HashMap::new(),
94            functions: HashMap::new(),
95            classes: HashMap::new(),
96            style_properties: HashMap::new(),
97            config_keys: HashMap::new(),
98            recent_edits: VecDeque::with_capacity(MAX_RECENT_EDITS),
99            last_mentioned: HashMap::new(),
100            last_updated: current_timestamp(),
101        }
102    }
103}
104
105/// Entity resolver for fuzzy matching
106pub struct EntityResolver {
107    /// The entity index
108    index: EntityIndex,
109}
110
111impl Default for EntityResolver {
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117impl EntityResolver {
118    /// Create a new entity resolver
119    pub fn new() -> Self {
120        Self { index: EntityIndex::default() }
121    }
122
123    /// Check if the entity index is empty
124    pub fn index_is_empty(&self) -> bool {
125        self.index.ui_components.is_empty() && self.index.functions.is_empty() && self.index.classes.is_empty()
126    }
127
128    /// Resolve a vague term to entity matches
129    pub fn resolve(&self, term: &str) -> Option<EntityMatch> {
130        let matches = self.find_entity_fuzzy(term);
131
132        // Return best match if any
133        matches.into_iter().max_by(|a, b| {
134            a.total_score()
135                .partial_cmp(&b.total_score())
136                .unwrap_or(std::cmp::Ordering::Equal)
137        })
138    }
139
140    /// Find entities using fuzzy matching
141    fn find_entity_fuzzy(&self, term: &str) -> Vec<EntityMatch> {
142        let term_lower = term.to_lowercase();
143        let mut matches = Vec::new();
144
145        // Search UI components
146        self.search_hashmap(&self.index.ui_components, &term_lower, &mut matches);
147
148        // Search functions
149        self.search_hashmap(&self.index.functions, &term_lower, &mut matches);
150
151        // Search classes
152        self.search_hashmap(&self.index.classes, &term_lower, &mut matches);
153
154        // Sort by total score and limit
155        matches.sort_by(|a, b| {
156            b.total_score()
157                .partial_cmp(&a.total_score())
158                .unwrap_or(std::cmp::Ordering::Equal)
159        });
160        matches.truncate(MAX_ENTITY_MATCHES);
161
162        matches
163    }
164
165    /// Search a hashmap for matching entities
166    fn search_hashmap(&self, map: &HashMap<String, Vec<FileLocation>>, term: &str, matches: &mut Vec<EntityMatch>) {
167        for (entity, locations) in map {
168            let entity_lower = entity.to_lowercase();
169
170            // Exact match
171            if entity_lower == term {
172                matches.push(EntityMatch {
173                    entity: entity.clone(),
174                    locations: locations.clone(),
175                    confidence: 1.0,
176                    recency_score: self.calculate_recency_score(&entity_lower),
177                    mention_score: self.calculate_mention_score(&entity_lower),
178                    proximity_score: 0.0,
179                });
180                continue;
181            }
182
183            // Case-insensitive substring match
184            if entity_lower.contains(term) {
185                let confidence = term.len() as f32 / entity_lower.len() as f32;
186                matches.push(EntityMatch {
187                    entity: entity.clone(),
188                    locations: locations.clone(),
189                    confidence: confidence * 0.8, // Slightly lower than exact
190                    recency_score: self.calculate_recency_score(&entity_lower),
191                    mention_score: self.calculate_mention_score(&entity_lower),
192                    proximity_score: 0.0,
193                });
194                continue;
195            }
196
197            // Fuzzy match using Damerau-Levenshtein distance
198            let distance = levenshtein_distance(term, &entity_lower);
199            if distance <= 2 {
200                let confidence = strsim::normalized_damerau_levenshtein(term, &entity_lower) as f32;
201                matches.push(EntityMatch {
202                    entity: entity.clone(),
203                    locations: locations.clone(),
204                    confidence: confidence * 0.6, // Lower confidence for fuzzy
205                    recency_score: self.calculate_recency_score(&entity_lower),
206                    mention_score: self.calculate_mention_score(&entity_lower),
207                    proximity_score: 0.0,
208                });
209            }
210        }
211    }
212
213    /// Calculate recency score based on recent edits
214    fn calculate_recency_score(&self, entity: &str) -> f32 {
215        let now = current_timestamp();
216
217        // Check if entity was recently edited (within 5 minutes)
218        if let Some(edit) = self.index.recent_edits.iter().rev().find(|e| e.entity.to_lowercase() == entity) {
219            let age_seconds = now.saturating_sub(edit.timestamp);
220            if age_seconds < 300 {
221                // Score decays over 5 minutes
222                return 0.3 * (1.0 - (age_seconds as f32 / 300.0));
223            }
224        }
225
226        0.0
227    }
228
229    /// Calculate mention score based on conversation history
230    fn calculate_mention_score(&self, entity: &str) -> f32 {
231        if let Some(&timestamp) = self.index.last_mentioned.get(entity) {
232            let now = current_timestamp();
233            let age_seconds = now.saturating_sub(timestamp);
234
235            // Score decays over 10 minutes
236            if age_seconds < 600 {
237                return 0.2 * (1.0 - (age_seconds as f32 / 600.0));
238            }
239        }
240
241        0.0
242    }
243
244    /// Get mutable access to the index for building
245    pub fn index_mut(&mut self) -> &mut EntityIndex {
246        &mut self.index
247    }
248}
249
250/// Calculate Damerau-Levenshtein distance between two strings.
251///
252/// Like Levenshtein but also handles transpositions, which improves entity
253/// matching for real-world typos (e.g. "hte" ↔ "the").
254/// Delegates to the battle-tested [`strsim`](https://docs.rs/strsim) implementation.
255fn levenshtein_distance(a: &str, b: &str) -> usize {
256    strsim::damerau_levenshtein(a, b)
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[test]
264    fn test_levenshtein_distance() {
265        assert_eq!(levenshtein_distance("", ""), 0);
266        assert_eq!(levenshtein_distance("hello", "hello"), 0);
267        assert_eq!(levenshtein_distance("hello", "hallo"), 1);
268        assert_eq!(levenshtein_distance("sidebar", "sidbar"), 1);
269        assert_eq!(levenshtein_distance("button", "btn"), 3);
270    }
271
272    #[test]
273    fn test_entity_match_scoring() {
274        let match1 = EntityMatch {
275            entity: "Sidebar".to_string(),
276            locations: vec![],
277            confidence: 1.0,
278            recency_score: 0.3,
279            mention_score: 0.2,
280            proximity_score: 0.0,
281        };
282
283        assert!((match1.total_score() - 1.5).abs() < f32::EPSILON);
284    }
285
286    #[tokio::test]
287    async fn test_entity_resolver_exact_match() {
288        let mut resolver = EntityResolver::new();
289
290        resolver.index_mut().ui_components.insert(
291            "Sidebar".to_string(),
292            vec![FileLocation {
293                path: PathBuf::from("src/Sidebar.tsx"),
294                line_start: 1,
295                line_end: 50,
296                content_preview: "export const Sidebar = () => {}".to_string(),
297            }],
298        );
299
300        let result = resolver.resolve("sidebar");
301        assert!(result.is_some());
302
303        let matched = result.unwrap();
304        assert_eq!(matched.entity, "Sidebar");
305        assert!((matched.confidence - 1.0).abs() < f32::EPSILON);
306    }
307}
vtcode_core/context/entity_resolver.rs

vtcode_core/context/
entity_resolver.rs
