kb/
embedding.rs

1//! OpenAI-compatible embedding generation.
2//!
3//! Mirrors the Haskell `KB.Embedding` module: configuration via the
4//! `KB_EMBEDDING_*` environment variables, an async client trait used by
5//! the HTTP and MCP handlers, and a packed little-endian `float32` BLOB
6//! codec compatible with `sqlite-vec`.
7//!
8//! When `KB_EMBEDDING_BASE_URL` is unset [`read_embedding_config_from_env`]
9//! returns `None`, which signals "embedding disabled" to the handler
10//! layer. Failures inside [`HttpEmbeddingClient::embed`] are reported as
11//! [`EmbeddingError`]; they never panic.
12
13#![allow(
14    clippy::significant_drop_tightening,
15    reason = "lock/connection/server guards are intentionally held for the operation's duration; early drop would break atomicity"
16)]
17
18use std::env;
19use std::time::Duration;
20
21use async_trait::async_trait;
22use serde::Deserialize;
23
24/// A failure while requesting an embedding from the endpoint.
25#[derive(Debug, thiserror::Error)]
26pub enum EmbeddingError {
27    /// The HTTP request failed or the response body could not be parsed.
28    #[error("embedding request failed: {0}")]
29    Http(#[from] reqwest::Error),
30    /// The endpoint returned a non-success status code.
31    #[error("embedding endpoint returned status {0}")]
32    Status(reqwest::StatusCode),
33    /// The endpoint returned a success response carrying no embedding data.
34    #[error("embeddings response contained no data")]
35    EmptyResponse,
36}
37
38/// A packed embedding blob had a length that is not a multiple of 4.
39#[derive(Debug, thiserror::Error)]
40#[error("embedding blob length {0} is not a multiple of 4")]
41pub struct DecodeEmbeddingError(pub usize);
42
43/// Connection details for an OpenAI-compatible embedding endpoint.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct EmbeddingConfig {
46    /// Base URL up to and including `/v1` (e.g. `http://localhost:1234/v1`).
47    pub base_url: String,
48    /// Model identifier passed in the `model` JSON field.
49    pub model: String,
50    /// Bearer token. LM Studio / Ollama ignore it; `OpenAI` requires it.
51    pub api_key: Option<String>,
52}
53
54const DEFAULT_MODEL: &str = "text-embedding-bge-large-en-v1.5";
55const ENV_BASE_URL: &str = "KB_EMBEDDING_BASE_URL";
56const ENV_MODEL: &str = "KB_EMBEDDING_MODEL";
57const ENV_API_KEY: &str = "KB_EMBEDDING_API_KEY";
58
59/// Build an [`EmbeddingConfig`] from process environment variables.
60///
61/// Returns `None` when `KB_EMBEDDING_BASE_URL` is unset. `KB_EMBEDDING_MODEL`
62/// defaults to `text-embedding-bge-large-en-v1.5`. `KB_EMBEDDING_API_KEY`
63/// is optional.
64#[must_use]
65#[allow(
66    clippy::disallowed_methods,
67    reason = "kb-server takes all runtime config from the environment (12-factor); the embedding endpoint/model are deployment config, not user-facing tunables, and are read once at the server entry edge (REPO_INVARIANTS.md #5)"
68)]
69pub fn read_embedding_config_from_env() -> Option<EmbeddingConfig> {
70    config_from_env_inputs(
71        env::var(ENV_BASE_URL).ok(),
72        env::var(ENV_MODEL).ok(),
73        env::var(ENV_API_KEY).ok(),
74    )
75}
76
77/// Pure variant of [`read_embedding_config_from_env`] — caller supplies
78/// the env var values. Used by tests so they don't need to mutate the
79/// process environment.
80fn config_from_env_inputs(
81    base_url: Option<String>,
82    model: Option<String>,
83    api_key: Option<String>,
84) -> Option<EmbeddingConfig> {
85    let base_url = base_url?;
86    if base_url.is_empty() {
87        return None;
88    }
89    Some(EmbeddingConfig {
90        base_url,
91        model: model
92            .filter(|s| !s.is_empty())
93            .unwrap_or_else(|| DEFAULT_MODEL.into()),
94        api_key: api_key.filter(|s| !s.is_empty()),
95    })
96}
97
98/// Asynchronous embedding generator. Mirrors the Haskell
99/// `EmbeddingClient = Text -> IO (Either Text (Vector Float))`.
100#[async_trait]
101pub trait EmbeddingClient: Send + Sync {
102    /// Compute an embedding vector for `input`.
103    ///
104    /// # Errors
105    ///
106    /// Returns [`EmbeddingError`] on network, status, or empty-response
107    /// failure. Must never panic.
108    async fn embed(&self, input: &str) -> Result<Vec<f32>, EmbeddingError>;
109}
110
111/// HTTP-backed [`EmbeddingClient`] built on `reqwest::Client`.
112pub struct HttpEmbeddingClient {
113    base_url: String,
114    model: String,
115    api_key: Option<String>,
116    client: reqwest::Client,
117}
118
119impl HttpEmbeddingClient {
120    fn endpoint(&self) -> String {
121        format!("{}/embeddings", self.base_url)
122    }
123}
124
125#[async_trait]
126impl EmbeddingClient for HttpEmbeddingClient {
127    async fn embed(&self, input: &str) -> Result<Vec<f32>, EmbeddingError> {
128        let body = serde_json::json!({
129            "model": self.model,
130            "input": input,
131        });
132        let mut req = self.client.post(self.endpoint()).json(&body);
133        if let Some(key) = &self.api_key {
134            req = req.bearer_auth(key);
135        }
136        let resp = req.send().await?;
137        let status = resp.status();
138        if !status.is_success() {
139            return Err(EmbeddingError::Status(status));
140        }
141        let parsed: EmbeddingResponse = resp.json().await?;
142        parsed
143            .data
144            .into_iter()
145            .next()
146            .map(|item| item.embedding)
147            .ok_or(EmbeddingError::EmptyResponse)
148    }
149}
150
151/// Construct an HTTP-backed [`EmbeddingClient`].
152///
153/// Builds a `reqwest::Client` with a 30s timeout. The returned client is
154/// `Send + Sync` and can be wrapped in `Arc<dyn EmbeddingClient>` and
155/// shared across worker tasks.
156///
157/// # Panics
158///
159/// Panics only if the underlying `reqwest::Client::builder()` fails —
160/// that path requires a TLS / system configuration error and is not
161/// expected at runtime.
162#[must_use]
163pub fn http_embedding_client(config: EmbeddingConfig) -> HttpEmbeddingClient {
164    let client = reqwest::Client::builder()
165        .timeout(Duration::from_secs(30))
166        .build()
167        // `build()` only fails if the system TLS backend cannot be
168        // initialized; `Client::new()` is the documented default and
169        // reproduces the same failure mode without an explicit `expect`.
170        .unwrap_or_else(|_| reqwest::Client::new());
171    HttpEmbeddingClient {
172        base_url: config.base_url,
173        model: config.model,
174        api_key: config.api_key,
175        client,
176    }
177}
178
179#[derive(Deserialize)]
180struct EmbeddingResponse {
181    data: Vec<EmbeddingItem>,
182}
183
184#[derive(Deserialize)]
185struct EmbeddingItem {
186    embedding: Vec<f32>,
187}
188
189/// Encode a slice of `f32` as packed little-endian bytes (sqlite-vec
190/// compatible).
191#[must_use]
192pub fn encode_embedding(v: &[f32]) -> Vec<u8> {
193    let mut out = Vec::with_capacity(v.len() * 4);
194    for f in v {
195        out.extend_from_slice(&f.to_le_bytes());
196    }
197    out
198}
199
200/// Decode packed little-endian `f32` bytes back into a vector.
201///
202/// # Errors
203///
204/// Returns an error if `bytes.len()` is not a multiple of 4.
205pub fn decode_embedding(bytes: &[u8]) -> Result<Vec<f32>, DecodeEmbeddingError> {
206    if !bytes.len().is_multiple_of(4) {
207        return Err(DecodeEmbeddingError(bytes.len()));
208    }
209    let n = bytes.len() / 4;
210    let mut out = Vec::with_capacity(n);
211    for chunk in bytes.chunks_exact(4) {
212        // `chunks_exact(4)` yields slices of exactly 4 bytes, so the
213        // conversion always succeeds; the `Ok` guard avoids indexing or
214        // an `expect` while preserving behavior.
215        if let Ok(arr) = <[u8; 4]>::try_from(chunk) {
216            out.push(f32::from_le_bytes(arr));
217        }
218    }
219    Ok(out)
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use mockito::Matcher;
226    use proptest::prelude::*;
227
228    // ── embedding_config ─────────────────────────────────────────────
229
230    #[test]
231    fn embedding_config_returns_none_when_url_unset() {
232        assert_eq!(config_from_env_inputs(None, None, None), None);
233    }
234
235    #[test]
236    fn embedding_config_returns_none_when_url_empty() {
237        assert_eq!(
238            config_from_env_inputs(Some(String::new()), None, None),
239            None
240        );
241    }
242
243    #[test]
244    fn embedding_config_uses_default_model_when_unset() {
245        let cfg = config_from_env_inputs(Some("http://x/v1".into()), None, None).unwrap();
246        assert_eq!(cfg.base_url, "http://x/v1");
247        assert_eq!(cfg.model, "text-embedding-bge-large-en-v1.5");
248        assert!(cfg.api_key.is_none());
249    }
250
251    #[test]
252    fn embedding_config_passes_through_explicit_model() {
253        let cfg = config_from_env_inputs(
254            Some("http://x/v1".into()),
255            Some("custom-model".into()),
256            None,
257        )
258        .unwrap();
259        assert_eq!(cfg.model, "custom-model");
260    }
261
262    #[test]
263    fn embedding_config_carries_api_key_when_set() {
264        let cfg = config_from_env_inputs(Some("http://x/v1".into()), None, Some("sk-abc".into()))
265            .unwrap();
266        assert_eq!(cfg.api_key.as_deref(), Some("sk-abc"));
267    }
268
269    #[test]
270    fn embedding_config_treats_empty_api_key_as_none() {
271        let cfg =
272            config_from_env_inputs(Some("http://x/v1".into()), None, Some(String::new())).unwrap();
273        assert!(cfg.api_key.is_none());
274    }
275
276    #[test]
277    fn embedding_config_falls_back_when_model_empty() {
278        let cfg =
279            config_from_env_inputs(Some("http://x/v1".into()), Some(String::new()), None).unwrap();
280        assert_eq!(cfg.model, "text-embedding-bge-large-en-v1.5");
281    }
282
283    // ── embedding_blob_codec ─────────────────────────────────────────
284
285    #[test]
286    fn embedding_blob_codec_empty_vec_roundtrips() {
287        let v: Vec<f32> = vec![];
288        assert_eq!(decode_embedding(&encode_embedding(&v)).unwrap(), v);
289    }
290
291    #[test]
292    fn embedding_blob_codec_basic_roundtrip() {
293        let v = vec![0.0, 1.0, -1.0, 0.5, f32::MIN, f32::MAX];
294        let bytes = encode_embedding(&v);
295        assert_eq!(bytes.len(), 4 * v.len());
296        assert_eq!(decode_embedding(&bytes).unwrap(), v);
297    }
298
299    #[test]
300    fn embedding_blob_codec_rejects_truncated_buffer() {
301        let bad: Vec<u8> = vec![0; 7];
302        let err = decode_embedding(&bad).unwrap_err();
303        assert!(matches!(err, DecodeEmbeddingError(7)), "got: {err:?}");
304    }
305
306    #[test]
307    fn embedding_blob_codec_is_little_endian() {
308        // 1.0_f32 in IEEE-754 LE bytes is 00 00 80 3F.
309        let bytes = encode_embedding(&[1.0_f32]);
310        assert_eq!(bytes, vec![0x00, 0x00, 0x80, 0x3F]);
311    }
312
313    proptest! {
314        #[test]
315        fn embedding_blob_codec_property_roundtrip(
316            v in proptest::collection::vec(any::<f32>(), 0..32)
317        ) {
318            let bytes = encode_embedding(&v);
319            let back = decode_embedding(&bytes).unwrap();
320            // f32 NaN is not == itself, so compare bit patterns.
321            prop_assert_eq!(back.len(), v.len());
322            for (a, b) in v.iter().zip(back.iter()) {
323                prop_assert_eq!(a.to_bits(), b.to_bits());
324            }
325        }
326    }
327
328    // ── embedding_client ─────────────────────────────────────────────
329
330    fn make_client(server_url: &str, api_key: Option<&str>) -> HttpEmbeddingClient {
331        http_embedding_client(EmbeddingConfig {
332            base_url: server_url.to_string(),
333            model: "test-model".into(),
334            api_key: api_key.map(str::to_string),
335        })
336    }
337
338    #[tokio::test]
339    async fn embedding_client_posts_model_and_input() {
340        let mut server = mockito::Server::new_async().await;
341        let _m = server
342            .mock("POST", "/embeddings")
343            .match_header("content-type", "application/json")
344            .match_body(Matcher::PartialJsonString(
345                r#"{"model":"test-model","input":"hello"}"#.into(),
346            ))
347            .with_status(200)
348            .with_body(r#"{"data":[{"embedding":[0.1,0.2,0.3]}]}"#)
349            .create_async()
350            .await;
351        let client = make_client(&server.url(), None);
352        let v = client.embed("hello").await.unwrap();
353        assert_eq!(v, vec![0.1_f32, 0.2, 0.3]);
354    }
355
356    #[tokio::test]
357    async fn embedding_client_sets_bearer_auth_when_api_key_present() {
358        let mut server = mockito::Server::new_async().await;
359        let _m = server
360            .mock("POST", "/embeddings")
361            .match_header("authorization", "Bearer sk-abc")
362            .with_status(200)
363            .with_body(r#"{"data":[{"embedding":[1.0]}]}"#)
364            .create_async()
365            .await;
366        let client = make_client(&server.url(), Some("sk-abc"));
367        let v = client.embed("hi").await.unwrap();
368        assert_eq!(v, vec![1.0_f32]);
369    }
370
371    #[tokio::test]
372    async fn embedding_client_omits_auth_when_api_key_absent() {
373        let mut server = mockito::Server::new_async().await;
374        let _m = server
375            .mock("POST", "/embeddings")
376            .match_header("authorization", Matcher::Missing)
377            .with_status(200)
378            .with_body(r#"{"data":[{"embedding":[2.0]}]}"#)
379            .create_async()
380            .await;
381        let client = make_client(&server.url(), None);
382        let v = client.embed("hi").await.unwrap();
383        assert_eq!(v, vec![2.0_f32]);
384    }
385
386    #[tokio::test]
387    async fn embedding_client_returns_err_on_empty_data() {
388        let mut server = mockito::Server::new_async().await;
389        let _m = server
390            .mock("POST", "/embeddings")
391            .with_status(200)
392            .with_body(r#"{"data":[]}"#)
393            .create_async()
394            .await;
395        let client = make_client(&server.url(), None);
396        let err = client.embed("hi").await.unwrap_err();
397        assert!(matches!(err, EmbeddingError::EmptyResponse), "got: {err:?}");
398    }
399
400    #[tokio::test]
401    async fn embedding_client_returns_err_on_http_500() {
402        let mut server = mockito::Server::new_async().await;
403        let _m = server
404            .mock("POST", "/embeddings")
405            .with_status(500)
406            .with_body("boom")
407            .create_async()
408            .await;
409        let client = make_client(&server.url(), None);
410        let err = client.embed("hi").await.unwrap_err();
411        assert!(matches!(err, EmbeddingError::Status(_)), "got: {err:?}");
412    }
413
414    #[tokio::test]
415    async fn embedding_client_returns_err_on_malformed_json() {
416        let mut server = mockito::Server::new_async().await;
417        let _m = server
418            .mock("POST", "/embeddings")
419            .with_status(200)
420            .with_body("not json")
421            .create_async()
422            .await;
423        let client = make_client(&server.url(), None);
424        let err = client.embed("hi").await.unwrap_err();
425        assert!(matches!(err, EmbeddingError::Http(_)), "got: {err:?}");
426    }
427
428    #[tokio::test]
429    async fn embedding_client_returns_err_on_unreachable_endpoint() {
430        // A random unused port; reqwest fails fast with a connection refused.
431        let client = make_client("http://127.0.0.1:1", None);
432        let err = client.embed("hi").await.unwrap_err();
433        assert!(matches!(err, EmbeddingError::Http(_)), "got: {err:?}");
434    }
435
436    // ── http_embedding_client_async (criterion: http-embedding-client-uses-async-reqwest) ─
437
438    /// Asserts the `HttpEmbeddingClient` is async-only — building inside an
439    /// outer #[`tokio::main`] runtime never panics, and the `.embed()` call
440    /// awaits async send/json APIs. Calling this test inside a tokio
441    /// runtime (which v5's `reqwest::blocking` client could not) demonstrates
442    /// the original bug is fixed at the unit level.
443    #[tokio::test]
444    async fn http_embedding_client_async_construction_does_not_panic_in_tokio() {
445        let _client = http_embedding_client(EmbeddingConfig {
446            base_url: "http://127.0.0.1:1/v1".into(),
447            model: "test-model".into(),
448            api_key: None,
449        });
450    }
451
452    // ── embedding_client_async (criterion: embedding-trait-async) ─────
453
454    /// Asserts the `EmbeddingClient` trait is async — a stub impl satisfying
455    /// the trait must use an async fn (or equivalent), and the call site
456    /// must await it. Compiles only against the new async signature.
457    struct StubAsyncClient;
458
459    #[async_trait]
460    impl EmbeddingClient for StubAsyncClient {
461        async fn embed(&self, _input: &str) -> Result<Vec<f32>, EmbeddingError> {
462            Ok(vec![1.0, 2.0, 3.0])
463        }
464    }
465
466    #[tokio::test]
467    async fn embedding_client_async_trait_method_is_awaitable() {
468        let client = StubAsyncClient;
469        let v = client.embed("anything").await.unwrap();
470        assert_eq!(v, vec![1.0_f32, 2.0, 3.0]);
471    }
472}
kb/embedding.rs

kb/
embedding.rs
