Skip to main content

synaptic_voyage/
lib.rs

1pub mod reranker;
2pub use reranker::{VoyageReranker, VoyageRerankerModel};
3
4use async_trait::async_trait;
5use serde_json::json;
6use synaptic_core::{Embeddings, SynapticError};
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub enum VoyageModel {
10    Voyage3Large,
11    Voyage3,
12    Voyage3Lite,
13    VoyageCode3,
14    VoyageFinance2,
15    Custom(String),
16}
17
18impl VoyageModel {
19    pub fn as_str(&self) -> &str {
20        match self {
21            VoyageModel::Voyage3Large => "voyage-3-large",
22            VoyageModel::Voyage3 => "voyage-3",
23            VoyageModel::Voyage3Lite => "voyage-3-lite",
24            VoyageModel::VoyageCode3 => "voyage-code-3",
25            VoyageModel::VoyageFinance2 => "voyage-finance-2",
26            VoyageModel::Custom(s) => s.as_str(),
27        }
28    }
29}
30
31impl std::fmt::Display for VoyageModel {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(f, "{}", self.as_str())
34    }
35}
36
37#[derive(Debug, Clone)]
38pub struct VoyageConfig {
39    pub api_key: String,
40    pub model: String,
41    pub base_url: String,
42    pub input_type: Option<String>,
43}
44
45impl VoyageConfig {
46    pub fn new(api_key: impl Into<String>, model: VoyageModel) -> Self {
47        Self {
48            api_key: api_key.into(),
49            model: model.to_string(),
50            base_url: "https://api.voyageai.com/v1".to_string(),
51            input_type: None,
52        }
53    }
54
55    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
56        self.base_url = url.into();
57        self
58    }
59
60    pub fn with_input_type(mut self, t: impl Into<String>) -> Self {
61        self.input_type = Some(t.into());
62        self
63    }
64}
65
66pub struct VoyageEmbeddings {
67    config: VoyageConfig,
68    client: reqwest::Client,
69}
70
71impl VoyageEmbeddings {
72    pub fn new(config: VoyageConfig) -> Self {
73        Self {
74            config,
75            client: reqwest::Client::new(),
76        }
77    }
78
79    async fn embed_batch(
80        &self,
81        texts: &[&str],
82        input_type: Option<&str>,
83    ) -> Result<Vec<Vec<f32>>, SynapticError> {
84        let mut body = json!({
85            "model": self.config.model,
86            "input": texts,
87        });
88        let itype = input_type.or(self.config.input_type.as_deref());
89        if let Some(t) = itype {
90            body["input_type"] = json!(t);
91        }
92        let resp = self
93            .client
94            .post(format!("{}/embeddings", self.config.base_url))
95            .header("Authorization", format!("Bearer {}", self.config.api_key))
96            .header("Content-Type", "application/json")
97            .json(&body)
98            .send()
99            .await
100            .map_err(|e| SynapticError::Embedding(format!("Voyage request: {e}")))?;
101        let status = resp.status().as_u16();
102        let json: serde_json::Value = resp
103            .json()
104            .await
105            .map_err(|e| SynapticError::Embedding(format!("Voyage parse: {e}")))?;
106        if status != 200 {
107            return Err(SynapticError::Embedding(format!(
108                "Voyage API error ({}): {}",
109                status, json
110            )));
111        }
112        parse_voyage_response(&json)
113    }
114}
115
116fn parse_voyage_response(body: &serde_json::Value) -> Result<Vec<Vec<f32>>, SynapticError> {
117    let data = body
118        .get("data")
119        .and_then(|d| d.as_array())
120        .ok_or_else(|| SynapticError::Embedding("missing 'data' field".to_string()))?;
121    let mut result = Vec::with_capacity(data.len());
122    for item in data {
123        let emb = item
124            .get("embedding")
125            .and_then(|e| e.as_array())
126            .ok_or_else(|| SynapticError::Embedding("missing 'embedding' field".to_string()))?
127            .iter()
128            .map(|v| v.as_f64().unwrap_or(0.0) as f32)
129            .collect();
130        result.push(emb);
131    }
132    Ok(result)
133}
134
135#[async_trait]
136impl Embeddings for VoyageEmbeddings {
137    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, SynapticError> {
138        self.embed_batch(texts, Some("document")).await
139    }
140
141    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, SynapticError> {
142        let mut results = self.embed_batch(&[text], Some("query")).await?;
143        results
144            .pop()
145            .ok_or_else(|| SynapticError::Embedding("empty response".to_string()))
146    }
147}