ruvector_gnn_node/
lib.rs

1//! Node.js bindings for Ruvector GNN via NAPI-RS
2//!
3//! This module provides JavaScript bindings for the Ruvector GNN library,
4//! enabling graph neural network operations, tensor compression, and
5//! differentiable search in Node.js applications.
6
7#![allow(clippy::all)]
8
9use napi::bindgen_prelude::*;
10use napi_derive::napi;
11use ruvector_gnn::{
12    compress::{
13        CompressedTensor as RustCompressedTensor, CompressionLevel as RustCompressionLevel,
14        TensorCompress as RustTensorCompress,
15    },
16    layer::RuvectorLayer as RustRuvectorLayer,
17    search::{
18        differentiable_search as rust_differentiable_search,
19        hierarchical_forward as rust_hierarchical_forward,
20    },
21};
22
23// ==================== RuvectorLayer Bindings ====================
24
25/// Graph Neural Network layer for HNSW topology
26#[napi]
27pub struct RuvectorLayer {
28    inner: RustRuvectorLayer,
29}
30
31#[napi]
32impl RuvectorLayer {
33    /// Create a new Ruvector GNN layer
34    ///
35    /// # Arguments
36    /// * `input_dim` - Dimension of input node embeddings
37    /// * `hidden_dim` - Dimension of hidden representations
38    /// * `heads` - Number of attention heads
39    /// * `dropout` - Dropout rate (0.0 to 1.0)
40    ///
41    /// # Example
42    /// ```javascript
43    /// const layer = new RuvectorLayer(128, 256, 4, 0.1);
44    /// ```
45    #[napi(constructor)]
46    pub fn new(input_dim: u32, hidden_dim: u32, heads: u32, dropout: f64) -> Result<Self> {
47        let inner = RustRuvectorLayer::new(
48            input_dim as usize,
49            hidden_dim as usize,
50            heads as usize,
51            dropout as f32,
52        )
53        .map_err(|e| Error::new(Status::InvalidArg, e.to_string()))?;
54
55        Ok(Self { inner })
56    }
57
58    /// Forward pass through the GNN layer
59    ///
60    /// # Arguments
61    /// * `node_embedding` - Current node's embedding (Float32Array)
62    /// * `neighbor_embeddings` - Embeddings of neighbor nodes (Array of Float32Array)
63    /// * `edge_weights` - Weights of edges to neighbors (Float32Array)
64    ///
65    /// # Returns
66    /// Updated node embedding as Float32Array
67    ///
68    /// # Example
69    /// ```javascript
70    /// const node = new Float32Array([1.0, 2.0, 3.0, 4.0]);
71    /// const neighbors = [new Float32Array([0.5, 1.0, 1.5, 2.0]), new Float32Array([2.0, 3.0, 4.0, 5.0])];
72    /// const weights = new Float32Array([0.3, 0.7]);
73    /// const output = layer.forward(node, neighbors, weights);
74    /// ```
75    #[napi]
76    pub fn forward(
77        &self,
78        node_embedding: Float32Array,
79        neighbor_embeddings: Vec<Float32Array>,
80        edge_weights: Float32Array,
81    ) -> Result<Float32Array> {
82        validate_dimension(
83            "node_embedding",
84            node_embedding.len(),
85            self.inner.input_dim(),
86        )?;
87        for neighbor in &neighbor_embeddings {
88            validate_dimension("neighbor_embedding", neighbor.len(), self.inner.input_dim())?;
89        }
90        let node_slice = node_embedding.as_ref();
91        let neighbors_vec: Vec<Vec<f32>> = neighbor_embeddings
92            .into_iter()
93            .map(|arr| arr.to_vec())
94            .collect();
95        let weights_slice = edge_weights.as_ref();
96
97        let result = self
98            .inner
99            .forward(node_slice, &neighbors_vec, weights_slice);
100
101        Ok(Float32Array::new(result))
102    }
103
104    /// Serialize the layer to JSON
105    #[napi]
106    pub fn to_json(&self) -> Result<String> {
107        serde_json::to_string(&self.inner).map_err(|e| {
108            Error::new(
109                Status::GenericFailure,
110                format!("Serialization error: {}", e),
111            )
112        })
113    }
114
115    /// Deserialize the layer from JSON
116    #[napi(factory)]
117    pub fn from_json(json: String) -> Result<Self> {
118        let inner: RustRuvectorLayer = serde_json::from_str(&json)
119            .map_err(|e| Error::new(Status::InvalidArg, format!("Deserialization error: {}", e)))?;
120        inner
121            .validate()
122            .map_err(|e| Error::new(Status::InvalidArg, e.to_string()))?;
123        Ok(Self { inner })
124    }
125}
126
127// ==================== TensorCompress Bindings ====================
128
129/// Compression level for tensor compression
130#[napi(object)]
131pub struct CompressionLevelConfig {
132    /// Type of compression: "none", "half", "pq8", "pq4", "binary"
133    pub level_type: String,
134    /// Scale factor (for "half" compression)
135    pub scale: Option<f64>,
136    /// Number of subvectors (for PQ compression)
137    pub subvectors: Option<u32>,
138    /// Number of centroids (for PQ8)
139    pub centroids: Option<u32>,
140    /// Outlier threshold (for PQ4)
141    pub outlier_threshold: Option<f64>,
142    /// Binary threshold (for binary compression)
143    pub threshold: Option<f64>,
144}
145
146impl CompressionLevelConfig {
147    fn to_rust(&self) -> Result<RustCompressionLevel> {
148        // The core stores these as u8; `as u8` would silently wrap 256 -> 0.
149        for (name, value) in [
150            ("subvectors", self.subvectors),
151            ("centroids", self.centroids),
152        ] {
153            if value.is_some_and(|v| v > u8::MAX as u32) {
154                return Err(Error::new(
155                    Status::InvalidArg,
156                    format!("{} must be at most {}", name, u8::MAX),
157                ));
158            }
159        }
160        match self.level_type.as_str() {
161            "none" => Ok(RustCompressionLevel::None),
162            "half" => Ok(RustCompressionLevel::Half {
163                scale: self.scale.unwrap_or(1.0) as f32,
164            }),
165            "pq8" => Ok(RustCompressionLevel::PQ8 {
166                subvectors: self.subvectors.unwrap_or(8) as u8,
167                centroids: self.centroids.unwrap_or(16) as u8,
168            }),
169            "pq4" => Ok(RustCompressionLevel::PQ4 {
170                subvectors: self.subvectors.unwrap_or(8) as u8,
171                outlier_threshold: self.outlier_threshold.unwrap_or(3.0) as f32,
172            }),
173            "binary" => Ok(RustCompressionLevel::Binary {
174                threshold: self.threshold.unwrap_or(0.0) as f32,
175            }),
176            _ => Err(Error::new(
177                Status::InvalidArg,
178                format!("Invalid compression level: {}", self.level_type),
179            )),
180        }
181    }
182}
183
184/// Tensor compressor with adaptive level selection
185#[napi]
186pub struct TensorCompress {
187    inner: RustTensorCompress,
188}
189
190#[napi]
191impl TensorCompress {
192    /// Create a new tensor compressor
193    ///
194    /// # Example
195    /// ```javascript
196    /// const compressor = new TensorCompress();
197    /// ```
198    #[napi(constructor)]
199    pub fn new() -> Self {
200        Self {
201            inner: RustTensorCompress::new(),
202        }
203    }
204
205    /// Compress an embedding based on access frequency
206    ///
207    /// # Arguments
208    /// * `embedding` - The input embedding vector (Float32Array)
209    /// * `access_freq` - Access frequency in range [0.0, 1.0]
210    ///
211    /// # Returns
212    /// Compressed tensor as JSON string
213    ///
214    /// # Example
215    /// ```javascript
216    /// const embedding = new Float32Array([1.0, 2.0, 3.0, 4.0]);
217    /// const compressed = compressor.compress(embedding, 0.5);
218    /// ```
219    #[napi]
220    pub fn compress(&self, embedding: Float32Array, access_freq: f64) -> Result<String> {
221        let embedding_slice = embedding.as_ref();
222
223        let compressed = self
224            .inner
225            .compress(embedding_slice, access_freq as f32)
226            .map_err(|e| Error::new(Status::GenericFailure, format!("Compression error: {}", e)))?;
227
228        serde_json::to_string(&compressed).map_err(|e| {
229            Error::new(
230                Status::GenericFailure,
231                format!("Serialization error: {}", e),
232            )
233        })
234    }
235
236    /// Compress with explicit compression level
237    ///
238    /// # Arguments
239    /// * `embedding` - The input embedding vector (Float32Array)
240    /// * `level` - Compression level configuration
241    ///
242    /// # Returns
243    /// Compressed tensor as JSON string
244    ///
245    /// # Example
246    /// ```javascript
247    /// const embedding = new Float32Array([1.0, 2.0, 3.0, 4.0]);
248    /// const level = { level_type: "half", scale: 1.0 };
249    /// const compressed = compressor.compressWithLevel(embedding, level);
250    /// ```
251    #[napi]
252    pub fn compress_with_level(
253        &self,
254        embedding: Float32Array,
255        level: CompressionLevelConfig,
256    ) -> Result<String> {
257        let embedding_slice = embedding.as_ref();
258        let rust_level = level.to_rust()?;
259
260        let compressed = self
261            .inner
262            .compress_with_level(embedding_slice, &rust_level)
263            .map_err(|e| Error::new(Status::GenericFailure, format!("Compression error: {}", e)))?;
264
265        serde_json::to_string(&compressed).map_err(|e| {
266            Error::new(
267                Status::GenericFailure,
268                format!("Serialization error: {}", e),
269            )
270        })
271    }
272
273    /// Decompress a compressed tensor
274    ///
275    /// # Arguments
276    /// * `compressed_json` - Compressed tensor as JSON string
277    ///
278    /// # Returns
279    /// Decompressed embedding vector as Float32Array
280    ///
281    /// # Example
282    /// ```javascript
283    /// const decompressed = compressor.decompress(compressed);
284    /// ```
285    #[napi]
286    pub fn decompress(&self, compressed_json: String) -> Result<Float32Array> {
287        let compressed: RustCompressedTensor =
288            serde_json::from_str(&compressed_json).map_err(|e| {
289                Error::new(
290                    Status::GenericFailure,
291                    format!("Deserialization error: {}", e),
292                )
293            })?;
294
295        let result = self.inner.decompress(&compressed).map_err(|e| {
296            Error::new(
297                Status::GenericFailure,
298                format!("Decompression error: {}", e),
299            )
300        })?;
301
302        Ok(Float32Array::new(result))
303    }
304}
305
306// ==================== Search Functions ====================
307
308/// Result from differentiable search
309#[napi(object)]
310pub struct SearchResult {
311    /// Indices of top-k candidates
312    pub indices: Vec<u32>,
313    /// Soft weights for top-k candidates
314    pub weights: Vec<f64>,
315}
316
317/// Differentiable search using soft attention mechanism
318///
319/// # Arguments
320/// * `query` - The query vector (Float32Array)
321/// * `candidate_embeddings` - List of candidate embedding vectors (Array of Float32Array)
322/// * `k` - Number of top results to return
323/// * `temperature` - Temperature for softmax (lower = sharper, higher = smoother)
324///
325/// # Returns
326/// Search result with indices and soft weights
327///
328/// # Example
329/// ```javascript
330/// const query = new Float32Array([1.0, 0.0, 0.0]);
331/// const candidates = [new Float32Array([1.0, 0.0, 0.0]), new Float32Array([0.9, 0.1, 0.0]), new Float32Array([0.0, 1.0, 0.0])];
332/// const result = differentiableSearch(query, candidates, 2, 1.0);
333/// console.log(result.indices); // [0, 1]
334/// console.log(result.weights); // [0.x, 0.y]
335/// ```
336#[napi]
337pub fn differentiable_search(
338    query: Float32Array,
339    candidate_embeddings: Vec<Float32Array>,
340    k: u32,
341    temperature: f64,
342) -> Result<SearchResult> {
343    for candidate in &candidate_embeddings {
344        validate_dimension("candidate_embedding", candidate.len(), query.len())?;
345    }
346    let query_slice = query.as_ref();
347    let candidates_vec: Vec<Vec<f32>> = candidate_embeddings
348        .into_iter()
349        .map(|arr| arr.to_vec())
350        .collect();
351
352    let (indices, weights) =
353        rust_differentiable_search(query_slice, &candidates_vec, k as usize, temperature as f32);
354
355    Ok(SearchResult {
356        indices: indices.iter().map(|&i| i as u32).collect(),
357        weights: weights.iter().map(|&w| w as f64).collect(),
358    })
359}
360
361/// Hierarchical forward pass through GNN layers
362///
363/// # Arguments
364/// * `query` - The query vector (Float32Array)
365/// * `layer_embeddings` - Embeddings organized by layer (Array of Array of Float32Array)
366/// * `gnn_layers_json` - JSON array of serialized GNN layers
367///
368/// # Returns
369/// Final embedding after hierarchical processing as Float32Array
370///
371/// # Example
372/// ```javascript
373/// const query = new Float32Array([1.0, 0.0]);
374/// const layerEmbeddings = [[new Float32Array([1.0, 0.0]), new Float32Array([0.0, 1.0])]];
375/// const layer1 = new RuvectorLayer(2, 2, 1, 0.0);
376/// const layers = [layer1.toJson()];
377/// const result = hierarchicalForward(query, layerEmbeddings, layers);
378/// ```
379#[napi]
380pub fn hierarchical_forward(
381    query: Float32Array,
382    layer_embeddings: Vec<Vec<Float32Array>>,
383    gnn_layers_json: Vec<String>,
384) -> Result<Float32Array> {
385    let query_slice = query.as_ref();
386
387    let embeddings_f32: Vec<Vec<Vec<f32>>> = layer_embeddings
388        .into_iter()
389        .map(|layer| layer.into_iter().map(|arr| arr.to_vec()).collect())
390        .collect();
391
392    let gnn_layers: Vec<RustRuvectorLayer> = gnn_layers_json
393        .iter()
394        .map(|json| {
395            let layer: RustRuvectorLayer = serde_json::from_str(json).map_err(|e| {
396                Error::new(
397                    Status::InvalidArg,
398                    format!("Layer deserialization error: {}", e),
399                )
400            })?;
401            layer
402                .validate()
403                .map_err(|e| Error::new(Status::InvalidArg, e.to_string()))?;
404            Ok(layer)
405        })
406        .collect::<Result<Vec<_>>>()?;
407
408    // Mirror the core's skipped empty layers and projection dimensions before
409    // entering code that assumes matching vector/matrix shapes.
410    let mut dimension = query.len();
411    for (embeddings, layer) in embeddings_f32.iter().zip(&gnn_layers) {
412        if embeddings.is_empty() {
413            continue;
414        }
415        validate_dimension("layer query", dimension, layer.input_dim())?;
416        for embedding in embeddings {
417            validate_dimension("layer embedding", embedding.len(), dimension)?;
418        }
419        dimension = layer.output_dim();
420    }
421
422    let result = rust_hierarchical_forward(query_slice, &embeddings_f32, &gnn_layers);
423
424    Ok(Float32Array::new(result))
425}
426
427// ==================== Helper Functions ====================
428
429fn validate_dimension(name: &str, actual: usize, expected: usize) -> Result<()> {
430    if actual != expected {
431        return Err(Error::new(
432            Status::InvalidArg,
433            format!("{} dimension must be {}, got {}", name, expected, actual),
434        ));
435    }
436    Ok(())
437}
438
439/// Get the compression level that would be selected for a given access frequency
440///
441/// # Arguments
442/// * `access_freq` - Access frequency in range [0.0, 1.0]
443///
444/// # Returns
445/// String describing the compression level: "none", "half", "pq8", "pq4", or "binary"
446///
447/// # Example
448/// ```javascript
449/// const level = getCompressionLevel(0.9); // "none" (hot data)
450/// const level2 = getCompressionLevel(0.5); // "half" (warm data)
451/// ```
452#[napi]
453pub fn get_compression_level(access_freq: f64) -> String {
454    if access_freq > 0.8 {
455        "none".to_string()
456    } else if access_freq > 0.4 {
457        "half".to_string()
458    } else if access_freq > 0.1 {
459        "pq8".to_string()
460    } else if access_freq > 0.01 {
461        "pq4".to_string()
462    } else {
463        "binary".to_string()
464    }
465}
466
467/// Module initialization
468#[napi]
469pub fn init() -> String {
470    "Ruvector GNN Node.js bindings initialized".to_string()
471}
ruvector_gnn_node/lib.rs

ruvector_gnn_node/
lib.rs
