codesynapse_core/
embed_pipeline.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::path::PathBuf;
4
5pub trait NodeEmbedder {
6    fn embed_nodes(&self, nodes: &[(&str, &str)]) -> HashMap<String, Vec<f32>>;
7}
8
9impl NodeEmbedder for crate::embedding::StaticEmbedder {
10    fn embed_nodes(&self, nodes: &[(&str, &str)]) -> HashMap<String, Vec<f32>> {
11        crate::embedding::StaticEmbedder::embed_nodes(self, nodes)
12    }
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16struct EmbedEntry {
17    hash: String,
18    embeddings: HashMap<String, Vec<f32>>,
19}
20
21pub struct EmbedCache {
22    path: PathBuf,
23    entries: HashMap<String, EmbedEntry>,
24}
25
26impl EmbedCache {
27    pub fn load(path: PathBuf) -> Self {
28        let entries = std::fs::read_to_string(&path)
29            .ok()
30            .and_then(|s| serde_json::from_str(&s).ok())
31            .unwrap_or_default();
32        Self { path, entries }
33    }
34
35    pub fn save(&self) -> Result<(), String> {
36        if let Some(parent) = self.path.parent() {
37            std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
38        }
39        let json = serde_json::to_string(&self.entries).map_err(|e| e.to_string())?;
40        std::fs::write(&self.path, json).map_err(|e| e.to_string())
41    }
42
43    pub fn get_if_fresh<'a>(
44        &'a self,
45        file_path: &str,
46        hash: &str,
47    ) -> Option<&'a HashMap<String, Vec<f32>>> {
48        self.entries
49            .get(file_path)
50            .filter(|e| e.hash == hash)
51            .map(|e| &e.embeddings)
52    }
53
54    pub fn insert(
55        &mut self,
56        file_path: String,
57        hash: String,
58        embeddings: HashMap<String, Vec<f32>>,
59    ) {
60        self.entries
61            .insert(file_path, EmbedEntry { hash, embeddings });
62    }
63}
64
65pub fn embed_file_nodes<E: NodeEmbedder>(
66    embedder: &E,
67    nodes: &[(&str, &str)],
68    file_path: &str,
69    hash: &str,
70    cache: &mut EmbedCache,
71) -> HashMap<String, Vec<f32>> {
72    if let Some(cached) = cache.get_if_fresh(file_path, hash) {
73        return cached.clone();
74    }
75    let embs = embedder.embed_nodes(nodes);
76    cache.insert(file_path.to_string(), hash.to_string(), embs.clone());
77    embs
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use tempfile::tempdir;
84
85    struct FakeEmbedder;
86
87    impl NodeEmbedder for FakeEmbedder {
88        fn embed_nodes(&self, nodes: &[(&str, &str)]) -> HashMap<String, Vec<f32>> {
89            nodes
90                .iter()
91                .map(|(id, _)| (id.to_string(), vec![1.0f32, 2.0, 3.0]))
92                .collect()
93        }
94    }
95
96    #[test]
97    fn test_embed_file_nodes_computes_embeddings() {
98        let tmp = tempdir().unwrap();
99        let cache_path = tmp.path().join("embed_cache.json");
100        let mut cache = EmbedCache::load(cache_path);
101        let embedder = FakeEmbedder;
102        let nodes = vec![("node1", "MyClass"), ("node2", "doThing")];
103        let result = embed_file_nodes(&embedder, &nodes, "src/foo.rs", "abc123", &mut cache);
104        assert_eq!(result.len(), 2);
105        assert!(result.contains_key("node1"));
106        assert_eq!(result["node1"], vec![1.0f32, 2.0, 3.0]);
107    }
108
109    #[test]
110    fn test_embed_file_nodes_cache_hit() {
111        let tmp = tempdir().unwrap();
112        let cache_path = tmp.path().join("embed_cache.json");
113        let mut cache = EmbedCache::load(cache_path);
114        let embedder = FakeEmbedder;
115        let nodes = vec![("n1", "Foo")];
116
117        let first = embed_file_nodes(&embedder, &nodes, "src/bar.rs", "hash1", &mut cache);
118        // Insert sentinel to detect if embedder called again
119        cache.insert(
120            "src/bar.rs".to_string(),
121            "hash1".to_string(),
122            [("n1".to_string(), vec![9.0f32])].into_iter().collect(),
123        );
124        let second = embed_file_nodes(&embedder, &nodes, "src/bar.rs", "hash1", &mut cache);
125        // second should return the sentinel (cache hit), not recompute
126        assert_eq!(second["n1"], vec![9.0f32]);
127        let _ = first;
128    }
129
130    #[test]
131    fn test_embed_file_nodes_cache_miss_on_changed_hash() {
132        let tmp = tempdir().unwrap();
133        let cache_path = tmp.path().join("embed_cache.json");
134        let mut cache = EmbedCache::load(cache_path);
135        let embedder = FakeEmbedder;
136        let nodes = vec![("n1", "Foo")];
137
138        embed_file_nodes(&embedder, &nodes, "src/baz.rs", "hash1", &mut cache);
139        // Change the stored entry to a sentinel
140        cache.insert(
141            "src/baz.rs".to_string(),
142            "hash1".to_string(),
143            [("n1".to_string(), vec![9.0f32])].into_iter().collect(),
144        );
145        // Now call with a different hash — should recompute (not return sentinel)
146        let result = embed_file_nodes(&embedder, &nodes, "src/baz.rs", "hash2", &mut cache);
147        assert_eq!(result["n1"], vec![1.0f32, 2.0, 3.0]);
148    }
149
150    #[test]
151    fn test_embed_cache_save_and_reload() {
152        let tmp = tempdir().unwrap();
153        let cache_path = tmp.path().join("embed_cache.json");
154
155        let mut cache = EmbedCache::load(cache_path.clone());
156        let embedder = FakeEmbedder;
157        let nodes = vec![("x", "Thing")];
158        embed_file_nodes(&embedder, &nodes, "src/x.rs", "deadbeef", &mut cache);
159        cache.save().unwrap();
160
161        let cache2 = EmbedCache::load(cache_path);
162        let hit = cache2.get_if_fresh("src/x.rs", "deadbeef");
163        assert!(hit.is_some());
164        assert_eq!(hit.unwrap()["x"], vec![1.0f32, 2.0, 3.0]);
165    }
166
167    #[test]
168    fn test_embed_cache_stale_hash_returns_none() {
169        let tmp = tempdir().unwrap();
170        let cache_path = tmp.path().join("embed_cache.json");
171        let mut cache = EmbedCache::load(cache_path);
172        cache.insert("f.rs".to_string(), "old".to_string(), Default::default());
173        assert!(cache.get_if_fresh("f.rs", "new").is_none());
174        assert!(cache.get_if_fresh("f.rs", "old").is_some());
175    }
176}
codesynapse_core/embed_pipeline.rs

codesynapse_core/
embed_pipeline.rs
