llm_embedding/
openai_compatible.rs

1//! OpenAI-compatible HTTP embedding provider.
2//!
3//! One transport covers every OpenAI-compatible endpoint (OpenAI, Gemini via
4//! its compatibility layer, Azure OpenAI, Ollama, self-hosted servers):
5//! provider differences are endpoint URLs plus optional bearer keys, selected
6//! through [`EmbeddingConfig`] presets, not through transport branches.
7
8use async_trait::async_trait;
9use serde::{Deserialize, Serialize};
10
11use crate::config::EmbeddingConfig;
12use crate::error::{EmbeddingError, Result};
13use crate::preprocessor::PreprocessorImpl;
14use crate::provider::{EmbeddingProvider, EmbeddingResult};
15
16/// Provider that talks to any OpenAI-compatible `/embeddings` endpoint.
17pub struct OpenAICompatibleProvider {
18    config: EmbeddingConfig,
19    preprocessor: PreprocessorImpl,
20    client: reqwest::Client,
21}
22
23impl std::fmt::Debug for OpenAICompatibleProvider {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("OpenAICompatibleProvider")
26            .field("model", &self.config.model)
27            .field("base_url", &self.config.base_url)
28            .finish()
29    }
30}
31
32/// Request body for the `/embeddings` endpoint.
33#[derive(Debug, Serialize)]
34struct EmbeddingRequest {
35    model: String,
36    input: Vec<String>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    encoding_format: Option<String>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    dimensions: Option<usize>,
41}
42
43/// Response body from the `/embeddings` endpoint.
44#[derive(Debug, Deserialize)]
45struct EmbeddingResponse {
46    data: Vec<EmbeddingData>,
47    #[serde(default)]
48    usage: Option<EmbeddingUsage>,
49}
50
51/// One embedding entry in the response.
52#[derive(Debug, Deserialize)]
53struct EmbeddingData {
54    index: usize,
55    embedding: Vec<f32>,
56}
57
58/// Token usage reported by the provider.
59#[derive(Debug, Deserialize, Default)]
60struct EmbeddingUsage {
61    #[serde(default)]
62    prompt_tokens: u64,
63    #[serde(default)]
64    total_tokens: u64,
65}
66
67impl OpenAICompatibleProvider {
68    /// Creates a provider; fails fast when the config is incomplete.
69    pub fn new(config: EmbeddingConfig) -> Result<Self> {
70        config.validate()?;
71        let client = reqwest::Client::builder()
72            .timeout(std::time::Duration::from_secs(config.timeout_secs))
73            .build()
74            .map_err(|err| EmbeddingError::Config(format!("failed to build HTTP client: {err}")))?;
75        let preprocessor = PreprocessorImpl::from_config(&config.preprocessor);
76        Ok(Self {
77            config,
78            preprocessor,
79            client,
80        })
81    }
82
83    /// Returns the provider configuration.
84    pub fn config(&self) -> &EmbeddingConfig {
85        &self.config
86    }
87
88    /// Returns the active text preprocessor.
89    pub fn preprocessor(&self) -> &PreprocessorImpl {
90        &self.preprocessor
91    }
92
93    /// Builds the request body, applying preprocessing first.
94    fn build_request(&self, texts: &[String]) -> EmbeddingRequest {
95        let borrowed: Vec<&str> = texts.iter().map(String::as_str).collect();
96        EmbeddingRequest {
97            model: self.config.model.clone(),
98            input: self.preprocessor.process_batch(&borrowed),
99            encoding_format: Some("float".into()),
100            dimensions: self.config.dimension,
101        }
102    }
103
104    /// Reassembles the response into input order and validates the contract.
105    fn parse_response(&self, response: EmbeddingResponse) -> Result<EmbeddingResult> {
106        let mut ordered = response.data;
107        ordered.sort_by_key(|item| item.index);
108        validate_embedding_data(&ordered, &self.config)?;
109
110        let usage = response.usage.unwrap_or_default();
111        Ok(EmbeddingResult {
112            embeddings: ordered.into_iter().map(|item| item.embedding).collect(),
113            prompt_tokens: usage.prompt_tokens,
114            total_tokens: usage.total_tokens,
115        })
116    }
117}
118
119/// Validates the provider response contract: exact count, sequential indexes,
120/// non-empty finite vectors, and the configured dimension when set.
121fn validate_embedding_data(data: &[EmbeddingData], config: &EmbeddingConfig) -> Result<()> {
122    for (expected_index, item) in data.iter().enumerate() {
123        if item.index != expected_index {
124            return Err(EmbeddingError::Decode(format!(
125                "embedding index mismatch: expected {expected_index}, received {}",
126                item.index
127            )));
128        }
129        if item.embedding.is_empty() {
130            return Err(EmbeddingError::InvalidRequest(format!(
131                "embedding at index {expected_index} is empty; check the dimension configuration"
132            )));
133        }
134        if let Some(dimension) = config.dimension {
135            if item.embedding.len() != dimension {
136                return Err(EmbeddingError::InvalidRequest(format!(
137                    "embedding dimension mismatch at index {expected_index}: expected {dimension}, received {}",
138                    item.embedding.len()
139                )));
140            }
141        }
142        if item.embedding.iter().any(|value| !value.is_finite()) {
143            return Err(EmbeddingError::InvalidRequest(format!(
144                "embedding at index {expected_index} contains a non-finite value"
145            )));
146        }
147    }
148    Ok(())
149}
150
151#[async_trait]
152impl EmbeddingProvider for OpenAICompatibleProvider {
153    async fn embed(&self, texts: &[String]) -> Result<EmbeddingResult> {
154        if texts.is_empty() {
155            return Ok(EmbeddingResult::default());
156        }
157
158        let request = self.build_request(texts);
159        let mut outgoing = self.client.post(&self.config.base_url).json(&request);
160        if let Some(api_key) = &self.config.api_key {
161            outgoing = outgoing.bearer_auth(api_key);
162        }
163        let response = outgoing.send().await?;
164
165        let status = response.status();
166        if !status.is_success() {
167            let body = response.text().await.unwrap_or_default();
168            return Err(EmbeddingError::Provider {
169                status: status.as_u16(),
170                message: body,
171            });
172        }
173
174        let decoded: EmbeddingResponse = response
175            .json()
176            .await
177            .map_err(|err| EmbeddingError::Decode(err.to_string()))?;
178        let result = self.parse_response(decoded)?;
179        if result.embeddings.len() != texts.len() {
180            return Err(EmbeddingError::Decode(format!(
181                "embedding count mismatch: expected {}, received {}",
182                texts.len(),
183                result.embeddings.len()
184            )));
185        }
186        Ok(result)
187    }
188
189    fn dimension(&self) -> usize {
190        self.config.dimension.unwrap_or(0)
191    }
192
193    fn model_name(&self) -> &str {
194        &self.config.model
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use crate::preprocessor::PreprocessorConfig;
202
203    fn test_config() -> EmbeddingConfig {
204        EmbeddingConfig::new("http://example.com/embeddings", "test-model").with_dimension(2)
205    }
206
207    #[test]
208    fn requires_complete_config() {
209        let missing_dimension = EmbeddingConfig::new("http://example.com", "model");
210        assert!(OpenAICompatibleProvider::new(missing_dimension).is_err());
211    }
212
213    #[test]
214    fn build_request_applies_preprocessor_and_dimensions() {
215        let config = EmbeddingConfig::new("http://example.com", "nomic-embed")
216            .with_dimension(768)
217            .with_preprocessor(PreprocessorConfig::Prefix {
218                prefix: "search_query: ".into(),
219            });
220        let provider = OpenAICompatibleProvider::new(config).expect("valid config");
221        let request = provider.build_request(&["rust".to_string()]);
222        assert_eq!(request.model, "nomic-embed");
223        assert_eq!(request.input, vec!["search_query: rust".to_string()]);
224        assert_eq!(request.encoding_format.as_deref(), Some("float"));
225        assert_eq!(request.dimensions, Some(768));
226    }
227
228    #[test]
229    fn parse_response_reorders_by_index() {
230        let provider = OpenAICompatibleProvider::new(test_config()).expect("valid config");
231        let response = EmbeddingResponse {
232            data: vec![
233                EmbeddingData {
234                    index: 1,
235                    embedding: vec![3.0, 4.0],
236                },
237                EmbeddingData {
238                    index: 0,
239                    embedding: vec![1.0, 2.0],
240                },
241            ],
242            usage: None,
243        };
244        let result = provider.parse_response(response).expect("ordered parse");
245        assert_eq!(result.embeddings[0], vec![1.0, 2.0]);
246        assert_eq!(result.embeddings[1], vec![3.0, 4.0]);
247    }
248
249    #[test]
250    fn parse_response_rejects_contract_violations() {
251        let provider = OpenAICompatibleProvider::new(test_config()).expect("valid config");
252
253        let duplicate = EmbeddingResponse {
254            data: vec![
255                EmbeddingData {
256                    index: 0,
257                    embedding: vec![1.0, 2.0],
258                },
259                EmbeddingData {
260                    index: 0,
261                    embedding: vec![3.0, 4.0],
262                },
263            ],
264            usage: None,
265        };
266        assert!(provider.parse_response(duplicate).is_err());
267
268        let wrong_dimension = EmbeddingResponse {
269            data: vec![EmbeddingData {
270                index: 0,
271                embedding: vec![1.0],
272            }],
273            usage: None,
274        };
275        assert!(provider.parse_response(wrong_dimension).is_err());
276
277        let non_finite = EmbeddingResponse {
278            data: vec![EmbeddingData {
279                index: 0,
280                embedding: vec![f32::NAN, 2.0],
281            }],
282            usage: None,
283        };
284        assert!(provider.parse_response(non_finite).is_err());
285    }
286
287    #[test]
288    fn accessors_expose_model_and_dimension() {
289        let provider = OpenAICompatibleProvider::new(test_config()).expect("valid config");
290        assert_eq!(provider.model_name(), "test-model");
291        assert_eq!(provider.dimension(), 2);
292        assert_eq!(provider.preprocessor().preprocess("x"), "x".to_string());
293    }
294}
llm_embedding/openai_compatible.rs

llm_embedding/
openai_compatible.rs
