ruvector_tiny_dancer_node/
lib.rs1#![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#[napi(object)]
25#[derive(Debug, Clone)]
26pub struct RouterConfig {
27 pub model_path: String,
29 pub confidence_threshold: Option<f64>,
31 pub max_uncertainty: Option<f64>,
33 pub enable_circuit_breaker: Option<bool>,
35 pub circuit_breaker_threshold: Option<u32>,
37 pub enable_quantization: Option<bool>,
39 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: None,
55 }
56 }
57}
58
59#[napi(object)]
61#[derive(Clone)]
62pub struct Candidate {
63 pub id: String,
65 pub embedding: Float32Array,
67 pub metadata: Option<String>,
69 pub created_at: Option<i64>,
71 pub access_count: Option<u32>,
73 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#[napi(object)]
102pub struct RoutingRequest {
103 pub query_embedding: Float32Array,
105 pub candidates: Vec<Candidate>,
107 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#[napi(object)]
135#[derive(Debug, Clone)]
136pub struct RoutingDecision {
137 pub candidate_id: String,
139 pub confidence: f64,
141 pub use_lightweight: bool,
143 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#[napi(object)]
160#[derive(Debug, Clone)]
161pub struct RoutingResponse {
162 pub decisions: Vec<RoutingDecision>,
164 pub inference_time_us: u32,
166 pub candidates_processed: u32,
168 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#[napi]
185pub struct Router {
186 inner: Arc<RwLock<CoreRouter>>,
187}
188
189#[napi]
190impl Router {
191 #[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 #[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 #[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 #[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#[napi]
280pub fn version() -> String {
281 env!("CARGO_PKG_VERSION").to_string()
282}
283
284#[napi]
286pub fn hello() -> String {
287 "Hello from Tiny Dancer Node.js bindings!".to_string()
288}
289
290#[napi(object)]
293pub struct DracoRowJs {
294 pub embedding: Vec<f64>,
295 pub scores: std::collections::HashMap<String, f64>,
296}
297
298#[napi(object)]
300pub struct TrainRouterOptions {
301 pub output_path: String,
303 pub input_dim: u32,
305 pub hidden_dim: Option<u32>,
307 pub epochs: Option<u32>,
309 pub learning_rate: Option<f64>,
311 pub tolerance: Option<f64>,
314}
315
316#[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#[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#[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}