gitcortex_mcp/
embeddings.rs

1//! Semantic search via local embeddings (AllMiniLM-L6-v2, 384 dims).
2//!
3//! Model is downloaded from HuggingFace on first use (~23 MB), cached in
4//! `$XDG_DATA_HOME/gitcortex/models` (never inside a repo). All subsequent
5//! starts load from cache.
6//!
7//! Vector index is persisted per-branch at:
8//!   `~/.local/share/gitcortex/{repo_id}/embeddings_{branch}.bin`
9//!
10//! Background indexer (`index_missing`) embeds nodes that don't yet have a
11//! vector. Call it once after `gcx serve` opens the store. Search stays
12//! text-only while the indexer runs; it automatically uses semantic hits once
13//! at least one vector is loaded.
14
15use std::collections::HashMap;
16use std::io::{BufWriter, Write};
17use std::path::{Path, PathBuf};
18
19use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
20use gitcortex_core::graph::Node;
21
22use crate::mcp::search::tokenize;
23
24/// Minimum cosine similarity to surface as a semantic hit.
25const SIMILARITY_THRESHOLD: f32 = 0.50;
26const DIM: usize = 384;
27
28// Binary format: magic + version + dim + count + entries
29const MAGIC: &[u8; 4] = b"GCXV";
30const FORMAT_VERSION: u32 = 2;
31
32// ── Vector index ──────────────────────────────────────────────────────────────
33
34pub struct SemanticIndex {
35    /// node_id → unit-normalised embedding
36    vectors: HashMap<String, Vec<f32>>,
37    path: PathBuf,
38}
39
40impl SemanticIndex {
41    pub fn load_or_create(path: &Path) -> Self {
42        let vectors = load_bin(path).unwrap_or_default();
43        if !vectors.is_empty() {
44            tracing::info!(
45                "semantic index loaded: {} vectors from {}",
46                vectors.len(),
47                path.display()
48            );
49        }
50        Self {
51            vectors,
52            path: path.to_owned(),
53        }
54    }
55
56    pub fn has(&self, node_id: &str) -> bool {
57        self.vectors.contains_key(node_id)
58    }
59
60    pub fn insert(&mut self, node_id: String, vec: Vec<f32>) {
61        self.vectors.insert(node_id, unit_normalise(vec));
62    }
63
64    pub fn len(&self) -> usize {
65        self.vectors.len()
66    }
67
68    pub fn is_empty(&self) -> bool {
69        self.vectors.is_empty()
70    }
71
72    /// Drop vectors whose node ID is not in `live_ids`. Node UUIDs regenerate
73    /// on every re-index, so without pruning the index file grows with
74    /// orphaned vectors that can still surface as (unresolvable) hits.
75    /// Returns the number of vectors removed.
76    pub fn retain_ids(&mut self, live_ids: &std::collections::HashSet<String>) -> usize {
77        let before = self.vectors.len();
78        self.vectors.retain(|id, _| live_ids.contains(id));
79        before - self.vectors.len()
80    }
81
82    pub fn save(&self) {
83        if let Err(e) = save_bin(&self.path, &self.vectors) {
84            tracing::warn!("failed to save semantic index: {e}");
85        }
86    }
87
88    /// Return up to `k` `(node_id, similarity)` pairs with cosine similarity ≥ SIMILARITY_THRESHOLD.
89    /// Query vector need not be pre-normalised — normalised internally.
90    pub fn top_k(&self, query_vec: &[f32], k: usize) -> Vec<(String, f32)> {
91        let q = unit_normalise(query_vec.to_vec());
92        let mut scores: Vec<(&String, f32)> = self
93            .vectors
94            .iter()
95            .map(|(id, v)| (id, dot(&q, v)))
96            .filter(|(_, s)| *s >= SIMILARITY_THRESHOLD)
97            .collect();
98        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
99        scores
100            .into_iter()
101            .take(k)
102            .map(|(id, s)| (id.clone(), s))
103            .collect()
104    }
105}
106
107// ── Embedder ──────────────────────────────────────────────────────────────────
108
109pub struct Embedder {
110    model: TextEmbedding,
111}
112
113impl Embedder {
114    /// Download (first run) or load (cached) AllMiniLM-L6-v2.
115    ///
116    /// `cache_dir` is where fastembed stores the downloaded model weights.
117    /// Pass `branch::models_dir()` so the cache lands in
118    /// `$XDG_DATA_HOME/gitcortex/models`, never inside a repo.
119    pub fn new(cache_dir: &Path) -> anyhow::Result<Self> {
120        std::fs::create_dir_all(cache_dir)?;
121        tracing::info!("initialising semantic embedder (AllMiniLM-L6-v2) …");
122        let model = TextEmbedding::try_new(
123            InitOptions::new(EmbeddingModel::AllMiniLML6V2)
124                .with_show_download_progress(false)
125                .with_cache_dir(cache_dir.to_path_buf()),
126        )?;
127        tracing::info!("semantic embedder ready");
128        Ok(Self { model })
129    }
130
131    pub fn embed_one(&self, text: &str) -> anyhow::Result<Vec<f32>> {
132        let mut out = self.model.embed(vec![text.to_owned()], None)?;
133        out.pop()
134            .ok_or_else(|| anyhow::anyhow!("embedder returned no vectors"))
135    }
136
137    /// Embed a batch of texts. Returns one vector per input in order.
138    pub fn embed_batch(&self, texts: Vec<String>) -> anyhow::Result<Vec<Vec<f32>>> {
139        self.model.embed(texts, None)
140    }
141}
142
143// ── Text representation for a node ────────────────────────────────────────────
144
145/// Build the text string that gets embedded for a node.
146///
147/// Appends tokenized identifier words (CamelCase/snake_case → space-separated
148/// lowercase) so NL queries like "validate token" match `validate_token`
149/// without relying on the model to unsplit glued identifiers.
150pub fn node_text(n: &Node) -> String {
151    let kind = n.kind.to_string();
152    let sig = &n.metadata.definition.signature;
153    let doc = n.metadata.definition.doc_comment.as_deref().unwrap_or("");
154
155    // Tokenize the simple name and the last segment of the qualified path.
156    let name_words = tokenize(&n.name).join(" ");
157    let qname_last = n
158        .qualified_name
159        .rsplit("::")
160        .next()
161        .unwrap_or(&n.qualified_name);
162    let qname_words = if qname_last != n.name {
163        tokenize(qname_last).join(" ")
164    } else {
165        String::new()
166    };
167
168    let mut parts = vec![kind.as_str(), n.qualified_name.as_str()];
169    if !sig.is_empty() {
170        parts.push(sig.as_str());
171    }
172    if !doc.is_empty() {
173        parts.push(doc);
174    }
175    parts.push(name_words.as_str());
176    if !qname_words.is_empty() {
177        parts.push(qname_words.as_str());
178    }
179    parts.join(" ")
180}
181
182// ── Math helpers ──────────────────────────────────────────────────────────────
183
184fn dot(a: &[f32], b: &[f32]) -> f32 {
185    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
186}
187
188fn unit_normalise(mut v: Vec<f32>) -> Vec<f32> {
189    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
190    if norm > f32::EPSILON {
191        for x in &mut v {
192            *x /= norm;
193        }
194    }
195    v
196}
197
198// ── Binary storage ────────────────────────────────────────────────────────────
199//
200// Layout (all integers little-endian):
201//   [4]  magic "GCXV"
202//   [4]  format version (u32)
203//   [4]  embedding dimension (u32)
204//   [4]  record count (u32)
205//   per record:
206//     [4]       id_len (u32)
207//     [id_len]  node_id (UTF-8)
208//     [dim × 4] f32 values
209
210fn load_bin(path: &Path) -> Option<HashMap<String, Vec<f32>>> {
211    let data = std::fs::read(path).ok()?;
212    let mut p = 0usize;
213
214    macro_rules! read_u32 {
215        () => {{
216            let b: [u8; 4] = data.get(p..p + 4)?.try_into().ok()?;
217            p += 4;
218            u32::from_le_bytes(b)
219        }};
220    }
221
222    if data.get(p..p + 4)? != MAGIC {
223        return None;
224    }
225    p += 4;
226
227    let ver = read_u32!();
228    if ver != FORMAT_VERSION {
229        return None;
230    }
231    let dim = read_u32!() as usize;
232    let count = read_u32!() as usize;
233
234    let mut map = HashMap::with_capacity(count);
235    for _ in 0..count {
236        let id_len = read_u32!() as usize;
237        let id = String::from_utf8(data.get(p..p + id_len)?.to_vec()).ok()?;
238        p += id_len;
239        let end = p + dim * 4;
240        let vec: Vec<f32> = data
241            .get(p..end)?
242            .chunks_exact(4)
243            .map(|b| f32::from_le_bytes(b.try_into().unwrap()))
244            .collect();
245        p = end;
246        map.insert(id, vec);
247    }
248    Some(map)
249}
250
251fn save_bin(path: &Path, vectors: &HashMap<String, Vec<f32>>) -> std::io::Result<()> {
252    if let Some(parent) = path.parent() {
253        std::fs::create_dir_all(parent)?;
254    }
255    let tmp = path.with_extension("tmp");
256    {
257        let f = std::fs::File::create(&tmp)?;
258        let mut w = BufWriter::new(f);
259        w.write_all(MAGIC)?;
260        w.write_all(&FORMAT_VERSION.to_le_bytes())?;
261        w.write_all(&(DIM as u32).to_le_bytes())?;
262        w.write_all(&(vectors.len() as u32).to_le_bytes())?;
263        for (id, vec) in vectors {
264            let id_b = id.as_bytes();
265            w.write_all(&(id_b.len() as u32).to_le_bytes())?;
266            w.write_all(id_b)?;
267            for &v in vec {
268                w.write_all(&v.to_le_bytes())?;
269            }
270        }
271        w.flush()?;
272    }
273    std::fs::rename(&tmp, path)?;
274    Ok(())
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use gitcortex_core::graph::{NodeId, NodeMetadata, Span};
281    use gitcortex_core::schema::NodeKind;
282    use std::path::PathBuf;
283
284    fn make_node(name: &str, qualified_name: &str, sig: &str, doc: &str) -> Node {
285        let mut meta = NodeMetadata::default();
286        meta.definition.signature = sig.to_owned();
287        meta.definition.doc_comment = if doc.is_empty() {
288            None
289        } else {
290            Some(doc.to_owned())
291        };
292        Node {
293            id: NodeId::default(),
294            kind: NodeKind::Function,
295            name: name.to_owned(),
296            qualified_name: qualified_name.to_owned(),
297            file: PathBuf::from("src/lib.rs"),
298            span: Span {
299                start_line: 1,
300                end_line: 5,
301            },
302            metadata: meta,
303        }
304    }
305
306    #[test]
307    fn node_text_contains_tokenized_words() {
308        let n = make_node(
309            "validate_token",
310            "auth::validate_token",
311            "fn validate_token(t: &str) -> bool",
312            "",
313        );
314        let text = node_text(&n);
315        assert!(
316            text.contains("validate token"),
317            "expected 'validate token' in: {text}"
318        );
319        assert!(
320            text.contains("auth::validate_token"),
321            "expected qualified name in: {text}"
322        );
323    }
324
325    #[test]
326    fn node_text_qualified_segment_tokenized_when_differs_from_name() {
327        let n = make_node("new", "http::HttpClient::new", "", "");
328        let text = node_text(&n);
329        assert!(text.contains("new"), "expected 'new' in: {text}");
330    }
331
332    #[test]
333    fn node_text_includes_doc_and_sig() {
334        let n = make_node(
335            "parse_json",
336            "util::parse_json",
337            "fn parse_json(s: &str) -> Value",
338            "Parse a JSON string.",
339        );
340        let text = node_text(&n);
341        assert!(text.contains("Parse a JSON string."));
342        assert!(text.contains("fn parse_json"));
343        assert!(text.contains("parse json"));
344    }
345
346    #[test]
347    fn load_bin_rejects_stale_version() {
348        let dir = tempfile::tempdir().unwrap();
349        let path = dir.path().join("test.bin");
350        let mut buf = Vec::new();
351        buf.extend_from_slice(b"GCXV");
352        buf.extend_from_slice(&1u32.to_le_bytes()); // old version
353        buf.extend_from_slice(&384u32.to_le_bytes());
354        buf.extend_from_slice(&0u32.to_le_bytes());
355        std::fs::write(&path, &buf).unwrap();
356        assert!(load_bin(&path).is_none(), "v1 file should be rejected");
357    }
358
359    #[test]
360    fn save_and_load_roundtrip() {
361        let dir = tempfile::tempdir().unwrap();
362        let path = dir.path().join("idx.bin");
363        let mut vecs: HashMap<String, Vec<f32>> = HashMap::new();
364        vecs.insert("node-1".to_owned(), vec![1.0; 384]);
365        vecs.insert("node-2".to_owned(), vec![0.5; 384]);
366        save_bin(&path, &vecs).unwrap();
367        let loaded = load_bin(&path).expect("should load v2 file");
368        assert_eq!(loaded.len(), 2);
369        assert!(loaded.contains_key("node-1"));
370    }
371
372    #[test]
373    fn top_k_returns_scores_in_range() {
374        let mut index = SemanticIndex {
375            vectors: HashMap::new(),
376            path: PathBuf::from("/tmp/unused"),
377        };
378        let v: Vec<f32> = {
379            let mut raw = vec![0.0f32; 384];
380            raw[0] = 1.0;
381            raw
382        };
383        index.insert("a".to_owned(), v.clone());
384        index.insert("b".to_owned(), v.clone());
385        let results = index.top_k(&v, 10);
386        assert_eq!(results.len(), 2);
387        for (_, score) in &results {
388            assert!(
389                *score >= SIMILARITY_THRESHOLD,
390                "score {score} below threshold"
391            );
392            assert!(*score <= 1.001, "score {score} above 1.0");
393        }
394        assert!(results[0].1 >= results[1].1);
395    }
396
397    #[test]
398    fn top_k_respects_k_limit() {
399        let mut index = SemanticIndex {
400            vectors: HashMap::new(),
401            path: PathBuf::from("/tmp/unused"),
402        };
403        let v: Vec<f32> = {
404            let mut raw = vec![0.0f32; 384];
405            raw[0] = 1.0;
406            raw
407        };
408        for i in 0..20u32 {
409            index.insert(format!("node-{i}"), v.clone());
410        }
411        let results = index.top_k(&v, 5);
412        assert_eq!(results.len(), 5);
413    }
414}
gitcortex_mcp/embeddings.rs

gitcortex_mcp/
embeddings.rs
