Skip to main content

clawft_core/embeddings/
api_embedder.rs

1//! API-based embedder using an OpenAI-compatible embeddings endpoint.
2//!
3//! [`ApiEmbedder`] generates embeddings by calling a remote LLM provider's
4//! `/embeddings` endpoint. When no API key is configured (or the call fails),
5//! it falls back to a deterministic SHA-256-based pseudo-embedding for
6//! development and testing.
7//!
8//! This module is gated behind the `rvf` feature flag.
9
10use std::collections::HashMap;
11
12use async_trait::async_trait;
13use sha2::{Digest, Sha256};
14use tracing::{debug, warn};
15
16use super::{Embedder, EmbeddingError};
17
18/// Default embedding dimension when using the hash fallback.
19const DEFAULT_DIMENSION: usize = 384;
20
21/// Configuration for the API embedder.
22#[derive(Debug, Clone)]
23pub struct ApiEmbedderConfig {
24    /// Base URL for the embedding endpoint (e.g. `https://api.openai.com/v1`).
25    pub base_url: String,
26    /// Environment variable holding the API key.
27    pub api_key_env: String,
28    /// The embedding model name (e.g. `text-embedding-3-small`).
29    pub model: String,
30    /// Embedding dimension override (if the provider supports it).
31    pub dimension: Option<usize>,
32    /// Extra HTTP headers for the request.
33    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
48/// Generates embeddings via a remote API endpoint, with hash fallback.
49///
50/// # API Mode
51///
52/// When a valid API key is available, sends requests to the provider's
53/// `/embeddings` endpoint in the OpenAI-compatible format.
54///
55/// # Fallback Mode
56///
57/// When no API key is set or the API call fails, produces deterministic
58/// pseudo-embeddings using SHA-256 hashing. These are NOT semantically
59/// meaningful but provide consistent, reproducible vectors for testing
60/// and development workflows.
61pub struct ApiEmbedder {
62    config: ApiEmbedderConfig,
63    http: reqwest::Client,
64    dim: usize,
65}
66
67impl ApiEmbedder {
68    /// Create a new `ApiEmbedder` with the given configuration.
69    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    /// Create an `ApiEmbedder` with default configuration (OpenAI, 384-dim).
79    pub fn with_defaults() -> Self {
80        Self::new(ApiEmbedderConfig::default())
81    }
82
83    /// Create a fallback-only embedder that never calls an API.
84    ///
85    /// Useful for testing or environments without API access.
86    pub fn hash_only(dimension: usize) -> Self {
87        Self::new(ApiEmbedderConfig {
88            api_key_env: String::new(), // empty => always falls back
89            dimension: Some(dimension),
90            ..ApiEmbedderConfig::default()
91        })
92    }
93
94    /// Resolve the API key from the environment.
95    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    /// Call the remote embeddings API for a single text.
103    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        // Parse OpenAI-format response: { data: [{ embedding: [...] }, ...] }
150        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    /// Generate a deterministic pseudo-embedding from text using SHA-256.
176    ///
177    /// This is NOT semantically meaningful -- it is a fallback for when no
178    /// API is available. The output is deterministic: same input always
179    /// produces the same vector.
180    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        // Hash the full text to seed the vector
189        let mut hasher = Sha256::new();
190        hasher.update(normalized.as_bytes());
191        let hash = hasher.finalize();
192
193        // Fill the vector by cycling through the hash bytes
194        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        // Normalize to unit length
206        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        // Try API first, fall back to hash
221        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        // No API key set, should fall back to hash
347        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}