lc_embeddings/
openai_compat.rs

1// lc-embeddings/src/openai_compat.rs
2//! Shared base class for OpenAI-compatible-protocol embedding clients (P1-5).
3//!
4//! DeepSeek and Qwen speak the same OpenAI `/embeddings` protocol (same request body,
5//! same `data[index]` alignment, same Bearer auth), and their sources were almost
6//! line-for-line duplicates. This module extracts the common implementation; DeepSeek/Qwen
7//! only configure URL / model / dimension / batch size via [`CompatSpec`].
8
9use crate::{EmbeddingError, Embeddings};
10use async_trait::async_trait;
11use serde::Deserialize;
12
13/// Abstraction over provider config fields — DeepSeek/Qwen config structs share field names.
14pub trait CompatConfigAccess {
15    /// Returns the API key
16    fn api_key(&self) -> &str;
17    /// Returns the Base URL
18    fn base_url(&self) -> &str;
19    /// Returns the model name
20    fn model(&self) -> &str;
21    /// Optional matryoshka output-dimension override (0.21.0 S5.1): when set,
22    /// the request carries a `dimensions` field. Default `None` — the request
23    /// body stays identical to pre-0.21.0 for providers that do not opt in.
24    fn dimensions(&self) -> Option<usize> {
25        None
26    }
27}
28
29/// Static specification for an OpenAI-compatible-protocol provider.
30///
31/// Implementing this trait grants the full embedding capability provided by
32/// `OpenAICompatEmbeddings`; it is the extension point for new OpenAI-compatible providers.
33pub trait CompatSpec: CompatConfigAccess + Sized + Default {
34    /// Environment variable name: API key (used in construction-time error messages, P1-3).
35    fn api_key_env() -> &'static str;
36    /// The batch limit for a single request.
37    fn batch_size() -> usize;
38    /// Vector dimension for a given model; unknown models must error (P1-2), never lying with a default.
39    fn dimension_for(model: &str) -> Result<usize, EmbeddingError>;
40    /// Per-instance validation beyond the static dimension check (P1-2/P1-3).
41    /// Default: valid. Used e.g. by Qwen to reject `dimensions` on models that
42    /// do not support the parameter.
43    fn validate(_config: &Self) -> Result<(), EmbeddingError> {
44        Ok(())
45    }
46    /// Constructs config from environment variables (reuses each config's from_env_result).
47    fn from_env_result() -> Result<Self, EmbeddingError>;
48}
49
50/// Builds the `/embeddings` request body.
51///
52/// Pure helper (0.21.0 S5.1): `dimensions` is only present when set — the
53/// serialized body stays identical to pre-0.21.0 otherwise (snapshotted in
54/// tests). `input` is a string for the single-text path and an array for
55/// batches, matching the original inline bodies.
56pub(crate) fn build_request_body(
57    model: &str,
58    input: serde_json::Value,
59    dimensions: Option<usize>,
60) -> serde_json::Value {
61    let mut body = serde_json::json!({
62        "model": model,
63        "input": input,
64    });
65    if let Some(d) = dimensions {
66        body["dimensions"] = serde_json::json!(d);
67    }
68    body
69}
70
71/// Generic OpenAI-compatible embedding client (P1-5).
72///
73/// Providers speaking the OpenAI `/embeddings` protocol (DeepSeek/Qwen, etc.) reuse this
74/// implementation via the [`CompatSpec`] config spec; `C` is each provider's config type.
75pub struct OpenAICompatEmbeddings<C: CompatConfigAccess + CompatSpec> {
76    config: C,
77    client: reqwest::Client,
78    dimension: usize,
79}
80
81impl<C: CompatConfigAccess + CompatSpec> std::fmt::Debug for OpenAICompatEmbeddings<C> {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.debug_struct("OpenAICompatEmbeddings")
84            .field("model", &self.config.model())
85            .field("dimension", &self.dimension)
86            .finish()
87    }
88}
89
90impl<C: CompatConfigAccess + CompatSpec> OpenAICompatEmbeddings<C> {
91    /// Fails fast at construction (P1-3): an empty API key errors immediately instead of
92    /// waiting until the request to 401; also validates the model dimension is known (P1-2).
93    pub fn new(config: C) -> Result<Self, EmbeddingError> {
94        if config.api_key().trim().is_empty() {
95            return Err(EmbeddingError::Config(format!(
96                "{} is empty",
97                C::api_key_env()
98            )));
99        }
100        C::validate(&config)?;
101        // 0.21.0 S5.1: an explicit matryoshka `dimensions` override wins over
102        // the model default (validated by the spec's `validate` hook).
103        let dimension = match config.dimensions() {
104            Some(d) => d,
105            None => C::dimension_for(config.model())?,
106        };
107        Ok(Self {
108            config,
109            client: reqwest::Client::new(),
110            dimension,
111        })
112    }
113
114    /// Creates from environment variables, returning a Result.
115    pub fn from_env_result() -> Result<Self, EmbeddingError> {
116        let config = C::from_env_result()?;
117        Self::new(config)
118    }
119}
120
121#[async_trait]
122impl<C: CompatConfigAccess + CompatSpec + Send + Sync> Embeddings for OpenAICompatEmbeddings<C> {
123    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
124        if text.trim().is_empty() {
125            return Err(EmbeddingError::EmptyInput);
126        }
127
128        let url = format!("{}/embeddings", self.config.base_url());
129
130        let body = build_request_body(
131            self.config.model(),
132            serde_json::json!(text),
133            self.config.dimensions(),
134        );
135
136        // P2-5: exponential backoff retry on 429/5xx.
137        let response = crate::retry::post_json_with_retry(
138            &self.client,
139            &url,
140            self.config.api_key(),
141            &body,
142            &crate::retry::DEFAULT_RETRY,
143        )
144        .await
145        .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
146
147        let status = response.status();
148        if !status.is_success() {
149            // P1-4: the error body must also error if reading fails; do not swallow it with unwrap_or_default().
150            let error_text = response.text().await.map_err(|e| {
151                EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
152            })?;
153            return Err(EmbeddingError::ApiError(format!(
154                "HTTP {}: {}",
155                status, error_text
156            )));
157        }
158
159        let embedding_response: EmbeddingResponse = response
160            .json()
161            .await
162            .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
163
164        let mut embedding = embedding_response
165            .data
166            .first()
167            .ok_or_else(|| EmbeddingError::ApiError("No embedding data in response".to_string()))?
168            .embedding
169            .clone();
170        // P2-8: uniform L2 normalization, guaranteeing unit length.
171        crate::l2_normalize(&mut embedding);
172        Ok(embedding)
173    }
174
175    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
176        // P1-1: an empty slice is not an error (nothing to do); only empty/all-whitespace texts error.
177        if texts.is_empty() {
178            return Ok(Vec::new());
179        }
180        if texts.iter().any(|t| t.trim().is_empty()) {
181            return Err(EmbeddingError::EmptyInput);
182        }
183
184        let url = format!("{}/embeddings", self.config.base_url());
185        let batch_size = C::batch_size().max(1);
186        // P0-1: collect item-by-item into Option slots, rejecting silent empty vectors. A chunk
187        // returning fewer/misaligned entries leaves None slots that error at the end, instead of
188        // producing zero vectors downstream treats as "dissimilar".
189        let mut all_results: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
190        let mut offset = 0;
191
192        for chunk in texts.chunks(batch_size) {
193            let body = build_request_body(
194                self.config.model(),
195                serde_json::json!(chunk),
196                self.config.dimensions(),
197            );
198
199            // P2-5: exponential backoff retry on 429/5xx.
200            let response = crate::retry::post_json_with_retry(
201                &self.client,
202                &url,
203                self.config.api_key(),
204                &body,
205                &crate::retry::DEFAULT_RETRY,
206            )
207            .await
208            .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
209
210            let status = response.status();
211            if !status.is_success() {
212                // P1-4: the error body must also error if reading fails; do not swallow it with unwrap_or_default().
213                let error_text = response.text().await.map_err(|e| {
214                    EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
215                })?;
216                return Err(EmbeddingError::ApiError(format!(
217                    "HTTP {}: {}",
218                    status, error_text
219                )));
220            }
221
222            let embedding_response: EmbeddingResponse = response
223                .json()
224                .await
225                .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
226
227            for item in embedding_response.data {
228                let global_index = offset + item.index as usize;
229                if global_index >= all_results.len() {
230                    // Provider index beyond the requested range = batch misalignment; error out.
231                    return Err(EmbeddingError::BatchMismatch {
232                        expected: all_results.len(),
233                        actual: global_index + 1,
234                    });
235                }
236                all_results[global_index] = Some(item.embedding);
237            }
238            offset += chunk.len();
239        }
240
241        // Unwrap into Result: any empty slot errors explicitly rather than leaving a zero vector; then apply uniform L2 normalization (P2-8).
242        all_results
243            .into_iter()
244            .map(|opt| {
245                let mut v = opt.ok_or(EmbeddingError::EmptyVectorInBatch)?;
246                crate::l2_normalize(&mut v);
247                Ok(v)
248            })
249            .collect()
250    }
251
252    fn dimension(&self) -> usize {
253        self.dimension
254    }
255
256    fn model_name(&self) -> &str {
257        self.config.model()
258    }
259}
260
261/// Embedding response body for the OpenAI-compatible protocol (shared by DeepSeek/Qwen).
262#[derive(Debug, Deserialize)]
263struct EmbeddingResponse {
264    data: Vec<EmbeddingData>,
265}
266
267#[derive(Debug, Deserialize)]
268struct EmbeddingData {
269    embedding: Vec<f32>,
270    index: i32,
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    /// 0.21.0 S5.1: without `dimensions` the body is identical to pre-0.21.0.
278    #[test]
279    fn request_body_without_dimensions_unchanged() {
280        let body = build_request_body("text-embedding-v1", serde_json::json!(["a", "b"]), None);
281        assert_eq!(
282            body,
283            serde_json::json!({"model": "text-embedding-v1", "input": ["a", "b"]})
284        );
285        assert!(
286            body.get("dimensions").is_none(),
287            "no `dimensions` key may leak into requests that did not opt in"
288        );
289    }
290
291    /// 0.21.0 S5.1: with `dimensions` the matryoshka override is carried in the body.
292    #[test]
293    fn request_body_with_dimensions() {
294        let body = build_request_body("qwen3-embedding-0.6b", serde_json::json!("text"), Some(512));
295        assert_eq!(body["dimensions"], 512);
296        assert_eq!(body["model"], "qwen3-embedding-0.6b");
297        assert_eq!(body["input"], "text");
298    }
299}
lc_embeddings/openai_compat.rs

lc_embeddings/
openai_compat.rs
