Skip to main content

ruvector_tiny_dancer_node/
lib.rs

1//! Node.js bindings for Tiny Dancer neural routing via NAPI-RS
2//!
3//! High-performance Rust neural routing with zero-copy buffer sharing,
4//! async/await support, and complete TypeScript type definitions.
5
6#![allow(clippy::all)]
7#![allow(clippy::pedantic)]
8
9use napi::bindgen_prelude::*;
10use napi_derive::napi;
11use parking_lot::RwLock;
12use ruvector_tiny_dancer_core::{
13    types::{
14        Candidate as CoreCandidate, RouterConfig as CoreRouterConfig,
15        RoutingDecision as CoreRoutingDecision, RoutingRequest as CoreRoutingRequest,
16        RoutingResponse as CoreRoutingResponse,
17    },
18    Router as CoreRouter,
19};
20use std::collections::HashMap;
21use std::sync::Arc;
22
23/// Router configuration
24#[napi(object)]
25#[derive(Debug, Clone)]
26pub struct RouterConfig {
27    /// Model path
28    pub model_path: String,
29    /// Confidence threshold (0.0 to 1.0)
30    pub confidence_threshold: Option<f64>,
31    /// Maximum uncertainty (0.0 to 1.0)
32    pub max_uncertainty: Option<f64>,
33    /// Enable circuit breaker
34    pub enable_circuit_breaker: Option<bool>,
35    /// Circuit breaker threshold
36    pub circuit_breaker_threshold: Option<u32>,
37    /// Enable quantization
38    pub enable_quantization: Option<bool>,
39    /// Database path
40    pub database_path: Option<String>,
41}
42
43impl From<RouterConfig> for CoreRouterConfig {
44    fn from(config: RouterConfig) -> Self {
45        CoreRouterConfig {
46            model_path: config.model_path,
47            confidence_threshold: config.confidence_threshold.unwrap_or(0.85) as f32,
48            max_uncertainty: config.max_uncertainty.unwrap_or(0.15) as f32,
49            enable_circuit_breaker: config.enable_circuit_breaker.unwrap_or(true),
50            circuit_breaker_threshold: config.circuit_breaker_threshold.unwrap_or(5),
51            enable_quantization: config.enable_quantization.unwrap_or(true),
52            database_path: config.database_path,
53            // VoI escalation gate (ADR-331) is not yet exposed over NAPI.
54            voi: None,
55        }
56    }
57}
58
59/// Candidate for routing
60#[napi(object)]
61#[derive(Clone)]
62pub struct Candidate {
63    /// Candidate ID
64    pub id: String,
65    /// Embedding vector
66    pub embedding: Float32Array,
67    /// Metadata (JSON string)
68    pub metadata: Option<String>,
69    /// Creation timestamp
70    pub created_at: Option<i64>,
71    /// Access count
72    pub access_count: Option<u32>,
73    /// Success rate (0.0 to 1.0)
74    pub success_rate: Option<f64>,
75}
76
77impl Candidate {
78    fn to_core(&self) -> Result<CoreCandidate> {
79        let metadata: HashMap<String, serde_json::Value> = if let Some(ref meta_str) = self.metadata
80        {
81            serde_json::from_str(meta_str)
82                .map_err(|e| Error::from_reason(format!("Invalid metadata JSON: {}", e)))?
83        } else {
84            HashMap::new()
85        };
86
87        Ok(CoreCandidate {
88            id: self.id.clone(),
89            embedding: self.embedding.to_vec(),
90            metadata,
91            created_at: self
92                .created_at
93                .unwrap_or_else(|| chrono::Utc::now().timestamp()),
94            access_count: self.access_count.unwrap_or(0) as u64,
95            success_rate: self.success_rate.unwrap_or(0.0) as f32,
96        })
97    }
98}
99
100/// Routing request
101#[napi(object)]
102pub struct RoutingRequest {
103    /// Query embedding
104    pub query_embedding: Float32Array,
105    /// Candidates to score
106    pub candidates: Vec<Candidate>,
107    /// Optional metadata (JSON string)
108    pub metadata: Option<String>,
109}
110
111impl RoutingRequest {
112    fn to_core(&self) -> Result<CoreRoutingRequest> {
113        let candidates: Result<Vec<CoreCandidate>> =
114            self.candidates.iter().map(|c| c.to_core()).collect();
115
116        let metadata = if let Some(ref meta_str) = self.metadata {
117            Some(
118                serde_json::from_str(meta_str)
119                    .map_err(|e| Error::from_reason(format!("Invalid metadata JSON: {}", e)))?,
120            )
121        } else {
122            None
123        };
124
125        Ok(CoreRoutingRequest {
126            query_embedding: self.query_embedding.to_vec(),
127            candidates: candidates?,
128            metadata,
129        })
130    }
131}
132
133/// Routing decision
134#[napi(object)]
135#[derive(Debug, Clone)]
136pub struct RoutingDecision {
137    /// Candidate ID
138    pub candidate_id: String,
139    /// Confidence score (0.0 to 1.0)
140    pub confidence: f64,
141    /// Whether to use lightweight model
142    pub use_lightweight: bool,
143    /// Uncertainty estimate (0.0 to 1.0)
144    pub uncertainty: f64,
145}
146
147impl From<CoreRoutingDecision> for RoutingDecision {
148    fn from(decision: CoreRoutingDecision) -> Self {
149        Self {
150            candidate_id: decision.candidate_id,
151            confidence: decision.confidence as f64,
152            use_lightweight: decision.use_lightweight,
153            uncertainty: decision.uncertainty as f64,
154        }
155    }
156}
157
158/// Routing response
159#[napi(object)]
160#[derive(Debug, Clone)]
161pub struct RoutingResponse {
162    /// Routing decisions
163    pub decisions: Vec<RoutingDecision>,
164    /// Total inference time in microseconds
165    pub inference_time_us: u32,
166    /// Number of candidates processed
167    pub candidates_processed: u32,
168    /// Feature engineering time in microseconds
169    pub feature_time_us: u32,
170}
171
172impl From<CoreRoutingResponse> for RoutingResponse {
173    fn from(response: CoreRoutingResponse) -> Self {
174        Self {
175            decisions: response.decisions.into_iter().map(Into::into).collect(),
176            inference_time_us: response.inference_time_us as u32,
177            candidates_processed: response.candidates_processed as u32,
178            feature_time_us: response.feature_time_us as u32,
179        }
180    }
181}
182
183/// Tiny Dancer neural router
184#[napi]
185pub struct Router {
186    inner: Arc<RwLock<CoreRouter>>,
187}
188
189#[napi]
190impl Router {
191    /// Create a new router with configuration
192    ///
193    /// # Example
194    /// ```javascript
195    /// const router = new Router({
196    ///   modelPath: './models/fastgrnn.safetensors',
197    ///   confidenceThreshold: 0.85,
198    ///   maxUncertainty: 0.15,
199    ///   enableCircuitBreaker: true
200    /// });
201    /// ```
202    #[napi(constructor)]
203    pub fn new(config: RouterConfig) -> Result<Self> {
204        let core_config: CoreRouterConfig = config.into();
205        let router = CoreRouter::new(core_config)
206            .map_err(|e| Error::from_reason(format!("Failed to create router: {}", e)))?;
207
208        Ok(Self {
209            inner: Arc::new(RwLock::new(router)),
210        })
211    }
212
213    /// Route a request through the neural routing system
214    ///
215    /// Returns routing decisions with confidence scores and model recommendations
216    ///
217    /// # Example
218    /// ```javascript
219    /// const response = await router.route({
220    ///   queryEmbedding: new Float32Array([0.1, 0.2, ...]),
221    ///   candidates: [
222    ///     { id: '1', embedding: new Float32Array([...]) },
223    ///     { id: '2', embedding: new Float32Array([...]) }
224    ///   ]
225    /// });
226    /// console.log('Top decision:', response.decisions[0]);
227    /// console.log('Inference time:', response.inferenceTimeUs, 'μs');
228    /// ```
229    #[napi]
230    pub async fn route(&self, request: RoutingRequest) -> Result<RoutingResponse> {
231        let core_request = request.to_core()?;
232        let router = self.inner.clone();
233
234        tokio::task::spawn_blocking(move || {
235            let router = router.read();
236            router.route(core_request)
237        })
238        .await
239        .map_err(|e| Error::from_reason(format!("Task failed: {}", e)))?
240        .map_err(|e| Error::from_reason(format!("Routing failed: {}", e)))
241        .map(Into::into)
242    }
243
244    /// Reload the model from disk (hot-reload)
245    ///
246    /// # Example
247    /// ```javascript
248    /// await router.reloadModel();
249    /// ```
250    #[napi]
251    pub async fn reload_model(&self) -> Result<()> {
252        let router = self.inner.clone();
253
254        tokio::task::spawn_blocking(move || {
255            let router = router.read();
256            router.reload_model()
257        })
258        .await
259        .map_err(|e| Error::from_reason(format!("Task failed: {}", e)))?
260        .map_err(|e| Error::from_reason(format!("Model reload failed: {}", e)))
261    }
262
263    /// Check circuit breaker status
264    ///
265    /// Returns true if the circuit is closed (healthy), false if open (unhealthy)
266    ///
267    /// # Example
268    /// ```javascript
269    /// const isHealthy = router.circuitBreakerStatus();
270    /// ```
271    #[napi]
272    pub fn circuit_breaker_status(&self) -> Option<bool> {
273        let router = self.inner.read();
274        router.circuit_breaker_status()
275    }
276}
277
278/// Get the version of the Tiny Dancer library
279#[napi]
280pub fn version() -> String {
281    env!("CARGO_PKG_VERSION").to_string()
282}
283
284/// Hello function for testing bindings
285#[napi]
286pub fn hello() -> String {
287    "Hello from Tiny Dancer Node.js bindings!".to_string()
288}
289
290/// One DRACO training row: a query embedding and the quality each model achieved
291/// on it (model id → quality, 0..1). Matches `@metaharness/router`'s row shape.
292#[napi(object)]
293pub struct DracoRowJs {
294    pub embedding: Vec<f64>,
295    pub scores: std::collections::HashMap<String, f64>,
296}
297
298/// Options for `trainRouter`.
299#[napi(object)]
300pub struct TrainRouterOptions {
301    /// Where to write the trained `.safetensors` model.
302    pub output_path: String,
303    /// Input feature dimension (must equal the embedding length).
304    pub input_dim: u32,
305    /// Hidden dimension (default 12).
306    pub hidden_dim: Option<u32>,
307    /// Training epochs (default 40).
308    pub epochs: Option<u32>,
309    /// Learning rate (default 0.05).
310    pub learning_rate: Option<f64>,
311    /// DRACO label tolerance: cheap model is "good enough" within this of the best
312    /// (default 0.05).
313    pub tolerance: Option<f64>,
314}
315
316/// Result of `trainRouter`.
317#[napi(object)]
318pub struct TrainRouterResult {
319    pub epochs_run: u32,
320    pub train_loss: f64,
321    pub train_accuracy: f64,
322    pub val_accuracy: f64,
323    pub model_path: String,
324    pub model_bytes: u32,
325}
326
327/// Train a FastGRNN router from a DRACO dataset and write it to a
328/// `.safetensors` file consumable by `new Router({ modelPath })`.
329///
330/// ```javascript
331/// const res = await trainRouter(rows, { haiku: 1, opus: 15 }, {
332///   outputPath: './router.safetensors', inputDim: 8, epochs: 40,
333/// });
334/// const router = new Router({ modelPath: res.modelPath });
335/// ```
336#[napi]
337pub async fn train_router(
338    rows: Vec<DracoRowJs>,
339    prices: std::collections::HashMap<String, f64>,
340    options: TrainRouterOptions,
341) -> Result<TrainRouterResult> {
342    use ruvector_tiny_dancer_core::model::{FastGRNN, FastGRNNConfig};
343    use ruvector_tiny_dancer_core::training::{DracoRow, Trainer, TrainingConfig, TrainingDataset};
344
345    tokio::task::spawn_blocking(move || -> std::result::Result<TrainRouterResult, String> {
346        let core_rows: Vec<DracoRow> = rows
347            .into_iter()
348            .map(|r| DracoRow {
349                embedding: r.embedding.into_iter().map(|v| v as f32).collect(),
350                scores: r.scores.into_iter().map(|(k, v)| (k, v as f32)).collect(),
351            })
352            .collect();
353        let core_prices: std::collections::HashMap<String, f32> =
354            prices.into_iter().map(|(k, v)| (k, v as f32)).collect();
355
356        let tolerance = options.tolerance.unwrap_or(0.05) as f32;
357        let dataset = TrainingDataset::from_draco(&core_rows, &core_prices, tolerance)
358            .map_err(|e| format!("dataset: {e}"))?;
359
360        let model_config = FastGRNNConfig {
361            input_dim: options.input_dim as usize,
362            hidden_dim: options.hidden_dim.unwrap_or(12) as usize,
363            output_dim: 1,
364            ..Default::default()
365        };
366        let train_config = TrainingConfig {
367            learning_rate: options.learning_rate.unwrap_or(0.05) as f32,
368            epochs: options.epochs.unwrap_or(40) as usize,
369            early_stopping_patience: None,
370            l2_reg: 0.0,
371            ..Default::default()
372        };
373
374        let mut model = FastGRNN::new(model_config.clone()).map_err(|e| format!("model: {e}"))?;
375        let metrics = Trainer::new(&model_config, train_config)
376            .train(&mut model, &dataset)
377            .map_err(|e| format!("train: {e}"))?;
378        model
379            .save(&options.output_path)
380            .map_err(|e| format!("save: {e}"))?;
381
382        let last = metrics.last().ok_or_else(|| "no metrics".to_string())?;
383        let model_bytes = std::fs::metadata(&options.output_path)
384            .map(|m| m.len() as u32)
385            .unwrap_or(0);
386
387        Ok(TrainRouterResult {
388            epochs_run: metrics.len() as u32,
389            train_loss: last.train_loss as f64,
390            train_accuracy: last.train_accuracy as f64,
391            val_accuracy: last.val_accuracy as f64,
392            model_path: options.output_path,
393            model_bytes,
394        })
395    })
396    .await
397    .map_err(|e| Error::from_reason(format!("Task failed: {}", e)))?
398    .map_err(Error::from_reason)
399}
400
401/// Score a query embedding with a trained FastGRNN model (raw forward pass).
402///
403/// Loads the `.safetensors` produced by {@link train_router} and runs the model
404/// directly on `embedding` (which must match the model's `input_dim`). Returns
405/// the sigmoid output in 0..1 — high means "the cheap model is good enough"
406/// (route to the cheaper model); low means route to a stronger model.
407///
408/// This is the inference path that matches `trainRouter` (trained on raw
409/// embeddings); it does not run `Router`'s feature engineering.
410#[napi]
411pub async fn score(model_path: String, embedding: Vec<f64>) -> Result<f64> {
412    use ruvector_tiny_dancer_core::model::FastGRNN;
413
414    tokio::task::spawn_blocking(move || -> std::result::Result<f64, String> {
415        let model = FastGRNN::load(&model_path).map_err(|e| format!("load: {e}"))?;
416        let feats: Vec<f32> = embedding.into_iter().map(|v| v as f32).collect();
417        let s = model
418            .forward(&feats, None)
419            .map_err(|e| format!("forward: {e}"))?;
420        Ok(s as f64)
421    })
422    .await
423    .map_err(|e| Error::from_reason(format!("Task failed: {}", e)))?
424    .map_err(Error::from_reason)
425}