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}