oxios_kernel/embedding/
api.rs

1//! API-based embedding provider (OpenAI-compatible).
2//!
3//! Calls a remote embedding endpoint (e.g. `POST /v1/embeddings`) and returns
4//! dense `f32` vectors suitable for sqlite-vec KNN and HNSW ANN search.
5//!
6//! ## Why this exists
7//!
8//! The default `TfIdfEmbeddingProvider` produces sparse vectors that
9//! `EmbeddingVector::to_f32_dense()` cannot convert to f32 (`embedding.rs:99`
10//! returns `None` for `Sparse`), so `SqliteMemoryStore::remember()` silently
11//! skips the vector insert — `memory_vectors_rowids` stays empty.
12//!
13//! `GgufEmbeddingProvider` (feature `embedding-gguf`) is aarch64-only and
14//! requires a 329MB model download.
15//!
16//! API embeddings: zero-dep (reqwest already in tree), cross-platform, and
17//! the user already has API keys configured for LLM providers.
18//!
19//! ## Config
20//!
21//! ```toml
22//! [embedding]
23//! endpoint = "https://api.openai.com/v1/embeddings"
24//! api_key  = ""               # empty → inherit from active LLM provider
25//! model    = "text-embedding-3-small"
26//! # dimensions = 1536         # optional; defaults per model
27//! ```
28//!
29//! ## Failure handling
30//!
31//! Network errors / non-2xx responses return `Err`. Callers in the write path
32//! (`SqliteMemoryStore::remember`) must treat this as non-fatal — store the
33//! text+FTS5 row and skip the vector. See Phase 2b in the design doc.
34
35use std::sync::Arc;
36use std::time::Duration;
37
38use anyhow::{Context, Result};
39use async_trait::async_trait;
40use parking_lot::Mutex;
41use serde::{Deserialize, Serialize};
42
43use crate::embedding::{EmbeddingProvider, EmbeddingVector};
44
45/// Default well-known models and their output dimensions.
46///
47/// Used when the user does not specify `dimensions` in config.
48pub fn default_dimensions(model: &str) -> Option<usize> {
49    match model {
50        "text-embedding-3-small" => Some(1536),
51        "text-embedding-3-large" => Some(3072),
52        "text-embedding-ada-002" => Some(1536),
53        _ => None,
54    }
55}
56
57/// HTTP client wrapper (single connection pool, configurable timeout).
58#[derive(Debug)]
59struct ApiClient {
60    inner: reqwest::Client,
61}
62
63impl ApiClient {
64    fn new() -> Result<Self> {
65        let inner = reqwest::Client::builder()
66            .timeout(Duration::from_secs(15))
67            .connect_timeout(Duration::from_secs(5))
68            .build()
69            .context("Failed to construct reqwest client for embedding API")?;
70        Ok(Self { inner })
71    }
72}
73
74/// Request body for `POST /v1/embeddings` (OpenAI-compatible).
75#[derive(Debug, Serialize)]
76struct EmbeddingRequest<'a> {
77    input: &'a str,
78    model: &'a str,
79    #[serde(skip_serializing_if = "Option::is_none")]
80    dimensions: Option<usize>,
81}
82
83/// Response body for `POST /v1/embeddings`.
84#[derive(Debug, Deserialize)]
85struct EmbeddingResponse {
86    /// List of embeddings, one per input. We send a single input.
87    data: Vec<EmbeddingData>,
88}
89
90#[derive(Debug, Deserialize)]
91struct EmbeddingData {
92    embedding: Vec<f32>,
93}
94
95/// API-based embedding provider.
96///
97/// Created at boot from `[embedding]` config. If config is unset, this
98/// provider is never instantiated and `TfIdfEmbeddingProvider` is used.
99pub struct ApiEmbeddingProvider {
100    /// HTTP endpoint, e.g. `https://api.openai.com/v1/embeddings`.
101    endpoint: String,
102    /// API key (Bearer token).
103    api_key: String,
104    /// Model name, e.g. `text-embedding-3-small`.
105    model: String,
106    /// Output dimensionality. Required for sqlite-vec table creation.
107    /// Resolved from config or `default_dimensions()`.
108    dimensions: usize,
109    /// HTTP client (lazy-initialized via Mutex on first use to avoid
110    /// propagating Client build errors through construction).
111    client: Mutex<Option<Arc<ApiClient>>>,
112}
113
114impl std::fmt::Debug for ApiEmbeddingProvider {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        f.debug_struct("ApiEmbeddingProvider")
117            .field("endpoint", &self.endpoint)
118            .field("model", &self.model)
119            .field("dimensions", &self.dimensions)
120            .finish_non_exhaustive()
121    }
122}
123
124impl ApiEmbeddingProvider {
125    /// Construct a new provider from the resolved config fields.
126    ///
127    /// Returns `Err` if `dimensions` cannot be determined.
128    pub fn new(
129        endpoint: String,
130        api_key: String,
131        model: String,
132        dimensions: Option<usize>,
133    ) -> Result<Self> {
134        let dimensions = match dimensions {
135            Some(d) if d > 0 => d,
136            _ => default_dimensions(&model)
137                .ok_or_else(|| anyhow::anyhow!(
138                    "Embedding model '{model}' has no known dimensionality; specify `dimensions` in [embedding] config"
139                ))?,
140        };
141        Ok(Self {
142            endpoint,
143            api_key,
144            model,
145            dimensions,
146            client: Mutex::new(None),
147        })
148    }
149
150    /// Get or lazily build the HTTP client.
151    fn http(&self) -> Result<Arc<ApiClient>> {
152        let mut guard = self.client.lock();
153        if let Some(c) = guard.as_ref() {
154            return Ok(Arc::clone(c));
155        }
156        let c = Arc::new(ApiClient::new()?);
157        *guard = Some(Arc::clone(&c));
158        Ok(c)
159    }
160}
161
162#[async_trait]
163impl EmbeddingProvider for ApiEmbeddingProvider {
164    async fn embed(&self, text: &str) -> Result<EmbeddingVector> {
165        let client = self.http()?;
166        let req = EmbeddingRequest {
167            input: text,
168            model: &self.model,
169            dimensions: None, // most providers ignore or default to model size
170        };
171        let resp = client
172            .inner
173            .post(&self.endpoint)
174            .bearer_auth(&self.api_key)
175            .json(&req)
176            .send()
177            .await
178            .context("Embedding API request failed")?;
179
180        let status = resp.status();
181        let body = resp
182            .text()
183            .await
184            .context("Failed to read embedding API response body")?;
185
186        if !status.is_success() {
187            anyhow::bail!(
188                "Embedding API returned HTTP {status}: {}",
189                truncate(&body, 200)
190            );
191        }
192
193        let parsed: EmbeddingResponse = serde_json::from_str(&body)
194            .context("Failed to parse embedding API response as JSON")?;
195
196        let vec = parsed
197            .data
198            .into_iter()
199            .next()
200            .ok_or_else(|| anyhow::anyhow!("Embedding API returned no vectors in response"))?
201            .embedding;
202
203        if vec.len() != self.dimensions {
204            anyhow::bail!(
205                "Embedding dimension mismatch: configured {} got {}",
206                self.dimensions,
207                vec.len()
208            );
209        }
210
211        Ok(EmbeddingVector::DenseF32(vec))
212    }
213
214    fn name(&self) -> &str {
215        "api-embedding"
216    }
217}
218
219/// Truncate a string for error messages without panicking on UTF-8 boundaries.
220fn truncate(s: &str, max_bytes: usize) -> String {
221    if s.len() <= max_bytes {
222        return s.to_string();
223    }
224    let mut end = max_bytes;
225    while end > 0 && !s.is_char_boundary(end) {
226        end -= 1;
227    }
228    format!("{}...", &s[..end])
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn test_default_dimensions() {
237        assert_eq!(default_dimensions("text-embedding-3-small"), Some(1536));
238        assert_eq!(default_dimensions("text-embedding-3-large"), Some(3072));
239        assert_eq!(default_dimensions("unknown-model"), None);
240    }
241
242    #[test]
243    fn test_truncate_ascii() {
244        assert_eq!(truncate("hello world", 5), "hello...");
245    }
246
247    #[test]
248    fn test_truncate_korean_safe_boundary() {
249        // "가나다라마" is 10 bytes; truncate at 7 should not panic
250        let s = "가나다라마";
251        let _ = truncate(s, 7);
252    }
253
254    #[test]
255    fn test_constructor_requires_known_dim() {
256        let p = ApiEmbeddingProvider::new(
257            "https://example.com".into(),
258            "k".into(),
259            "text-embedding-3-small".into(),
260            None,
261        )
262        .unwrap();
263        assert_eq!(p.dimensions, 1536);
264    }
265
266    #[test]
267    fn test_constructor_explicit_dim() {
268        let p = ApiEmbeddingProvider::new(
269            "https://example.com".into(),
270            "k".into(),
271            "custom-model".into(),
272            Some(768),
273        )
274        .unwrap();
275        assert_eq!(p.dimensions, 768);
276    }
277
278    #[test]
279    fn test_constructor_unknown_model_no_dim_fails() {
280        let r = ApiEmbeddingProvider::new(
281            "https://example.com".into(),
282            "k".into(),
283            "custom-model".into(),
284            None,
285        );
286        assert!(r.is_err());
287    }
288}
oxios_kernel/embedding/api.rs

oxios_kernel/embedding/
api.rs
