clawft_core/embeddings/
api_embedder.rs1use std::collections::HashMap;
11
12use async_trait::async_trait;
13use sha2::{Digest, Sha256};
14use tracing::{debug, warn};
15
16use super::{Embedder, EmbeddingError};
17
18const DEFAULT_DIMENSION: usize = 384;
20
21#[derive(Debug, Clone)]
23pub struct ApiEmbedderConfig {
24 pub base_url: String,
26 pub api_key_env: String,
28 pub model: String,
30 pub dimension: Option<usize>,
32 pub headers: HashMap<String, String>,
34}
35
36impl Default for ApiEmbedderConfig {
37 fn default() -> Self {
38 Self {
39 base_url: "https://api.openai.com/v1".into(),
40 api_key_env: "OPENAI_API_KEY".into(),
41 model: "text-embedding-3-small".into(),
42 dimension: Some(DEFAULT_DIMENSION),
43 headers: HashMap::new(),
44 }
45 }
46}
47
48pub struct ApiEmbedder {
62 config: ApiEmbedderConfig,
63 http: reqwest::Client,
64 dim: usize,
65}
66
67impl ApiEmbedder {
68 pub fn new(config: ApiEmbedderConfig) -> Self {
70 let dim = config.dimension.unwrap_or(DEFAULT_DIMENSION);
71 Self {
72 config,
73 http: reqwest::Client::new(),
74 dim,
75 }
76 }
77
78 pub fn with_defaults() -> Self {
80 Self::new(ApiEmbedderConfig::default())
81 }
82
83 pub fn hash_only(dimension: usize) -> Self {
87 Self::new(ApiEmbedderConfig {
88 api_key_env: String::new(), dimension: Some(dimension),
90 ..ApiEmbedderConfig::default()
91 })
92 }
93
94 fn resolve_api_key(&self) -> Option<String> {
96 if self.config.api_key_env.is_empty() {
97 return None;
98 }
99 std::env::var(&self.config.api_key_env).ok()
100 }
101
102 async fn call_api(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
104 let api_key = self.resolve_api_key().ok_or_else(|| {
105 EmbeddingError::Internal(format!("API key not set: {}", self.config.api_key_env))
106 })?;
107
108 let base = self.config.base_url.trim_end_matches('/');
109 let url = format!("{base}/embeddings");
110
111 let mut body = serde_json::json!({
112 "input": texts,
113 "model": self.config.model,
114 });
115
116 if let Some(dim) = self.config.dimension {
117 body["dimensions"] = serde_json::json!(dim);
118 }
119
120 let mut req = self
121 .http
122 .post(&url)
123 .header("Authorization", format!("Bearer {api_key}"))
124 .header("Content-Type", "application/json");
125
126 for (k, v) in &self.config.headers {
127 req = req.header(k.as_str(), v.as_str());
128 }
129
130 let response = req
131 .json(&body)
132 .send()
133 .await
134 .map_err(|e| EmbeddingError::Internal(format!("HTTP request failed: {e}")))?;
135
136 if !response.status().is_success() {
137 let status = response.status();
138 let body_text = response.text().await.unwrap_or_default();
139 return Err(EmbeddingError::Internal(format!(
140 "API error HTTP {status}: {body_text}"
141 )));
142 }
143
144 let resp_json: serde_json::Value = response
145 .json()
146 .await
147 .map_err(|e| EmbeddingError::Internal(format!("failed to parse API response: {e}")))?;
148
149 let data = resp_json
151 .get("data")
152 .and_then(|d| d.as_array())
153 .ok_or_else(|| EmbeddingError::Internal("response missing 'data' array".into()))?;
154
155 let mut results = Vec::with_capacity(data.len());
156 for item in data {
157 let embedding = item
158 .get("embedding")
159 .and_then(|e| e.as_array())
160 .ok_or_else(|| {
161 EmbeddingError::Internal("response item missing 'embedding'".into())
162 })?;
163
164 let vec: Vec<f32> = embedding
165 .iter()
166 .map(|v| v.as_f64().unwrap_or(0.0) as f32)
167 .collect();
168
169 results.push(vec);
170 }
171
172 Ok(results)
173 }
174
175 pub fn hash_fallback(&self, text: &str) -> Vec<f32> {
181 let normalized = text.to_lowercase();
182 let mut vector = vec![0.0f32; self.dim];
183
184 if normalized.trim().is_empty() {
185 return vector;
186 }
187
188 let mut hasher = Sha256::new();
190 hasher.update(normalized.as_bytes());
191 let hash = hasher.finalize();
192
193 for (i, val) in vector.iter_mut().enumerate() {
195 let byte_idx = i % hash.len();
196 let bit_idx = (i / hash.len()) % 8;
197 let byte = hash[byte_idx];
198 if (byte >> bit_idx) & 1 == 1 {
199 *val = 1.0;
200 } else {
201 *val = -1.0;
202 }
203 }
204
205 let norm: f32 = vector.iter().map(|x| x * x).sum::<f32>().sqrt();
207 if norm > 0.0 {
208 for val in &mut vector {
209 *val /= norm;
210 }
211 }
212
213 vector
214 }
215}
216
217#[async_trait]
218impl Embedder for ApiEmbedder {
219 async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
220 match self.call_api(&[text]).await {
222 Ok(mut results) if !results.is_empty() => {
223 debug!(model = %self.config.model, "API embedding succeeded");
224 Ok(results.remove(0))
225 }
226 Ok(_) => {
227 warn!("API returned empty embeddings, using hash fallback");
228 Ok(self.hash_fallback(text))
229 }
230 Err(e) => {
231 debug!(error = %e, "API embedding failed, using hash fallback");
232 Ok(self.hash_fallback(text))
233 }
234 }
235 }
236
237 async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
238 let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
239
240 match self.call_api(&text_refs).await {
241 Ok(results) if results.len() == texts.len() => {
242 debug!(
243 model = %self.config.model,
244 count = texts.len(),
245 "batch API embedding succeeded"
246 );
247 Ok(results)
248 }
249 Ok(_) | Err(_) => {
250 debug!("batch API embedding failed, using hash fallback");
251 Ok(texts.iter().map(|t| self.hash_fallback(t)).collect())
252 }
253 }
254 }
255
256 fn dimension(&self) -> usize {
257 self.dim
258 }
259
260 fn name(&self) -> &str {
261 &self.config.model
262 }
263}
264
265#[cfg(test)]
266mod tests {
267 use super::*;
268
269 #[test]
270 fn hash_fallback_deterministic() {
271 let embedder = ApiEmbedder::hash_only(384);
272 let e1 = embedder.hash_fallback("hello world");
273 let e2 = embedder.hash_fallback("hello world");
274 assert_eq!(e1, e2, "same text must produce identical hash embeddings");
275 }
276
277 #[test]
278 fn hash_fallback_different_text() {
279 let embedder = ApiEmbedder::hash_only(384);
280 let e1 = embedder.hash_fallback("hello world");
281 let e2 = embedder.hash_fallback("goodbye moon");
282 assert_ne!(e1, e2, "different text should produce different embeddings");
283 }
284
285 #[test]
286 fn hash_fallback_correct_dimension() {
287 let embedder = ApiEmbedder::hash_only(128);
288 let emb = embedder.hash_fallback("test text");
289 assert_eq!(emb.len(), 128);
290 }
291
292 #[test]
293 fn hash_fallback_unit_norm() {
294 let embedder = ApiEmbedder::hash_only(384);
295 let emb = embedder.hash_fallback("the quick brown fox");
296 let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
297 assert!(
298 (norm - 1.0).abs() < 0.01,
299 "L2 norm should be ~1.0, got {norm}"
300 );
301 }
302
303 #[test]
304 fn hash_fallback_empty_string() {
305 let embedder = ApiEmbedder::hash_only(384);
306 let emb = embedder.hash_fallback("");
307 assert_eq!(emb.len(), 384);
308 let sum: f32 = emb.iter().map(|x| x.abs()).sum();
309 assert!(
310 sum < f32::EPSILON,
311 "empty string should produce zero vector, sum={sum}"
312 );
313 }
314
315 #[test]
316 fn hash_fallback_case_insensitive() {
317 let embedder = ApiEmbedder::hash_only(384);
318 let e1 = embedder.hash_fallback("Hello World");
319 let e2 = embedder.hash_fallback("hello world");
320 assert_eq!(e1, e2, "hash fallback should be case-insensitive");
321 }
322
323 #[test]
324 fn default_config() {
325 let config = ApiEmbedderConfig::default();
326 assert_eq!(config.base_url, "https://api.openai.com/v1");
327 assert_eq!(config.api_key_env, "OPENAI_API_KEY");
328 assert_eq!(config.model, "text-embedding-3-small");
329 assert_eq!(config.dimension, Some(384));
330 }
331
332 #[test]
333 fn with_defaults_creates_embedder() {
334 let embedder = ApiEmbedder::with_defaults();
335 assert_eq!(embedder.dimension(), 384);
336 }
337
338 #[test]
339 fn hash_only_creates_embedder() {
340 let embedder = ApiEmbedder::hash_only(256);
341 assert_eq!(embedder.dimension(), 256);
342 }
343
344 #[tokio::test]
345 async fn embed_falls_back_without_api_key() {
346 let embedder = ApiEmbedder::hash_only(64);
348 let result = embedder.embed("test text").await;
349 assert!(result.is_ok());
350 let emb = result.unwrap();
351 assert_eq!(emb.len(), 64);
352 }
353
354 #[tokio::test]
355 async fn embed_batch_falls_back_without_api_key() {
356 let embedder = ApiEmbedder::hash_only(64);
357 let texts = vec!["hello".to_string(), "world".to_string()];
358 let result = embedder.embed_batch(&texts).await;
359 assert!(result.is_ok());
360 let batch = result.unwrap();
361 assert_eq!(batch.len(), 2);
362 assert_eq!(batch[0].len(), 64);
363 assert_eq!(batch[1].len(), 64);
364 }
365
366 #[tokio::test]
367 async fn embed_batch_consistency_with_single() {
368 let embedder = ApiEmbedder::hash_only(64);
369 let texts = vec!["hello world".to_string(), "goodbye moon".to_string()];
370 let batch = embedder.embed_batch(&texts).await.unwrap();
371
372 let e1 = embedder.embed("hello world").await.unwrap();
373 let e2 = embedder.embed("goodbye moon").await.unwrap();
374
375 assert_eq!(batch[0], e1);
376 assert_eq!(batch[1], e2);
377 }
378
379 #[tokio::test]
380 async fn trait_dimension_matches() {
381 let embedder = ApiEmbedder::hash_only(128);
382 let emb = embedder.embed("test").await.unwrap();
383 assert_eq!(emb.len(), embedder.dimension());
384 }
385}