agent_framework_azure/
embeddings.rs

1//! Azure OpenAI embeddings client.
2//!
3//! Rust equivalent of upstream's Azure branch of the OpenAI embedding client
4//! (`agent_framework_openai/_embedding_client.py`): the deployment-scoped
5//! `POST {endpoint}/openai/deployments/{deployment}/embeddings` endpoint with
6//! `api-key` or Entra ID auth. The wire body/response are OpenAI-shaped, so
7//! request extras and response parsing are shared with
8//! `agent-framework-openai` rather than duplicated.
9
10use std::sync::Arc;
11
12use agent_framework_core::client::EmbeddingClient;
13use agent_framework_core::error::{Error, Result};
14use agent_framework_core::types::{EmbeddingGenerationOptions, GeneratedEmbeddings};
15use serde_json::{json, Map, Value};
16
17use crate::credential::TokenCredential;
18use crate::{parse_retry_after, DEFAULT_API_VERSION};
19
20enum Auth {
21    ApiKey(String),
22    Credential(Arc<dyn TokenCredential>),
23}
24
25/// An Azure OpenAI embeddings client
26/// (`POST {endpoint}/openai/deployments/{deployment}/embeddings`).
27#[derive(Clone)]
28pub struct AzureOpenAIEmbeddingClient {
29    inner: Arc<Inner>,
30}
31
32struct Inner {
33    http: reqwest::Client,
34    endpoint: String,
35    deployment: String,
36    api_version: String,
37    auth: Auth,
38}
39
40impl AzureOpenAIEmbeddingClient {
41    /// Create a client authenticating with a static API key
42    /// (`api-key` header).
43    pub fn new(
44        endpoint: impl Into<String>,
45        deployment: impl Into<String>,
46        api_key: impl Into<String>,
47    ) -> Self {
48        Self {
49            inner: Arc::new(Inner {
50                http: reqwest::Client::new(),
51                endpoint: endpoint.into(),
52                deployment: deployment.into(),
53                api_version: DEFAULT_API_VERSION.to_string(),
54                auth: Auth::ApiKey(api_key.into()),
55            }),
56        }
57    }
58
59    /// Create a client authenticating via a [`TokenCredential`]
60    /// (`Authorization: Bearer <token>`, e.g. Microsoft Entra ID).
61    pub fn with_token_credential(
62        endpoint: impl Into<String>,
63        deployment: impl Into<String>,
64        credential: Arc<dyn TokenCredential>,
65    ) -> Self {
66        Self {
67            inner: Arc::new(Inner {
68                http: reqwest::Client::new(),
69                endpoint: endpoint.into(),
70                deployment: deployment.into(),
71                api_version: DEFAULT_API_VERSION.to_string(),
72                auth: Auth::Credential(credential),
73            }),
74        }
75    }
76
77    /// Build an API-key-authenticated client from `AZURE_OPENAI_ENDPOINT`,
78    /// `AZURE_OPENAI_API_KEY`, `AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME`, and
79    /// optional `AZURE_OPENAI_API_VERSION`.
80    pub fn from_env() -> Result<Self> {
81        let endpoint = std::env::var("AZURE_OPENAI_ENDPOINT")
82            .map_err(|_| Error::Configuration("AZURE_OPENAI_ENDPOINT is not set".into()))?;
83        let api_key = std::env::var("AZURE_OPENAI_API_KEY")
84            .map_err(|_| Error::Configuration("AZURE_OPENAI_API_KEY is not set".into()))?;
85        let deployment = std::env::var("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME").map_err(|_| {
86            Error::Configuration("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME is not set".into())
87        })?;
88        let mut client = Self::new(endpoint, deployment, api_key);
89        if let Ok(v) = std::env::var("AZURE_OPENAI_API_VERSION") {
90            client = client.with_api_version(v);
91        }
92        Ok(client)
93    }
94
95    /// Override the API version (default [`DEFAULT_API_VERSION`]).
96    ///
97    /// [`DEFAULT_API_VERSION`]: crate::AzureOpenAIClient::api_version
98    pub fn with_api_version(mut self, api_version: impl Into<String>) -> Self {
99        arc_inner(&mut self.inner).api_version = api_version.into();
100        self
101    }
102
103    /// The deployment name this client targets.
104    pub fn deployment(&self) -> &str {
105        &self.inner.deployment
106    }
107
108    fn url(&self) -> String {
109        format!(
110            "{}/openai/deployments/{}/embeddings?api-version={}",
111            self.inner.endpoint.trim_end_matches('/'),
112            self.inner.deployment,
113            self.inner.api_version,
114        )
115    }
116
117    fn build_body(&self, values: &[String], options: Option<&EmbeddingGenerationOptions>) -> Value {
118        let mut body = Map::new();
119        // The deployment in the URL already selects the model; only send
120        // `model` if the caller explicitly asked for a specific one (same
121        // convention as the chat client).
122        if let Some(model) = options.and_then(|o| o.model.as_ref()) {
123            body.insert("model".into(), json!(model));
124        }
125        body.insert("input".into(), json!(values));
126        if let Some(options) = options {
127            if let Some(dimensions) = options.dimensions {
128                body.insert("dimensions".into(), json!(dimensions));
129            }
130            for key in ["encoding_format", "user"] {
131                if let Some(v) = options.additional_properties.get(key) {
132                    body.insert(key.into(), v.clone());
133                }
134            }
135        }
136        Value::Object(body)
137    }
138
139    async fn auth_header(&self) -> Result<(&'static str, String)> {
140        match &self.inner.auth {
141            Auth::ApiKey(key) => Ok(("api-key", key.clone())),
142            Auth::Credential(credential) => {
143                let token = credential.get_token().await?;
144                Ok(("Authorization", format!("Bearer {token}")))
145            }
146        }
147    }
148}
149
150fn arc_inner(inner: &mut Arc<Inner>) -> &mut Inner {
151    if Arc::strong_count(inner) != 1 {
152        *inner = Arc::new(Inner {
153            http: inner.http.clone(),
154            endpoint: inner.endpoint.clone(),
155            deployment: inner.deployment.clone(),
156            api_version: inner.api_version.clone(),
157            auth: match &inner.auth {
158                Auth::ApiKey(k) => Auth::ApiKey(k.clone()),
159                Auth::Credential(c) => Auth::Credential(c.clone()),
160            },
161        });
162    }
163    Arc::get_mut(inner).expect("just ensured unique")
164}
165
166#[async_trait::async_trait]
167impl EmbeddingClient for AzureOpenAIEmbeddingClient {
168    async fn get_embeddings(
169        &self,
170        values: Vec<String>,
171        options: Option<EmbeddingGenerationOptions>,
172    ) -> Result<GeneratedEmbeddings> {
173        let body = self.build_body(&values, options.as_ref());
174        let (header_name, header_value) = self.auth_header().await?;
175        let resp = self
176            .inner
177            .http
178            .post(self.url())
179            .header(header_name, header_value)
180            .json(&body)
181            .send()
182            .await
183            .map_err(|e| Error::service(format!("request failed: {e}")))?;
184        if !resp.status().is_success() {
185            let status = resp.status();
186            let retry_after = parse_retry_after(resp.headers());
187            let text = resp.text().await.unwrap_or_default();
188            // Wire-compatible with OpenAI; classification shared verbatim.
189            return Err(agent_framework_openai::classify_service_error(
190                status.as_u16(),
191                &text,
192                format!("Azure OpenAI API error {status}: {text}"),
193                retry_after,
194            ));
195        }
196        let value: Value = resp
197            .json()
198            .await
199            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
200        agent_framework_openai::embeddings::parse_embeddings_response(&value)
201    }
202
203    fn model(&self) -> Option<&str> {
204        Some(&self.inner.deployment)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211
212    #[test]
213    fn url_is_deployment_scoped_with_api_version() {
214        let client = AzureOpenAIEmbeddingClient::new(
215            "https://example.openai.azure.com/",
216            "embed-dep",
217            "key",
218        )
219        .with_api_version("2024-10-21");
220        assert_eq!(
221            client.url(),
222            "https://example.openai.azure.com/openai/deployments/embed-dep/embeddings?api-version=2024-10-21"
223        );
224    }
225
226    #[test]
227    fn body_omits_model_unless_explicit() {
228        let client = AzureOpenAIEmbeddingClient::new("https://e", "dep", "key");
229        let body = client.build_body(&["x".into()], None);
230        assert!(body.get("model").is_none());
231        assert_eq!(body["input"], json!(["x"]));
232
233        let options = EmbeddingGenerationOptions::new()
234            .with_model("text-embedding-3-large")
235            .with_dimensions(64);
236        let body = client.build_body(&["x".into()], Some(&options));
237        assert_eq!(body["model"], "text-embedding-3-large");
238        assert_eq!(body["dimensions"], 64);
239    }
240}
agent_framework_azure/embeddings.rs

agent_framework_azure/
embeddings.rs
