agent_framework_openai/
embeddings.rs

1//! OpenAI embeddings client.
2//!
3//! Rust equivalent of upstream's `OpenAIEmbeddingClient`
4//! (`agent_framework_openai/_embedding_client.py`): the
5//! [`POST /v1/embeddings`](https://platform.openai.com/docs/api-reference/embeddings)
6//! endpoint, batching all input values into one request. Works against any
7//! OpenAI-compatible server via [`OpenAIEmbeddingClient::with_base_url`].
8
9use std::sync::Arc;
10
11use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
12use base64::Engine as _;
13
14use agent_framework_core::client::EmbeddingClient;
15use agent_framework_core::error::{Error, Result};
16use agent_framework_core::types::{
17    Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails,
18};
19use serde_json::{json, Map, Value};
20
21use crate::{classify_service_error, parse_retry_after, DEFAULT_BASE_URL};
22
23/// An OpenAI (or OpenAI-compatible) embeddings client.
24///
25/// ```no_run
26/// # use agent_framework_openai::OpenAIEmbeddingClient;
27/// # use agent_framework_core::client::EmbeddingClient;
28/// # async fn demo() -> agent_framework_core::error::Result<()> {
29/// let client = OpenAIEmbeddingClient::from_env("text-embedding-3-small")?;
30/// let batch = client
31///     .get_embeddings(vec!["Hello, world!".into()], None)
32///     .await?;
33/// println!("{} dims", batch[0].dimensions());
34/// # Ok(())
35/// # }
36/// ```
37#[derive(Clone)]
38pub struct OpenAIEmbeddingClient {
39    inner: Arc<Inner>,
40}
41
42struct Inner {
43    http: reqwest::Client,
44    api_key: String,
45    base_url: String,
46    model: String,
47    organization: Option<String>,
48}
49
50impl OpenAIEmbeddingClient {
51    /// Create a client for the given API key and default embedding model.
52    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
53        Self {
54            inner: Arc::new(Inner {
55                http: reqwest::Client::new(),
56                api_key: api_key.into(),
57                base_url: DEFAULT_BASE_URL.to_string(),
58                model: model.into(),
59                organization: None,
60            }),
61        }
62    }
63
64    /// Build a client from the environment: `OPENAI_API_KEY` (required) and
65    /// optional `OPENAI_BASE_URL`. The model falls back through
66    /// `OPENAI_EMBEDDING_MODEL` when the argument is empty — mirroring
67    /// upstream's `model or OPENAI_EMBEDDING_MODEL` resolution.
68    pub fn from_env(model: impl Into<String>) -> Result<Self> {
69        let key = std::env::var("OPENAI_API_KEY")
70            .map_err(|_| Error::Configuration("OPENAI_API_KEY is not set".into()))?;
71        let mut model = model.into();
72        if model.is_empty() {
73            model = std::env::var("OPENAI_EMBEDDING_MODEL").map_err(|_| {
74                Error::Configuration(
75                    "no embedding model: pass one or set OPENAI_EMBEDDING_MODEL".into(),
76                )
77            })?;
78        }
79        let mut client = Self::new(key, model);
80        if let Ok(base) = std::env::var("OPENAI_BASE_URL") {
81            client = client.with_base_url(base);
82        }
83        Ok(client)
84    }
85
86    /// Override the base URL (for OpenAI-compatible servers).
87    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
88        arc_inner(&mut self.inner).base_url = base_url.into();
89        self
90    }
91
92    /// Set the organization header.
93    pub fn with_organization(mut self, org: impl Into<String>) -> Self {
94        arc_inner(&mut self.inner).organization = Some(org.into());
95        self
96    }
97
98    fn build_body(&self, values: &[String], options: Option<&EmbeddingGenerationOptions>) -> Value {
99        let mut body = Map::new();
100        let model = options
101            .and_then(|o| o.model.clone())
102            .unwrap_or_else(|| self.inner.model.clone());
103        body.insert("model".into(), json!(model));
104        body.insert("input".into(), json!(values));
105        if let Some(options) = options {
106            if let Some(dimensions) = options.dimensions {
107                body.insert("dimensions".into(), json!(dimensions));
108            }
109            // Provider-specific extras understood by this endpoint
110            // (`encoding_format`, `user`) are forwarded verbatim.
111            for key in ["encoding_format", "user"] {
112                if let Some(v) = options.additional_properties.get(key) {
113                    body.insert(key.into(), v.clone());
114                }
115            }
116        }
117        Value::Object(body)
118    }
119}
120
121/// `Arc::make_mut` over an `Inner` that is not `Clone`-derivable field-wise —
122/// manual clone keeps the `reqwest::Client` (itself an `Arc` internally)
123/// shared.
124fn arc_inner(inner: &mut Arc<Inner>) -> &mut Inner {
125    if Arc::strong_count(inner) != 1 {
126        *inner = Arc::new(Inner {
127            http: inner.http.clone(),
128            api_key: inner.api_key.clone(),
129            base_url: inner.base_url.clone(),
130            model: inner.model.clone(),
131            organization: inner.organization.clone(),
132        });
133    }
134    Arc::get_mut(inner).expect("just ensured unique")
135}
136
137/// Decode one `embedding` field into a vector.
138///
139/// The field is a numeric array under the default `encoding_format: "float"`,
140/// but a **base64 string** when the caller asks for `encoding_format:
141/// "base64"` — a packed little-endian `f32` array, per the OpenAI embeddings
142/// API and the Azure AI Inference surface that mirrors it. Every client here
143/// forwards `encoding_format` verbatim, so refusing the encoded form would
144/// reject a perfectly successful response for a documented option.
145fn parse_embedding_vector(field: Option<&Value>) -> Result<Vec<f32>> {
146    match field {
147        Some(Value::Array(values)) => Ok(values
148            .iter()
149            .map(|v| v.as_f64().unwrap_or_default() as f32)
150            .collect()),
151        Some(Value::String(encoded)) => {
152            let bytes = BASE64_STANDARD.decode(encoded.as_bytes()).map_err(|e| {
153                Error::service(format!("embeddings item has undecodable base64: {e}"))
154            })?;
155            if bytes.len() % 4 != 0 {
156                return Err(Error::service(format!(
157                    "base64 embedding decodes to {} bytes, not a whole number of f32s",
158                    bytes.len()
159                )));
160            }
161            // `split_first_chunk::<4>` hands back a `&[u8; 4]` directly, so
162            // `from_le_bytes` needs no length check or fallible conversion.
163            // (`chunks_exact(4)` reads the same but draws clippy's
164            // `chunks_exact_to_as_chunks`, whose suggested `as_chunks` only
165            // stabilised at this crate's exact MSRV — `split_first_chunk` has
166            // been stable since 1.77, well under it.)
167            let mut vector = Vec::with_capacity(bytes.len() / 4);
168            let mut rest = bytes.as_slice();
169            while let Some((chunk, tail)) = rest.split_first_chunk::<4>() {
170                vector.push(f32::from_le_bytes(*chunk));
171                rest = tail;
172            }
173            Ok(vector)
174        }
175        _ => Err(Error::service("embeddings item missing 'embedding' vector")),
176    }
177}
178
179/// Parse an OpenAI-shaped embeddings response
180/// (`{"data": [{"embedding": [...], "index": n}], "model": .., "usage": ..}`)
181/// into a [`GeneratedEmbeddings`], restoring input order via `index`.
182///
183/// `embedding` may be a numeric array (the default `encoding_format:
184/// "float"`) or a base64 string of packed little-endian `f32`s
185/// (`encoding_format: "base64"`); both decode to the same vector.
186pub fn parse_embeddings_response(value: &Value) -> Result<GeneratedEmbeddings> {
187    let model = value.get("model").and_then(Value::as_str);
188    let data = value
189        .get("data")
190        .and_then(Value::as_array)
191        .ok_or_else(|| Error::service("embeddings response missing 'data' array"))?;
192
193    let mut indexed: Vec<(usize, Embedding)> = Vec::with_capacity(data.len());
194    for (position, item) in data.iter().enumerate() {
195        let vector = parse_embedding_vector(item.get("embedding"))?;
196        let index = item
197            .get("index")
198            .and_then(Value::as_u64)
199            .map(|i| i as usize)
200            .unwrap_or(position);
201        indexed.push((
202            index,
203            Embedding {
204                vector,
205                model: model.map(String::from),
206            },
207        ));
208    }
209    indexed.sort_by_key(|(i, _)| *i);
210
211    let mut batch = GeneratedEmbeddings::new(indexed.into_iter().map(|(_, e)| e).collect());
212    if let Some(usage) = value.get("usage") {
213        let input = usage.get("prompt_tokens").and_then(Value::as_u64);
214        let total = usage.get("total_tokens").and_then(Value::as_u64);
215        if input.is_some() || total.is_some() {
216            batch.usage = Some(UsageDetails {
217                input_token_count: input,
218                total_token_count: total,
219                ..Default::default()
220            });
221        }
222    }
223    Ok(batch)
224}
225
226#[async_trait::async_trait]
227impl EmbeddingClient for OpenAIEmbeddingClient {
228    async fn get_embeddings(
229        &self,
230        values: Vec<String>,
231        options: Option<EmbeddingGenerationOptions>,
232    ) -> Result<GeneratedEmbeddings> {
233        let body = self.build_body(&values, options.as_ref());
234        let url = format!("{}/embeddings", self.inner.base_url.trim_end_matches('/'));
235        let mut req = self
236            .inner
237            .http
238            .post(&url)
239            .bearer_auth(&self.inner.api_key)
240            .json(&body);
241        if let Some(org) = &self.inner.organization {
242            req = req.header("OpenAI-Organization", org);
243        }
244        let resp = req
245            .send()
246            .await
247            .map_err(|e| Error::service(format!("request failed: {e}")))?;
248        if !resp.status().is_success() {
249            let status = resp.status();
250            let retry_after = parse_retry_after(resp.headers());
251            let text = resp.text().await.unwrap_or_default();
252            return Err(classify_service_error(
253                status.as_u16(),
254                &text,
255                format!("OpenAI API error {status}: {text}"),
256                retry_after,
257            ));
258        }
259        let value: Value = resp
260            .json()
261            .await
262            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
263        parse_embeddings_response(&value)
264    }
265
266    fn model(&self) -> Option<&str> {
267        Some(&self.inner.model)
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    #[test]
276    fn build_body_includes_model_input_and_dimensions() {
277        let client = OpenAIEmbeddingClient::new("sk-test", "text-embedding-3-small");
278        let options = EmbeddingGenerationOptions::new().with_dimensions(256);
279        let body = client.build_body(&["a".into(), "b".into()], Some(&options));
280        assert_eq!(body["model"], "text-embedding-3-small");
281        assert_eq!(body["input"], json!(["a", "b"]));
282        assert_eq!(body["dimensions"], 256);
283    }
284
285    #[test]
286    fn build_body_option_model_overrides_default() {
287        let client = OpenAIEmbeddingClient::new("sk-test", "text-embedding-3-small");
288        let options = EmbeddingGenerationOptions::new().with_model("text-embedding-3-large");
289        let body = client.build_body(&["a".into()], Some(&options));
290        assert_eq!(body["model"], "text-embedding-3-large");
291    }
292
293    #[test]
294    fn build_body_forwards_known_additional_properties_only() {
295        let client = OpenAIEmbeddingClient::new("sk-test", "m");
296        let mut options = EmbeddingGenerationOptions::new();
297        options
298            .additional_properties
299            .insert("encoding_format".into(), json!("float"));
300        options
301            .additional_properties
302            .insert("unrelated".into(), json!(true));
303        let body = client.build_body(&["a".into()], Some(&options));
304        assert_eq!(body["encoding_format"], "float");
305        assert!(body.get("unrelated").is_none());
306    }
307
308    #[test]
309    fn parse_response_restores_index_order_and_usage() {
310        let value = json!({
311            "model": "text-embedding-3-small",
312            "data": [
313                { "index": 1, "embedding": [0.3, 0.4] },
314                { "index": 0, "embedding": [0.1, 0.2] },
315            ],
316            "usage": { "prompt_tokens": 5, "total_tokens": 5 }
317        });
318        let batch = parse_embeddings_response(&value).unwrap();
319        assert_eq!(batch.len(), 2);
320        assert_eq!(batch[0].vector, vec![0.1, 0.2]);
321        assert_eq!(batch[1].vector, vec![0.3, 0.4]);
322        assert_eq!(batch[0].model.as_deref(), Some("text-embedding-3-small"));
323        let usage = batch.usage.as_ref().unwrap();
324        assert_eq!(usage.input_token_count, Some(5));
325        assert_eq!(usage.total_token_count, Some(5));
326    }
327
328    #[test]
329    fn parse_response_missing_data_errors() {
330        assert!(parse_embeddings_response(&json!({})).is_err());
331    }
332
333    /// Every client here forwards `encoding_format` verbatim, so a caller can
334    /// ask for `"base64"` and get embeddings back as packed little-endian f32
335    /// strings. Rejecting those would fail a successful response for a
336    /// documented option, so the parser accepts both encodings.
337    #[test]
338    fn parse_response_decodes_base64_embeddings() {
339        let vector: [f32; 3] = [1.0, -2.5, 0.125];
340        let mut bytes = Vec::new();
341        for f in vector {
342            bytes.extend_from_slice(&f.to_le_bytes());
343        }
344        let encoded = BASE64_STANDARD.encode(&bytes);
345
346        let value = json!({
347            "data": [{ "index": 0, "embedding": encoded }],
348            "model": "text-embedding-3-small",
349        });
350        let batch = parse_embeddings_response(&value).unwrap();
351        assert_eq!(batch.embeddings[0].vector, vector.to_vec());
352    }
353
354    #[test]
355    fn parse_response_mixes_base64_and_array_items_by_index() {
356        let encoded = BASE64_STANDARD.encode(2.0f32.to_le_bytes());
357        let value = json!({
358            "data": [
359                { "index": 1, "embedding": encoded },
360                { "index": 0, "embedding": [1.0] },
361            ],
362        });
363        let batch = parse_embeddings_response(&value).unwrap();
364        assert_eq!(batch.embeddings[0].vector, vec![1.0]);
365        assert_eq!(batch.embeddings[1].vector, vec![2.0]);
366    }
367
368    #[test]
369    fn parse_response_rejects_undecodable_or_misaligned_base64() {
370        // Not valid base64 at all.
371        let bad = json!({ "data": [{ "index": 0, "embedding": "!!!not base64!!!" }] });
372        assert!(parse_embeddings_response(&bad).is_err());
373
374        // Valid base64, but not a whole number of f32s — better to say so
375        // than to silently hand back a truncated vector.
376        let misaligned = BASE64_STANDARD.encode([1u8, 2, 3, 4, 5]);
377        let value = json!({ "data": [{ "index": 0, "embedding": misaligned }] });
378        let err = parse_embeddings_response(&value).unwrap_err();
379        assert!(err.to_string().contains("whole number of f32s"), "{err}");
380    }
381}
agent_framework_openai/embeddings.rs

agent_framework_openai/
embeddings.rs
