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

kb/
embedding.rs
