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 agent_framework_core::client::EmbeddingClient;
12use agent_framework_core::error::{Error, Result};
13use agent_framework_core::types::{
14    Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails,
15};
16use serde_json::{json, Map, Value};
17
18use crate::{classify_service_error, parse_retry_after, DEFAULT_BASE_URL};
19
20/// An OpenAI (or OpenAI-compatible) embeddings client.
21///
22/// ```no_run
23/// # use agent_framework_openai::OpenAIEmbeddingClient;
24/// # use agent_framework_core::client::EmbeddingClient;
25/// # async fn demo() -> agent_framework_core::error::Result<()> {
26/// let client = OpenAIEmbeddingClient::from_env("text-embedding-3-small")?;
27/// let batch = client
28///     .get_embeddings(vec!["Hello, world!".into()], None)
29///     .await?;
30/// println!("{} dims", batch[0].dimensions());
31/// # Ok(())
32/// # }
33/// ```
34#[derive(Clone)]
35pub struct OpenAIEmbeddingClient {
36    inner: Arc<Inner>,
37}
38
39struct Inner {
40    http: reqwest::Client,
41    api_key: String,
42    base_url: String,
43    model: String,
44    organization: Option<String>,
45}
46
47impl OpenAIEmbeddingClient {
48    /// Create a client for the given API key and default embedding model.
49    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
50        Self {
51            inner: Arc::new(Inner {
52                http: reqwest::Client::new(),
53                api_key: api_key.into(),
54                base_url: DEFAULT_BASE_URL.to_string(),
55                model: model.into(),
56                organization: None,
57            }),
58        }
59    }
60
61    /// Build a client from the environment: `OPENAI_API_KEY` (required) and
62    /// optional `OPENAI_BASE_URL`. The model falls back through
63    /// `OPENAI_EMBEDDING_MODEL` when the argument is empty — mirroring
64    /// upstream's `model or OPENAI_EMBEDDING_MODEL` resolution.
65    pub fn from_env(model: impl Into<String>) -> Result<Self> {
66        let key = std::env::var("OPENAI_API_KEY")
67            .map_err(|_| Error::Configuration("OPENAI_API_KEY is not set".into()))?;
68        let mut model = model.into();
69        if model.is_empty() {
70            model = std::env::var("OPENAI_EMBEDDING_MODEL").map_err(|_| {
71                Error::Configuration(
72                    "no embedding model: pass one or set OPENAI_EMBEDDING_MODEL".into(),
73                )
74            })?;
75        }
76        let mut client = Self::new(key, model);
77        if let Ok(base) = std::env::var("OPENAI_BASE_URL") {
78            client = client.with_base_url(base);
79        }
80        Ok(client)
81    }
82
83    /// Override the base URL (for OpenAI-compatible servers).
84    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
85        arc_inner(&mut self.inner).base_url = base_url.into();
86        self
87    }
88
89    /// Set the organization header.
90    pub fn with_organization(mut self, org: impl Into<String>) -> Self {
91        arc_inner(&mut self.inner).organization = Some(org.into());
92        self
93    }
94
95    fn build_body(&self, values: &[String], options: Option<&EmbeddingGenerationOptions>) -> Value {
96        let mut body = Map::new();
97        let model = options
98            .and_then(|o| o.model.clone())
99            .unwrap_or_else(|| self.inner.model.clone());
100        body.insert("model".into(), json!(model));
101        body.insert("input".into(), json!(values));
102        if let Some(options) = options {
103            if let Some(dimensions) = options.dimensions {
104                body.insert("dimensions".into(), json!(dimensions));
105            }
106            // Provider-specific extras understood by this endpoint
107            // (`encoding_format`, `user`) are forwarded verbatim.
108            for key in ["encoding_format", "user"] {
109                if let Some(v) = options.additional_properties.get(key) {
110                    body.insert(key.into(), v.clone());
111                }
112            }
113        }
114        Value::Object(body)
115    }
116}
117
118/// `Arc::make_mut` over an `Inner` that is not `Clone`-derivable field-wise —
119/// manual clone keeps the `reqwest::Client` (itself an `Arc` internally)
120/// shared.
121fn arc_inner(inner: &mut Arc<Inner>) -> &mut Inner {
122    if Arc::strong_count(inner) != 1 {
123        *inner = Arc::new(Inner {
124            http: inner.http.clone(),
125            api_key: inner.api_key.clone(),
126            base_url: inner.base_url.clone(),
127            model: inner.model.clone(),
128            organization: inner.organization.clone(),
129        });
130    }
131    Arc::get_mut(inner).expect("just ensured unique")
132}
133
134/// Parse an OpenAI-shaped embeddings response
135/// (`{"data": [{"embedding": [...], "index": n}], "model": .., "usage": ..}`)
136/// into a [`GeneratedEmbeddings`], restoring input order via `index`.
137pub fn parse_embeddings_response(value: &Value) -> Result<GeneratedEmbeddings> {
138    let model = value.get("model").and_then(Value::as_str);
139    let data = value
140        .get("data")
141        .and_then(Value::as_array)
142        .ok_or_else(|| Error::service("embeddings response missing 'data' array"))?;
143
144    let mut indexed: Vec<(usize, Embedding)> = Vec::with_capacity(data.len());
145    for (position, item) in data.iter().enumerate() {
146        let vector: Vec<f32> = item
147            .get("embedding")
148            .and_then(Value::as_array)
149            .ok_or_else(|| Error::service("embeddings item missing 'embedding' vector"))?
150            .iter()
151            .map(|v| v.as_f64().unwrap_or_default() as f32)
152            .collect();
153        let index = item
154            .get("index")
155            .and_then(Value::as_u64)
156            .map(|i| i as usize)
157            .unwrap_or(position);
158        indexed.push((
159            index,
160            Embedding {
161                vector,
162                model: model.map(String::from),
163            },
164        ));
165    }
166    indexed.sort_by_key(|(i, _)| *i);
167
168    let mut batch = GeneratedEmbeddings::new(indexed.into_iter().map(|(_, e)| e).collect());
169    if let Some(usage) = value.get("usage") {
170        let input = usage.get("prompt_tokens").and_then(Value::as_u64);
171        let total = usage.get("total_tokens").and_then(Value::as_u64);
172        if input.is_some() || total.is_some() {
173            batch.usage = Some(UsageDetails {
174                input_token_count: input,
175                total_token_count: total,
176                ..Default::default()
177            });
178        }
179    }
180    Ok(batch)
181}
182
183#[async_trait::async_trait]
184impl EmbeddingClient for OpenAIEmbeddingClient {
185    async fn get_embeddings(
186        &self,
187        values: Vec<String>,
188        options: Option<EmbeddingGenerationOptions>,
189    ) -> Result<GeneratedEmbeddings> {
190        let body = self.build_body(&values, options.as_ref());
191        let url = format!("{}/embeddings", self.inner.base_url.trim_end_matches('/'));
192        let mut req = self
193            .inner
194            .http
195            .post(&url)
196            .bearer_auth(&self.inner.api_key)
197            .json(&body);
198        if let Some(org) = &self.inner.organization {
199            req = req.header("OpenAI-Organization", org);
200        }
201        let resp = req
202            .send()
203            .await
204            .map_err(|e| Error::service(format!("request failed: {e}")))?;
205        if !resp.status().is_success() {
206            let status = resp.status();
207            let retry_after = parse_retry_after(resp.headers());
208            let text = resp.text().await.unwrap_or_default();
209            return Err(classify_service_error(
210                status.as_u16(),
211                &text,
212                format!("OpenAI API error {status}: {text}"),
213                retry_after,
214            ));
215        }
216        let value: Value = resp
217            .json()
218            .await
219            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
220        parse_embeddings_response(&value)
221    }
222
223    fn model(&self) -> Option<&str> {
224        Some(&self.inner.model)
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn build_body_includes_model_input_and_dimensions() {
234        let client = OpenAIEmbeddingClient::new("sk-test", "text-embedding-3-small");
235        let options = EmbeddingGenerationOptions::new().with_dimensions(256);
236        let body = client.build_body(&["a".into(), "b".into()], Some(&options));
237        assert_eq!(body["model"], "text-embedding-3-small");
238        assert_eq!(body["input"], json!(["a", "b"]));
239        assert_eq!(body["dimensions"], 256);
240    }
241
242    #[test]
243    fn build_body_option_model_overrides_default() {
244        let client = OpenAIEmbeddingClient::new("sk-test", "text-embedding-3-small");
245        let options = EmbeddingGenerationOptions::new().with_model("text-embedding-3-large");
246        let body = client.build_body(&["a".into()], Some(&options));
247        assert_eq!(body["model"], "text-embedding-3-large");
248    }
249
250    #[test]
251    fn build_body_forwards_known_additional_properties_only() {
252        let client = OpenAIEmbeddingClient::new("sk-test", "m");
253        let mut options = EmbeddingGenerationOptions::new();
254        options
255            .additional_properties
256            .insert("encoding_format".into(), json!("float"));
257        options
258            .additional_properties
259            .insert("unrelated".into(), json!(true));
260        let body = client.build_body(&["a".into()], Some(&options));
261        assert_eq!(body["encoding_format"], "float");
262        assert!(body.get("unrelated").is_none());
263    }
264
265    #[test]
266    fn parse_response_restores_index_order_and_usage() {
267        let value = json!({
268            "model": "text-embedding-3-small",
269            "data": [
270                { "index": 1, "embedding": [0.3, 0.4] },
271                { "index": 0, "embedding": [0.1, 0.2] },
272            ],
273            "usage": { "prompt_tokens": 5, "total_tokens": 5 }
274        });
275        let batch = parse_embeddings_response(&value).unwrap();
276        assert_eq!(batch.len(), 2);
277        assert_eq!(batch[0].vector, vec![0.1, 0.2]);
278        assert_eq!(batch[1].vector, vec![0.3, 0.4]);
279        assert_eq!(batch[0].model.as_deref(), Some("text-embedding-3-small"));
280        let usage = batch.usage.as_ref().unwrap();
281        assert_eq!(usage.input_token_count, Some(5));
282        assert_eq!(usage.total_token_count, Some(5));
283    }
284
285    #[test]
286    fn parse_response_missing_data_errors() {
287        assert!(parse_embeddings_response(&json!({})).is_err());
288    }
289}
agent_framework_openai/embeddings.rs

agent_framework_openai/
embeddings.rs
