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}
22
23/// Static specification for an OpenAI-compatible-protocol provider.
24///
25/// Implementing this trait grants the full embedding capability provided by
26/// `OpenAICompatEmbeddings`; it is the extension point for new OpenAI-compatible providers.
27pub trait CompatSpec: CompatConfigAccess + Sized + Default {
28    /// Environment variable name: API key (used in construction-time error messages, P1-3).
29    fn api_key_env() -> &'static str;
30    /// The batch limit for a single request.
31    fn batch_size() -> usize;
32    /// Vector dimension for a given model; unknown models must error (P1-2), never lying with a default.
33    fn dimension_for(model: &str) -> Result<usize, EmbeddingError>;
34    /// Constructs config from environment variables (reuses each config's from_env_result).
35    fn from_env_result() -> Result<Self, EmbeddingError>;
36}
37
38/// Generic OpenAI-compatible embedding client (P1-5).
39///
40/// Providers speaking the OpenAI `/embeddings` protocol (DeepSeek/Qwen, etc.) reuse this
41/// implementation via the [`CompatSpec`] config spec; `C` is each provider's config type.
42pub struct OpenAICompatEmbeddings<C: CompatConfigAccess + CompatSpec> {
43    config: C,
44    client: reqwest::Client,
45    dimension: usize,
46}
47
48impl<C: CompatConfigAccess + CompatSpec> std::fmt::Debug for OpenAICompatEmbeddings<C> {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("OpenAICompatEmbeddings")
51            .field("model", &self.config.model())
52            .field("dimension", &self.dimension)
53            .finish()
54    }
55}
56
57impl<C: CompatConfigAccess + CompatSpec> OpenAICompatEmbeddings<C> {
58    /// Fails fast at construction (P1-3): an empty API key errors immediately instead of
59    /// waiting until the request to 401; also validates the model dimension is known (P1-2).
60    pub fn new(config: C) -> Result<Self, EmbeddingError> {
61        if config.api_key().trim().is_empty() {
62            return Err(EmbeddingError::Config(format!(
63                "{} is empty",
64                C::api_key_env()
65            )));
66        }
67        let dimension = C::dimension_for(config.model())?;
68        Ok(Self {
69            config,
70            client: reqwest::Client::new(),
71            dimension,
72        })
73    }
74
75    /// Creates from environment variables, returning a Result.
76    pub fn from_env_result() -> Result<Self, EmbeddingError> {
77        let config = C::from_env_result()?;
78        Self::new(config)
79    }
80}
81
82#[async_trait]
83impl<C: CompatConfigAccess + CompatSpec + Send + Sync> Embeddings for OpenAICompatEmbeddings<C> {
84    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
85        if text.trim().is_empty() {
86            return Err(EmbeddingError::EmptyInput);
87        }
88
89        let url = format!("{}/embeddings", self.config.base_url());
90
91        let body = serde_json::json!({
92            "model": self.config.model(),
93            "input": text,
94        });
95
96        // P2-5: exponential backoff retry on 429/5xx.
97        let response = crate::retry::post_json_with_retry(
98            &self.client,
99            &url,
100            self.config.api_key(),
101            &body,
102            &crate::retry::DEFAULT_RETRY,
103        )
104        .await
105        .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
106
107        let status = response.status();
108        if !status.is_success() {
109            // P1-4: the error body must also error if reading fails; do not swallow it with unwrap_or_default().
110            let error_text = response.text().await.map_err(|e| {
111                EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
112            })?;
113            return Err(EmbeddingError::ApiError(format!(
114                "HTTP {}: {}",
115                status, error_text
116            )));
117        }
118
119        let embedding_response: EmbeddingResponse = response
120            .json()
121            .await
122            .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
123
124        let mut embedding = embedding_response
125            .data
126            .first()
127            .ok_or_else(|| EmbeddingError::ApiError("No embedding data in response".to_string()))?
128            .embedding
129            .clone();
130        // P2-8: uniform L2 normalization, guaranteeing unit length.
131        crate::l2_normalize(&mut embedding);
132        Ok(embedding)
133    }
134
135    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
136        // P1-1: an empty slice is not an error (nothing to do); only empty/all-whitespace texts error.
137        if texts.is_empty() {
138            return Ok(Vec::new());
139        }
140        if texts.iter().any(|t| t.trim().is_empty()) {
141            return Err(EmbeddingError::EmptyInput);
142        }
143
144        let url = format!("{}/embeddings", self.config.base_url());
145        let batch_size = C::batch_size().max(1);
146        // P0-1: collect item-by-item into Option slots, rejecting silent empty vectors. A chunk
147        // returning fewer/misaligned entries leaves None slots that error at the end, instead of
148        // producing zero vectors downstream treats as "dissimilar".
149        let mut all_results: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
150        let mut offset = 0;
151
152        for chunk in texts.chunks(batch_size) {
153            let body = serde_json::json!({
154                "model": self.config.model(),
155                "input": chunk,
156            });
157
158            // P2-5: exponential backoff retry on 429/5xx.
159            let response = crate::retry::post_json_with_retry(
160                &self.client,
161                &url,
162                self.config.api_key(),
163                &body,
164                &crate::retry::DEFAULT_RETRY,
165            )
166            .await
167            .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
168
169            let status = response.status();
170            if !status.is_success() {
171                // P1-4: the error body must also error if reading fails; do not swallow it with unwrap_or_default().
172                let error_text = response.text().await.map_err(|e| {
173                    EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
174                })?;
175                return Err(EmbeddingError::ApiError(format!(
176                    "HTTP {}: {}",
177                    status, error_text
178                )));
179            }
180
181            let embedding_response: EmbeddingResponse = response
182                .json()
183                .await
184                .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
185
186            for item in embedding_response.data {
187                let global_index = offset + item.index as usize;
188                if global_index >= all_results.len() {
189                    // Provider index beyond the requested range = batch misalignment; error out.
190                    return Err(EmbeddingError::BatchMismatch {
191                        expected: all_results.len(),
192                        actual: global_index + 1,
193                    });
194                }
195                all_results[global_index] = Some(item.embedding);
196            }
197            offset += chunk.len();
198        }
199
200        // Unwrap into Result: any empty slot errors explicitly rather than leaving a zero vector; then apply uniform L2 normalization (P2-8).
201        all_results
202            .into_iter()
203            .map(|opt| {
204                let mut v = opt.ok_or(EmbeddingError::EmptyVectorInBatch)?;
205                crate::l2_normalize(&mut v);
206                Ok(v)
207            })
208            .collect()
209    }
210
211    fn dimension(&self) -> usize {
212        self.dimension
213    }
214
215    fn model_name(&self) -> &str {
216        self.config.model()
217    }
218}
219
220/// Embedding response body for the OpenAI-compatible protocol (shared by DeepSeek/Qwen).
221#[derive(Debug, Deserialize)]
222struct EmbeddingResponse {
223    data: Vec<EmbeddingData>,
224}
225
226#[derive(Debug, Deserialize)]
227struct EmbeddingData {
228    embedding: Vec<f32>,
229    index: i32,
230}
lc_embeddings/openai_compat.rs

lc_embeddings/
openai_compat.rs
