lean_ctx/core/
semantic_cache.rs

1use std::collections::{HashMap, HashSet};
2use std::path::PathBuf;
3
4use serde::{Deserialize, Serialize};
5
6/// Recompute global IDF after this many mutations; searches refresh IDF when `idf_dirty`.
7const IDF_REBUILD_BATCH: u32 = 100;
8/// Hard cap on cached entries. Matches the search cap so the index stays bounded and
9/// never grows past what `find_similar` will actually search (previously the structure
10/// grew one entry per distinct file ever read, with no eviction).
11const MAX_SEMANTIC_ENTRIES: usize = 200;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SemanticCacheEntry {
15    pub path: String,
16    pub tfidf_vector: Vec<(String, f64)>,
17    pub token_count: usize,
18    pub access_count: u32,
19    pub last_session: String,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize, Default)]
23pub struct SemanticCacheIndex {
24    pub entries: Vec<SemanticCacheEntry>,
25    pub idf: HashMap<String, f64>,
26    pub total_docs: usize,
27    /// Documents containing each term (unique terms per entry).
28    #[serde(default)]
29    pub term_document_freq: HashMap<String, usize>,
30    #[serde(default)]
31    idf_dirty: bool,
32    #[serde(default)]
33    mutations_since_idf_rebuild: u32,
34}
35
36impl SemanticCacheIndex {
37    pub fn add_file(&mut self, path: &str, content: &str, session_id: &str) {
38        let tf = compute_tf(content);
39        let token_count = content.split_whitespace().count();
40
41        if let Some(existing) = self.entries.iter_mut().find(|e| e.path == path) {
42            remove_doc_terms(&mut self.term_document_freq, &existing.tfidf_vector);
43            existing.tfidf_vector = tf.iter().map(|(k, v)| (k.clone(), *v)).collect();
44            existing.token_count = token_count;
45            existing.access_count += 1;
46            existing.last_session = session_id.to_string();
47            add_doc_terms(&mut self.term_document_freq, &existing.tfidf_vector);
48        } else {
49            if self.entries.len() >= MAX_SEMANTIC_ENTRIES {
50                // Evict the lowest-access_count entry to bound the index. Decrement its
51                // terms from the DF map via the existing helper so IDF stays correct.
52                if let Some(victim_idx) = self
53                    .entries
54                    .iter()
55                    .enumerate()
56                    .min_by_key(|(_, e)| e.access_count)
57                    .map(|(i, _)| i)
58                {
59                    let victim = self.entries.swap_remove(victim_idx);
60                    remove_doc_terms(&mut self.term_document_freq, &victim.tfidf_vector);
61                }
62            }
63            let tf_vec: Vec<(String, f64)> = tf.iter().map(|(k, v)| (k.clone(), *v)).collect();
64            add_doc_terms(&mut self.term_document_freq, &tf_vec);
65            self.entries.push(SemanticCacheEntry {
66                path: path.to_string(),
67                tfidf_vector: tf_vec,
68                token_count,
69                access_count: 1,
70                last_session: session_id.to_string(),
71            });
72        }
73
74        self.total_docs = self.entries.len();
75        self.note_idf_mutation();
76    }
77
78    fn note_idf_mutation(&mut self) {
79        self.idf_dirty = true;
80        self.mutations_since_idf_rebuild = self.mutations_since_idf_rebuild.saturating_add(1);
81        if self.mutations_since_idf_rebuild >= IDF_REBUILD_BATCH {
82            self.recompute_idf_from_df();
83            self.idf_dirty = false;
84            self.mutations_since_idf_rebuild = 0;
85        }
86    }
87
88    fn recompute_idf_from_df(&mut self) {
89        self.idf.clear();
90        let n = self.total_docs as f64;
91        if n <= 0.0 {
92            return;
93        }
94        for (term, count) in &self.term_document_freq {
95            let idf = (n / (*count as f64 + 1.0)).ln() + 1.0;
96            self.idf.insert(term.clone(), idf);
97        }
98    }
99
100    fn rebuild_df_from_entries(&mut self) {
101        self.term_document_freq.clear();
102        for entry in &self.entries {
103            add_doc_terms(&mut self.term_document_freq, &entry.tfidf_vector);
104        }
105    }
106
107    fn repair_after_deserialize(&mut self) {
108        self.total_docs = self.entries.len();
109        if self.term_document_freq.is_empty() && !self.entries.is_empty() {
110            self.rebuild_df_from_entries();
111            self.idf_dirty = true;
112        }
113    }
114
115    fn ensure_idf_for_search(&mut self) {
116        if self.idf_dirty {
117            self.recompute_idf_from_df();
118            self.idf_dirty = false;
119            self.mutations_since_idf_rebuild = 0;
120        }
121    }
122
123    pub fn find_similar(&mut self, content: &str, threshold: f64) -> Vec<(String, f64)> {
124        if self.entries.len() > MAX_SEMANTIC_ENTRIES {
125            return Vec::new();
126        }
127
128        self.ensure_idf_for_search();
129
130        let query_tf = compute_tf(content);
131        let query_vec = self.tfidf_vector(&query_tf);
132
133        let mut results: Vec<(String, f64)> = self
134            .entries
135            .iter()
136            .filter_map(|entry| {
137                let entry_vec = self.tfidf_vector_from_stored(&entry.tfidf_vector);
138                let sim = cosine_similarity(&query_vec, &entry_vec);
139                if sim >= threshold {
140                    Some((entry.path.clone(), sim))
141                } else {
142                    None
143                }
144            })
145            .collect();
146
147        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
148        results
149    }
150
151    pub fn suggest_warmup(&self, top_n: usize) -> Vec<String> {
152        let mut ranked: Vec<(&SemanticCacheEntry, f64)> = self
153            .entries
154            .iter()
155            .map(|e| {
156                let score = e.access_count as f64 * 0.6 + e.token_count as f64 * 0.0001;
157                (e, score)
158            })
159            .collect();
160
161        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
162
163        ranked
164            .into_iter()
165            .take(top_n)
166            .map(|(e, _)| e.path.clone())
167            .collect()
168    }
169
170    fn tfidf_vector(&self, tf: &HashMap<String, f64>) -> HashMap<String, f64> {
171        tf.iter()
172            .map(|(term, freq)| {
173                let idf = self.idf.get(term).copied().unwrap_or(1.0);
174                (term.clone(), freq * idf)
175            })
176            .collect()
177    }
178
179    fn tfidf_vector_from_stored(&self, stored: &[(String, f64)]) -> HashMap<String, f64> {
180        stored
181            .iter()
182            .map(|(term, freq)| {
183                let idf = self.idf.get(term).copied().unwrap_or(1.0);
184                (term.clone(), freq * idf)
185            })
186            .collect()
187    }
188
189    pub fn save(&self, project_root: &str) -> Result<(), String> {
190        let path = index_path(project_root);
191        if let Some(dir) = path.parent() {
192            std::fs::create_dir_all(dir).map_err(|e| e.to_string())?;
193        }
194        let json = serde_json::to_string(self).map_err(|e| e.to_string())?;
195        std::fs::write(&path, json).map_err(|e| e.to_string())
196    }
197
198    pub fn load(project_root: &str) -> Option<Self> {
199        let path = index_path(project_root);
200        let content = std::fs::read_to_string(&path)
201            .or_else(|_| {
202                let legacy = legacy_index_path(project_root);
203                if legacy == path {
204                    return Err(std::io::Error::new(
205                        std::io::ErrorKind::NotFound,
206                        "same path",
207                    ));
208                }
209                let data = std::fs::read_to_string(&legacy)?;
210                let _ = std::fs::copy(&legacy, &path);
211                Ok(data)
212            })
213            .ok()?;
214        let mut index: SemanticCacheIndex = serde_json::from_str(&content).ok()?;
215        index.repair_after_deserialize();
216        Some(index)
217    }
218
219    pub fn load_or_create(project_root: &str) -> Self {
220        Self::load(project_root).unwrap_or_default()
221    }
222}
223
224fn remove_doc_terms(df: &mut HashMap<String, usize>, tf_vec: &[(String, f64)]) {
225    let unique: HashSet<&str> = tf_vec.iter().map(|(k, _)| k.as_str()).collect();
226    for term in unique {
227        if let Some(c) = df.get_mut(term) {
228            *c = c.saturating_sub(1);
229            if *c == 0 {
230                df.remove(term);
231            }
232        }
233    }
234}
235
236fn add_doc_terms(df: &mut HashMap<String, usize>, tf_vec: &[(String, f64)]) {
237    let unique: HashSet<&str> = tf_vec.iter().map(|(k, _)| k.as_str()).collect();
238    for term in unique {
239        *df.entry(term.to_string()).or_default() += 1;
240    }
241}
242
243fn compute_tf(content: &str) -> HashMap<String, f64> {
244    let mut counts: HashMap<String, usize> = HashMap::new();
245    let mut total = 0usize;
246
247    for word in content.split(|c: char| !c.is_alphanumeric() && c != '_') {
248        let w = word.to_lowercase();
249        if w.len() >= 2 {
250            *counts.entry(w).or_default() += 1;
251            total += 1;
252        }
253    }
254
255    if total == 0 {
256        return HashMap::new();
257    }
258
259    counts
260        .into_iter()
261        .map(|(term, count)| (term, count as f64 / total as f64))
262        .collect()
263}
264
265fn cosine_similarity(a: &HashMap<String, f64>, b: &HashMap<String, f64>) -> f64 {
266    let mut dot = 0.0f64;
267    let mut norm_a = 0.0f64;
268    let mut norm_b = 0.0f64;
269
270    for (term, val) in a {
271        norm_a += val * val;
272        if let Some(bval) = b.get(term) {
273            dot += val * bval;
274        }
275    }
276    for val in b.values() {
277        norm_b += val * val;
278    }
279
280    let denom = norm_a.sqrt() * norm_b.sqrt();
281    if denom < 1e-10 {
282        return 0.0;
283    }
284    dot / denom
285}
286
287fn index_path(project_root: &str) -> PathBuf {
288    let hash = crate::core::project_hash::hash_project_root(project_root);
289    crate::core::data_dir::lean_ctx_data_dir()
290        .unwrap_or_default()
291        .join("semantic_cache")
292        .join(format!("{hash}.json"))
293}
294
295fn legacy_index_path(project_root: &str) -> PathBuf {
296    use md5::{Digest, Md5};
297    let mut hasher = Md5::new();
298    hasher.update(project_root.as_bytes());
299    let hash = format!("{:x}", hasher.finalize());
300    crate::core::data_dir::lean_ctx_data_dir()
301        .unwrap_or_default()
302        .join("semantic_cache")
303        .join(format!("{hash}.json"))
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309
310    #[test]
311    fn compute_tf_basic() {
312        let tf = compute_tf("fn handle_request request response handle");
313        assert!(tf.contains_key("handle"));
314        assert!(tf.contains_key("request"));
315        assert!(tf["handle"] > 0.0);
316    }
317
318    #[test]
319    fn cosine_identical() {
320        let mut a = HashMap::new();
321        a.insert("hello".to_string(), 1.0);
322        a.insert("world".to_string(), 0.5);
323        let sim = cosine_similarity(&a, &a);
324        assert!((sim - 1.0).abs() < 0.001);
325    }
326
327    #[test]
328    fn cosine_orthogonal() {
329        let mut a = HashMap::new();
330        a.insert("hello".to_string(), 1.0);
331        let mut b = HashMap::new();
332        b.insert("world".to_string(), 1.0);
333        let sim = cosine_similarity(&a, &b);
334        assert!(sim.abs() < 0.001);
335    }
336
337    #[test]
338    fn add_and_find_similar() {
339        let mut index = SemanticCacheIndex::default();
340        index.add_file(
341            "auth.rs",
342            "fn validate_token check jwt expiry auth login",
343            "s1",
344        );
345        index.add_file(
346            "db.rs",
347            "fn connect_database pool query insert delete",
348            "s1",
349        );
350
351        let results = index.find_similar("validate auth token jwt", 0.1);
352        assert!(!results.is_empty());
353        assert_eq!(results[0].0, "auth.rs");
354    }
355
356    #[test]
357    fn warmup_suggestions() {
358        let mut index = SemanticCacheIndex::default();
359        index.add_file("hot.rs", "frequently accessed file", "s1");
360        index.entries[0].access_count = 50;
361        index.add_file("cold.rs", "rarely used", "s1");
362
363        let warmup = index.suggest_warmup(1);
364        assert_eq!(warmup.len(), 1);
365        assert_eq!(warmup[0], "hot.rs");
366    }
367}
lean_ctx/core/semantic_cache.rs

lean_ctx/core/
semantic_cache.rs
