#![allow(clippy::missing_errors_doc)]
use wasm_bindgen::prelude::*;
use crate::layer1_echo::{EchoLayer, InMemoryVectorStore, MockEmbeddingProvider};
use crate::layer2_speculator::RuleBasedSpeculator;
use crate::layer3_judge::{AdvancedClaimExtractor, JudgeConfig, JudgeImpl, MockSmtVerifier};
use crate::pipeline::{Pipeline, PipelineConfig};
use crate::types::{Document, Query};
#[wasm_bindgen(start)]
pub fn init() {
#[cfg(feature = "wasm")]
{
console_error_panic_hook::set_once();
}
}
#[wasm_bindgen]
pub struct WasmRagEngine {
pipeline: Pipeline<
EchoLayer<MockEmbeddingProvider, InMemoryVectorStore>,
RuleBasedSpeculator,
JudgeImpl<AdvancedClaimExtractor, MockSmtVerifier>,
>,
}
#[wasm_bindgen]
impl WasmRagEngine {
#[wasm_bindgen(constructor)]
#[must_use]
pub fn new(dimension: usize) -> Self {
let echo = EchoLayer::new(
MockEmbeddingProvider::new(dimension),
InMemoryVectorStore::new(dimension),
);
let speculator = RuleBasedSpeculator::default();
let judge = JudgeImpl::new(
AdvancedClaimExtractor::new(),
MockSmtVerifier::default(),
JudgeConfig::default(),
);
let config = PipelineConfig::default();
Self {
pipeline: Pipeline::new(echo, speculator, judge, config),
}
}
pub async fn index(
&mut self,
content: String,
title: Option<String>,
) -> Result<String, JsValue> {
let mut doc = Document::new(content);
if let Some(t) = title {
doc = doc.with_title(t);
}
let id = doc.id.clone();
crate::layer1_echo::Echo::index(self.pipeline.echo_mut(), doc)
.await
.map_err(|e| JsValue::from_str(&e.to_string()))?;
Ok(id.to_string())
}
pub async fn query(&self, query_text: String, top_k: usize) -> Result<String, JsValue> {
let query = Query::new(query_text).with_top_k(top_k);
let results =
crate::layer1_echo::Echo::search(self.pipeline.echo(), &query.text, top_k, None)
.await
.map_err(|e| JsValue::from_str(&e.to_string()))?;
let json =
serde_json::to_string(&results).map_err(|e| JsValue::from_str(&e.to_string()))?;
Ok(json)
}
pub async fn count(&self) -> usize {
crate::layer1_echo::Echo::count(self.pipeline.echo()).await
}
pub async fn clear(&mut self) -> Result<(), JsValue> {
crate::layer1_echo::Echo::clear(self.pipeline.echo_mut())
.await
.map_err(|e| JsValue::from_str(&e.to_string()))
}
}
#[wasm_bindgen]
#[must_use]
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
crate::layer1_echo::similarity::cosine_similarity(a, b)
}
#[wasm_bindgen]
#[must_use]
pub fn normalize_vector(v: &[f32]) -> Vec<f32> {
crate::layer1_echo::similarity::normalize(v)
}
#[wasm_bindgen]
pub fn parse_config(json: &str) -> Result<JsValue, JsValue> {
let config: crate::config::OxiRagConfig =
serde_json::from_str(json).map_err(|e| JsValue::from_str(&e.to_string()))?;
serde_wasm_bindgen::to_value(&config).map_err(|e| JsValue::from_str(&e.to_string()))
}
#[wasm_bindgen]
#[must_use]
pub fn version() -> String {
env!("CARGO_PKG_VERSION").to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use wasm_bindgen_test::*;
wasm_bindgen_test_configure!(run_in_browser);
#[wasm_bindgen_test]
async fn test_wasm_engine_creation() {
let engine = WasmRagEngine::new(64);
assert_eq!(engine.count().await, 0);
}
#[wasm_bindgen_test]
fn test_cosine_similarity() {
let a = vec![1.0, 0.0];
let b = vec![1.0, 0.0];
let sim = cosine_similarity(&a, &b);
assert!((sim - 1.0).abs() < 1e-6);
}
#[wasm_bindgen_test]
fn test_version() {
let v = version();
assert!(!v.is_empty());
}
}