rig_core/providers/openai/
embedding.rs

1use super::{client::ApiResponse, completion::Usage};
2use crate::embeddings::EmbeddingError;
3use crate::http_client::HttpClientExt;
4use crate::{embeddings, http_client};
5use serde::{Deserialize, Serialize};
6use serde_json::json;
7
8// ================================================================
9// OpenAI Embedding API
10// ================================================================
11/// `text-embedding-3-large` embedding model
12pub const TEXT_EMBEDDING_3_LARGE: &str = "text-embedding-3-large";
13/// `text-embedding-3-small` embedding model
14pub const TEXT_EMBEDDING_3_SMALL: &str = "text-embedding-3-small";
15/// `text-embedding-ada-002` embedding model
16pub const TEXT_EMBEDDING_ADA_002: &str = "text-embedding-ada-002";
17
18#[derive(Debug, Deserialize)]
19pub struct EmbeddingResponse {
20    pub object: String,
21    pub data: Vec<EmbeddingData>,
22    pub model: String,
23    pub usage: Usage,
24}
25
26#[derive(Debug, Deserialize, Clone, Serialize)]
27#[serde(rename_all = "snake_case")]
28pub enum EncodingFormat {
29    Float,
30    Base64,
31}
32
33#[derive(Debug, Deserialize)]
34pub struct EmbeddingData {
35    pub object: String,
36    pub embedding: Vec<serde_json::Number>,
37    pub index: usize,
38}
39
40#[doc(hidden)]
41#[derive(Clone)]
42pub struct GenericEmbeddingModel<Ext = super::OpenAIResponsesExt, H = reqwest::Client> {
43    client: crate::client::Client<Ext, H>,
44    pub model: String,
45    pub encoding_format: Option<EncodingFormat>,
46    pub user: Option<String>,
47    ndims: usize,
48}
49
50/// The embedding model struct for OpenAI's Embeddings API.
51///
52/// This preserves the historical public generic shape where the first generic
53/// parameter is the HTTP client type.
54pub type EmbeddingModel<H = reqwest::Client> = GenericEmbeddingModel<super::OpenAIResponsesExt, H>;
55
56fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
57    match identifier {
58        TEXT_EMBEDDING_3_LARGE => Some(3_072),
59        TEXT_EMBEDDING_3_SMALL | TEXT_EMBEDDING_ADA_002 => Some(1_536),
60        _ => None,
61    }
62}
63
64impl<Ext, H> embeddings::EmbeddingModel for GenericEmbeddingModel<Ext, H>
65where
66    crate::client::Client<Ext, H>: HttpClientExt + Clone + std::fmt::Debug + Send + 'static,
67    Ext: crate::client::Provider + Clone + 'static,
68    H: Clone + Default + std::fmt::Debug + 'static,
69{
70    const MAX_DOCUMENTS: usize = 1024;
71
72    type Client = crate::client::Client<Ext, H>;
73
74    fn make(client: &Self::Client, model: impl Into<String>, ndims: Option<usize>) -> Self {
75        let model = model.into();
76        let dims = ndims
77            .or(model_dimensions_from_identifier(&model))
78            .unwrap_or_default();
79
80        Self::new(client.clone(), model, dims)
81    }
82
83    fn ndims(&self) -> usize {
84        self.ndims
85    }
86
87    async fn embed_texts(
88        &self,
89        documents: impl IntoIterator<Item = String>,
90    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
91        let documents: Vec<String> = documents.into_iter().collect();
92        let response = self.embed_texts_with_usage(documents).await?;
93        Ok(response.embeddings)
94    }
95
96    async fn embed_texts_with_usage(
97        &self,
98        documents: impl IntoIterator<Item = String>,
99    ) -> Result<embeddings::EmbeddingResponse, EmbeddingError> {
100        let documents: Vec<String> = documents.into_iter().collect();
101
102        let mut body = json!({
103            "model": self.model,
104            "input": documents,
105        });
106
107        let body_object = body.as_object_mut().ok_or_else(|| {
108            EmbeddingError::ResponseError("embedding request body must be a JSON object".into())
109        })?;
110
111        if self.ndims > 0 && self.model.as_str() != TEXT_EMBEDDING_ADA_002 {
112            body_object.insert("dimensions".to_owned(), json!(self.ndims));
113        }
114
115        if let Some(encoding_format) = &self.encoding_format {
116            body_object.insert("encoding_format".to_owned(), json!(encoding_format));
117        }
118
119        if let Some(user) = &self.user {
120            body_object.insert("user".to_owned(), json!(user));
121        }
122
123        let body = serde_json::to_vec(&body)?;
124
125        let req = self
126            .client
127            .post("/embeddings")?
128            .body(body)
129            .map_err(|e| EmbeddingError::HttpError(e.into()))?;
130
131        let response = self.client.send(req).await?;
132
133        let status = response.status();
134        if status.is_success() {
135            let response_body: Vec<u8> = response.into_body().await?;
136            let parsed: ApiResponse<EmbeddingResponse> = serde_json::from_slice(&response_body)?;
137
138            match parsed {
139                ApiResponse::Ok(response) => {
140                    tracing::info!(target: "rig",
141                        "OpenAI embedding token usage: {:?}",
142                        response.usage
143                    );
144
145                    if response.data.len() != documents.len() {
146                        return Err(EmbeddingError::ResponseError(
147                            "Response data length does not match input length".into(),
148                        ));
149                    }
150
151                    let usage = crate::completion::Usage {
152                        input_tokens: response.usage.prompt_tokens as u64,
153                        output_tokens: 0,
154                        total_tokens: response.usage.total_tokens as u64,
155                        cached_input_tokens: response
156                            .usage
157                            .prompt_tokens_details
158                            .as_ref()
159                            .map_or(0, |d| d.cached_tokens as u64),
160                        cache_creation_input_tokens: 0,
161                        tool_use_prompt_tokens: 0,
162                        reasoning_tokens: 0,
163                    };
164
165                    let embeddings = response
166                        .data
167                        .into_iter()
168                        .zip(documents.into_iter())
169                        .map(|(embedding, document)| embeddings::Embedding {
170                            document,
171                            vec: embedding
172                                .embedding
173                                .into_iter()
174                                .filter_map(|n| n.as_f64())
175                                .collect(),
176                        })
177                        .collect();
178
179                    Ok(embeddings::EmbeddingResponse { embeddings, usage })
180                }
181                ApiResponse::Err(err) => {
182                    tracing::warn!(message = %err.message, "provider returned an error response");
183                    Err(EmbeddingError::from_http_response(
184                        status,
185                        String::from_utf8_lossy(&response_body).into_owned(),
186                    ))
187                }
188            }
189        } else {
190            let text = http_client::text(response).await?;
191            Err(EmbeddingError::from_http_response(status, text))
192        }
193    }
194}
195
196impl<Ext, H> GenericEmbeddingModel<Ext, H>
197where
198    Ext: crate::client::Provider,
199{
200    pub fn new(
201        client: crate::client::Client<Ext, H>,
202        model: impl Into<String>,
203        ndims: usize,
204    ) -> Self {
205        Self {
206            client,
207            model: model.into(),
208            encoding_format: None,
209            ndims,
210            user: None,
211        }
212    }
213
214    pub fn with_model(client: crate::client::Client<Ext, H>, model: &str, ndims: usize) -> Self {
215        Self {
216            client,
217            model: model.into(),
218            encoding_format: None,
219            ndims,
220            user: None,
221        }
222    }
223
224    pub fn with_encoding_format(
225        client: crate::client::Client<Ext, H>,
226        model: &str,
227        ndims: usize,
228        encoding_format: EncodingFormat,
229    ) -> Self {
230        Self {
231            client,
232            model: model.into(),
233            encoding_format: Some(encoding_format),
234            ndims,
235            user: None,
236        }
237    }
238
239    pub fn encoding_format(mut self, encoding_format: EncodingFormat) -> Self {
240        self.encoding_format = Some(encoding_format);
241        self
242    }
243
244    pub fn user(mut self, user: impl Into<String>) -> Self {
245        self.user = Some(user.into());
246        self
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use crate::client::EmbeddingsClient;
254    use crate::embeddings::EmbeddingModel as _;
255    use crate::providers::openai::CompletionsClient;
256    use crate::test_utils::RecordingHttpClient;
257
258    #[tokio::test]
259    async fn embedding_preserves_raw_provider_error_json_on_api_error_envelope() {
260        let body = r#"{"message":"embedding quota exceeded","type":"insufficient_quota"}"#;
261        let http_client =
262            RecordingHttpClient::with_error_response(http::StatusCode::ACCEPTED, body);
263        let client = CompletionsClient::builder()
264            .api_key("test-key")
265            .http_client(http_client)
266            .build()
267            .expect("build client");
268        let model = client.embedding_model("text-embedding-3-small");
269
270        let error = model
271            .embed_texts(["hello".to_string()])
272            .await
273            .expect_err("embedding should fail with provider error envelope");
274
275        match &error {
276            EmbeddingError::ProviderResponse(stored) => {
277                assert_eq!(stored.body, body);
278                assert_eq!(stored.status, Some(http::StatusCode::ACCEPTED));
279                assert_eq!(error.provider_response_body(), Some(body));
280                let json = error
281                    .provider_response_json()
282                    .expect("raw body should be valid JSON")
283                    .expect("parsed JSON should be present");
284                assert_eq!(json["type"], "insufficient_quota");
285            }
286            other => panic!("expected ProviderResponse, got {other:?}"),
287        }
288    }
289
290    #[tokio::test]
291    async fn embedding_http_non_success_preserves_status_and_body() {
292        let body = r#"{"error":{"message":"invalid api key","type":"invalid_request_error"}}"#;
293        let http_client =
294            RecordingHttpClient::with_error_response(http::StatusCode::UNAUTHORIZED, body);
295        let client = CompletionsClient::builder()
296            .api_key("test-key")
297            .http_client(http_client)
298            .build()
299            .expect("build client");
300        let model = client.embedding_model("text-embedding-3-small");
301
302        let error = model
303            .embed_texts(["hello".to_string()])
304            .await
305            .expect_err("embedding should fail with non-success status");
306
307        assert!(matches!(error, EmbeddingError::HttpError(_)));
308        assert_eq!(
309            error.provider_response_status(),
310            Some(http::StatusCode::UNAUTHORIZED)
311        );
312        assert_eq!(error.provider_response_body(), Some(body));
313    }
314}
rig_core/providers/openai/embedding.rs

rig_core/providers/openai/
embedding.rs
