nusy_codegraph/
embeddings.rs

1//! Embedding generation and semantic search for code objects.
2//!
3//! Re-exports shared types from `nusy-graph-query` and adds codegraph-specific
4//! functions for embedding CodeNodes and attaching embeddings to RecordBatches.
5
6use crate::schema::{CODE_EMBEDDING_DIM, CodeNode, CodeNodeKind};
7use arrow::array::{Array, Float32Array, RecordBatch};
8use std::sync::Arc;
9
10// Re-export shared types so existing consumers don't break.
11pub use nusy_graph_query::embedding::{EmbeddingError, EmbeddingProvider, cosine_similarity};
12
13pub type Result<T> = std::result::Result<T, EmbeddingError>;
14
15/// Deterministic hash-based embedding provider for codegraph (768-dim).
16pub struct HashEmbeddingProvider;
17
18impl EmbeddingProvider for HashEmbeddingProvider {
19    fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
20        Ok(texts
21            .iter()
22            .map(|t| nusy_graph_query::hash_to_vector(t, CODE_EMBEDDING_DIM as usize))
23            .collect())
24    }
25
26    fn embed(&self, text: &str) -> Result<Vec<f32>> {
27        Ok(nusy_graph_query::hash_to_vector(
28            text,
29            CODE_EMBEDDING_DIM as usize,
30        ))
31    }
32
33    fn dim(&self) -> usize {
34        CODE_EMBEDDING_DIM as usize
35    }
36}
37
38/// Build the embeddable text for a CodeNode.
39///
40/// Concatenates signature + docstring for functions/methods,
41/// name + docstring for classes, docstring for modules.
42pub fn node_to_embed_text(node: &CodeNode) -> Option<String> {
43    let parts: Vec<&str> = [
44        node.signature.as_deref(),
45        node.docstring.as_deref(),
46        Some(node.name.as_str()),
47    ]
48    .into_iter()
49    .flatten()
50    .collect();
51
52    if parts.is_empty() || (parts.len() == 1 && parts[0] == node.name) {
53        match node.kind {
54            CodeNodeKind::File | CodeNodeKind::Module => {
55                if node.docstring.is_some() {
56                    Some(parts.join(" "))
57                } else {
58                    None
59                }
60            }
61            _ => None,
62        }
63    } else {
64        Some(parts.join(" "))
65    }
66}
67
68/// Embed all CodeNodes that have embeddable text.
69///
70/// Returns a map from node ID to embedding vector.
71pub fn embed_nodes(
72    nodes: &[CodeNode],
73    provider: &dyn EmbeddingProvider,
74) -> Result<Vec<(String, Vec<f32>)>> {
75    let embeddable: Vec<(String, String)> = nodes
76        .iter()
77        .filter_map(|n| node_to_embed_text(n).map(|text| (n.id.clone(), text)))
78        .collect();
79
80    if embeddable.is_empty() {
81        return Ok(Vec::new());
82    }
83
84    let texts: Vec<String> = embeddable.iter().map(|(_, t)| t.clone()).collect();
85    let vectors = provider.embed_batch(&texts)?;
86
87    for vec in &vectors {
88        if vec.len() != provider.dim() {
89            return Err(EmbeddingError::DimensionMismatch {
90                expected: provider.dim(),
91                actual: vec.len(),
92            });
93        }
94    }
95
96    Ok(embeddable
97        .into_iter()
98        .zip(vectors)
99        .map(|((id, _), vec)| (id, vec))
100        .collect())
101}
102
103/// Update a CodeNodes RecordBatch with embedding vectors.
104///
105/// Replaces the null embedding column with actual vectors for nodes
106/// that have embeddings. Nodes without embeddings remain null.
107pub fn attach_embeddings(
108    batch: &RecordBatch,
109    embeddings: &[(String, Vec<f32>)],
110) -> Result<RecordBatch> {
111    use arrow::array::{FixedSizeListArray, StringArray};
112    use arrow::buffer::{BooleanBuffer, NullBuffer};
113    use arrow::datatypes::{DataType, Field};
114
115    let dim = CODE_EMBEDDING_DIM as usize;
116    let n = batch.num_rows();
117
118    let embed_map: std::collections::HashMap<&str, &Vec<f32>> = embeddings
119        .iter()
120        .map(|(id, vec)| (id.as_str(), vec))
121        .collect();
122
123    let ids = batch
124        .column(crate::schema::node_col::ID)
125        .as_any()
126        .downcast_ref::<StringArray>()
127        .ok_or_else(|| EmbeddingError::Provider("id column is not StringArray".to_string()))?;
128
129    let mut values = Vec::with_capacity(n * dim);
130    let mut validity = Vec::with_capacity(n);
131
132    for i in 0..n {
133        let id = ids.value(i);
134        if let Some(vec) = embed_map.get(id) {
135            values.extend_from_slice(vec);
136            validity.push(true);
137        } else {
138            values.extend(std::iter::repeat_n(0.0f32, dim));
139            validity.push(false);
140        }
141    }
142
143    let embedding_field = Arc::new(Field::new("item", DataType::Float32, false));
144    let embedding_array = FixedSizeListArray::try_new(
145        embedding_field,
146        CODE_EMBEDDING_DIM,
147        Arc::new(Float32Array::from(values)),
148        Some(NullBuffer::new(BooleanBuffer::from(validity))),
149    )
150    .map_err(|e| EmbeddingError::Provider(e.to_string()))?;
151
152    let mut columns: Vec<Arc<dyn Array>> = Vec::new();
153    for col_idx in 0..batch.num_columns() {
154        if col_idx == crate::schema::node_col::EMBEDDING {
155            columns.push(Arc::new(embedding_array.clone()));
156        } else {
157            columns.push(batch.column(col_idx).clone());
158        }
159    }
160
161    RecordBatch::try_new(batch.schema(), columns)
162        .map_err(|e| EmbeddingError::Provider(e.to_string()))
163}
164
165/// A search result from semantic search.
166#[derive(Debug, Clone)]
167pub struct SearchResult {
168    /// The CodeNode ID.
169    pub id: String,
170    /// The node name.
171    pub name: String,
172    /// The node kind.
173    pub kind: CodeNodeKind,
174    /// Cosine similarity score (0.0 to 1.0 for unit vectors).
175    pub score: f32,
176}
177
178/// Semantic search over embedded CodeNodes.
179///
180/// Embeds the query text, computes cosine similarity against all
181/// embedded nodes, and returns the top-k results.
182pub fn semantic_search(
183    nodes: &[CodeNode],
184    embeddings: &[(String, Vec<f32>)],
185    query: &str,
186    provider: &dyn EmbeddingProvider,
187    top_k: usize,
188) -> Result<Vec<SearchResult>> {
189    let query_vec = provider.embed(query)?;
190
191    let mut results: Vec<SearchResult> = embeddings
192        .iter()
193        .filter_map(|(id, vec)| {
194            let score = cosine_similarity(&query_vec, vec);
195            let node = nodes.iter().find(|n| n.id == *id)?;
196            Some(SearchResult {
197                id: id.clone(),
198                name: node.name.clone(),
199                kind: node.kind,
200                score,
201            })
202        })
203        .collect();
204
205    results.sort_by(|a, b| {
206        b.score
207            .partial_cmp(&a.score)
208            .unwrap_or(std::cmp::Ordering::Equal)
209    });
210    results.truncate(top_k);
211
212    Ok(results)
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218
219    fn sample_nodes() -> Vec<CodeNode> {
220        vec![
221            CodeNode {
222                id: "func:brain/signal.py::fuse".to_string(),
223                kind: CodeNodeKind::Function,
224                parent_id: None,
225                name: "fuse".to_string(),
226                signature: Some("def fuse(signals: list) -> dict".to_string()),
227                docstring: Some("Fuse signals from multiple sources.".to_string()),
228                body_hash: None,
229                body: None,
230                loc: Some(20),
231                cyclomatic_complexity: Some(5),
232                coverage_pct: None,
233                last_modified: None,
234                ..Default::default()
235            },
236            CodeNode {
237                id: "func:brain/train.py::train_lora".to_string(),
238                kind: CodeNodeKind::Function,
239                parent_id: None,
240                name: "train_lora".to_string(),
241                signature: Some("def train_lora(model, data) -> None".to_string()),
242                docstring: Some("Train a LoRA adapter on the model.".to_string()),
243                body_hash: None,
244                body: None,
245                loc: Some(50),
246                cyclomatic_complexity: Some(8),
247                coverage_pct: None,
248                last_modified: None,
249                ..Default::default()
250            },
251            CodeNode {
252                id: "class:brain/store.py::Store".to_string(),
253                kind: CodeNodeKind::Class,
254                parent_id: None,
255                name: "Store".to_string(),
256                signature: Some("class Store".to_string()),
257                docstring: Some("Knowledge store for persisting graph data.".to_string()),
258                body_hash: None,
259                body: None,
260                loc: Some(100),
261                cyclomatic_complexity: None,
262                coverage_pct: None,
263                last_modified: None,
264                ..Default::default()
265            },
266            CodeNode {
267                id: "file:brain/empty.py".to_string(),
268                kind: CodeNodeKind::File,
269                parent_id: None,
270                name: "empty.py".to_string(),
271                signature: None,
272                docstring: None,
273                body_hash: None,
274                body: None,
275                loc: Some(1),
276                cyclomatic_complexity: None,
277                coverage_pct: None,
278                last_modified: None,
279                ..Default::default()
280            },
281        ]
282    }
283
284    #[test]
285    fn test_hash_embedding_provider_deterministic() {
286        let provider = HashEmbeddingProvider;
287        let v1 = provider.embed("hello world").unwrap();
288        let v2 = provider.embed("hello world").unwrap();
289        assert_eq!(v1, v2);
290        assert_eq!(v1.len(), CODE_EMBEDDING_DIM as usize);
291    }
292
293    #[test]
294    fn test_hash_embedding_unit_length() {
295        let provider = HashEmbeddingProvider;
296        let v = provider.embed("test input").unwrap();
297        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
298        assert!(
299            (norm - 1.0).abs() < 1e-5,
300            "Vector should be unit length, got norm={norm}"
301        );
302    }
303
304    #[test]
305    fn test_cosine_similarity_identical() {
306        let v = vec![1.0, 0.0, 0.0];
307        assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6);
308    }
309
310    #[test]
311    fn test_cosine_similarity_orthogonal() {
312        let a = vec![1.0, 0.0, 0.0];
313        let b = vec![0.0, 1.0, 0.0];
314        assert!(cosine_similarity(&a, &b).abs() < 1e-6);
315    }
316
317    #[test]
318    fn test_cosine_similarity_opposite() {
319        let a = vec![1.0, 0.0];
320        let b = vec![-1.0, 0.0];
321        assert!((cosine_similarity(&a, &b) + 1.0).abs() < 1e-6);
322    }
323
324    #[test]
325    fn test_node_to_embed_text() {
326        let nodes = sample_nodes();
327        let text = node_to_embed_text(&nodes[0]).expect("should embed");
328        assert!(text.contains("fuse"));
329        assert!(text.contains("signals"));
330
331        let text = node_to_embed_text(&nodes[3]);
332        assert!(text.is_none(), "File without docstring should not embed");
333    }
334
335    #[test]
336    fn test_embed_nodes() {
337        let nodes = sample_nodes();
338        let provider = HashEmbeddingProvider;
339        let embeddings = embed_nodes(&nodes, &provider).unwrap();
340
341        assert_eq!(embeddings.len(), 3);
342        for (_, vec) in &embeddings {
343            assert_eq!(vec.len(), CODE_EMBEDDING_DIM as usize);
344        }
345    }
346
347    #[test]
348    fn test_semantic_search() {
349        let nodes = sample_nodes();
350        let provider = HashEmbeddingProvider;
351        let embeddings = embed_nodes(&nodes, &provider).unwrap();
352
353        let results = semantic_search(&nodes, &embeddings, "signal fusion", &provider, 3).unwrap();
354
355        assert!(!results.is_empty());
356        assert!(results.len() <= 3);
357
358        for r in &results {
359            assert!(r.score >= -1.0 && r.score <= 1.0);
360        }
361        for w in results.windows(2) {
362            assert!(w[0].score >= w[1].score);
363        }
364    }
365
366    #[test]
367    fn test_attach_embeddings_to_batch() {
368        use crate::schema::build_code_nodes_batch;
369
370        let nodes = sample_nodes();
371        let batch = build_code_nodes_batch(&nodes).expect("build batch");
372
373        let emb_col = batch.column(crate::schema::node_col::EMBEDDING);
374        for i in 0..batch.num_rows() {
375            assert!(emb_col.is_null(i), "Row {i} should be null initially");
376        }
377
378        let provider = HashEmbeddingProvider;
379        let embeddings = embed_nodes(&nodes, &provider).unwrap();
380        let updated = attach_embeddings(&batch, &embeddings).expect("attach");
381
382        assert_eq!(updated.num_rows(), batch.num_rows());
383        assert_eq!(updated.num_columns(), batch.num_columns());
384
385        let emb_col = updated.column(crate::schema::node_col::EMBEDDING);
386        assert!(!emb_col.is_null(0), "fuse should be embedded");
387        assert!(emb_col.is_null(3), "file without docstring should be null");
388    }
389
390    #[test]
391    fn test_embed_batch_consistency() {
392        let provider = HashEmbeddingProvider;
393        let texts = vec!["hello".to_string(), "world".to_string()];
394        let batch_result = provider.embed_batch(&texts).unwrap();
395        let single_1 = provider.embed("hello").unwrap();
396        let single_2 = provider.embed("world").unwrap();
397        assert_eq!(batch_result[0], single_1);
398        assert_eq!(batch_result[1], single_2);
399    }
400
401    #[test]
402    fn test_cosine_similarity_empty() {
403        assert_eq!(cosine_similarity(&[], &[]), 0.0);
404    }
405
406    #[test]
407    fn test_cosine_similarity_length_mismatch() {
408        assert_eq!(cosine_similarity(&[1.0], &[1.0, 2.0]), 0.0);
409    }
410}
nusy_codegraph/embeddings.rs

nusy_codegraph/
embeddings.rs
