1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{
4 RemoteProviderBase, build_openai_chat_messages, check_http_status, resolve_api_key,
5};
6use crate::traits::{
7 EmbedResult, EmbeddingModel, GenerationOptions, GenerationResult, GeneratorModel,
8 LoadedModelHandle, Message, ModelProvider, ProviderCapabilities, ProviderHealth, TokenUsage,
9};
10use async_trait::async_trait;
11use reqwest::Client;
12use serde_json::json;
13use std::sync::Arc;
14
15pub struct RemoteAzureOpenAIProvider {
21 base: RemoteProviderBase,
22}
23
24impl Default for RemoteAzureOpenAIProvider {
25 fn default() -> Self {
26 Self {
27 base: RemoteProviderBase::new(),
28 }
29 }
30}
31
32impl RemoteAzureOpenAIProvider {
33 pub fn new() -> Self {
34 Self::default()
35 }
36
37 #[cfg(test)]
38 fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
39 self.base.insert_test_breaker(key, age);
40 }
41
42 #[cfg(test)]
43 fn breaker_count(&self) -> usize {
44 self.base.breaker_count()
45 }
46
47 #[cfg(test)]
48 fn force_cleanup_now_for_test(&self) {
49 self.base.force_cleanup_now_for_test();
50 }
51}
52
53#[derive(Clone)]
56struct AzureResolvedOptions {
57 api_key: String,
58 resource_name: String,
59 api_version: String,
60 embedding_dimensions: Option<u32>,
61}
62
63impl AzureResolvedOptions {
64 fn from_spec(spec: &ModelAliasSpec) -> Result<Self> {
65 let api_key = resolve_api_key(&spec.options, "api_key_env", "AZURE_OPENAI_API_KEY")?;
66
67 let resource_name = spec
68 .options
69 .get("resource_name")
70 .and_then(|v| v.as_str())
71 .ok_or_else(|| {
72 RuntimeError::Config(
73 "Option 'resource_name' is required for Azure OpenAI provider".to_string(),
74 )
75 })?
76 .to_string();
77
78 let api_version = spec
79 .options
80 .get("api_version")
81 .and_then(|v| v.as_str())
82 .unwrap_or("2024-10-21")
83 .to_string();
84
85 let embedding_dimensions = spec
86 .options
87 .get("embedding_dimensions")
88 .and_then(|v| v.as_u64())
89 .map(|v| v as u32);
90
91 Ok(Self {
92 api_key,
93 resource_name,
94 api_version,
95 embedding_dimensions,
96 })
97 }
98
99 fn embed_url(&self, deployment: &str) -> String {
100 format!(
101 "https://{}.openai.azure.com/openai/deployments/{}/embeddings?api-version={}",
102 self.resource_name, deployment, self.api_version
103 )
104 }
105
106 fn chat_url(&self, deployment: &str) -> String {
107 format!(
108 "https://{}.openai.azure.com/openai/deployments/{}/chat/completions?api-version={}",
109 self.resource_name, deployment, self.api_version
110 )
111 }
112}
113
114#[async_trait]
115impl ModelProvider for RemoteAzureOpenAIProvider {
116 fn provider_id(&self) -> &'static str {
117 "remote/azure-openai"
118 }
119
120 fn capabilities(&self) -> ProviderCapabilities {
121 ProviderCapabilities {
122 supported_tasks: vec![ModelTask::Embed, ModelTask::Generate],
123 }
124 }
125
126 async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
127 let cb = self.base.circuit_breaker_for(spec);
128 let resolved = AzureResolvedOptions::from_spec(spec)?;
129
130 match spec.task {
131 ModelTask::Embed => {
132 let dims = resolved.embedding_dimensions.unwrap_or(1536);
133 let model = AzureOpenAIEmbeddingModel {
134 client: self.base.client.clone(),
135 cb: cb.clone(),
136 deployment: spec.model_id.clone(),
137 options: resolved,
138 dimensions: dims,
139 };
140 let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
141 Ok(Arc::new(handle) as LoadedModelHandle)
142 }
143 ModelTask::Generate => {
144 let model = AzureOpenAIGeneratorModel {
145 client: self.base.client.clone(),
146 cb,
147 deployment: spec.model_id.clone(),
148 options: resolved,
149 };
150 let handle: Arc<dyn GeneratorModel> = Arc::new(model);
151 Ok(Arc::new(handle) as LoadedModelHandle)
152 }
153 ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
154 "Azure OpenAI provider does not support task Raw".to_string(),
155 )),
156 _ => Err(RuntimeError::CapabilityMismatch(format!(
157 "Azure OpenAI provider does not support task {:?}",
158 spec.task
159 ))),
160 }
161 }
162
163 async fn health(&self) -> ProviderHealth {
164 ProviderHealth::Healthy
165 }
166}
167
168struct AzureOpenAIEmbeddingModel {
169 client: Client,
170 cb: crate::reliability::CircuitBreakerWrapper,
171 deployment: String,
172 options: AzureResolvedOptions,
173 dimensions: u32,
174}
175
176#[async_trait]
177impl EmbeddingModel for AzureOpenAIEmbeddingModel {
178 async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
179 let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
180
181 self.cb
182 .call(move || async move {
183 let url = self.options.embed_url(&self.deployment);
184
185 let response = self
186 .client
187 .post(&url)
188 .header("api-key", &self.options.api_key)
189 .json(&json!({
190 "input": texts
191 }))
192 .send()
193 .await
194 .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
195
196 let body: serde_json::Value = check_http_status("Azure OpenAI", response)?
197 .json()
198 .await
199 .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
200
201 let mut vectors = Vec::new();
202 if let Some(data) = body.get("data").and_then(|d| d.as_array()) {
203 for item in data {
204 if let Some(embedding) = item.get("embedding").and_then(|e| e.as_array()) {
205 let vec: Vec<f32> = embedding
206 .iter()
207 .filter_map(|v| v.as_f64().map(|f| f as f32))
208 .collect();
209 vectors.push(vec);
210 }
211 }
212 }
213
214 let usage = body.get("usage").map(|u| {
217 let prompt = u["prompt_tokens"].as_u64().unwrap_or(0) as usize;
218 let total = u["total_tokens"].as_u64().unwrap_or(prompt as u64) as usize;
219 TokenUsage {
220 prompt_tokens: prompt,
221 completion_tokens: 0,
222 total_tokens: total,
223 }
224 });
225
226 Ok(EmbedResult { vectors, usage })
227 })
228 .await
229 }
230
231 fn dimensions(&self) -> u32 {
232 self.dimensions
233 }
234}
235
236impl crate::traits::ModelInfo for AzureOpenAIEmbeddingModel {
237 fn model_id(&self) -> &str {
238 &self.deployment
239 }
240}
241
242struct AzureOpenAIGeneratorModel {
243 client: Client,
244 cb: crate::reliability::CircuitBreakerWrapper,
245 deployment: String,
246 options: AzureResolvedOptions,
247}
248
249impl crate::traits::ModelInfo for AzureOpenAIGeneratorModel {
250 fn model_id(&self) -> &str {
251 &self.deployment
252 }
253}
254
255#[async_trait]
256impl GeneratorModel for AzureOpenAIGeneratorModel {
257 async fn generate(
258 &self,
259 messages: &[Message],
260 options: GenerationOptions,
261 ) -> Result<GenerationResult> {
262 let messages = build_openai_chat_messages(messages);
263
264 self.cb
265 .call(move || async move {
266 let url = self.options.chat_url(&self.deployment);
267
268 let mut body = json!({
269 "messages": messages,
270 });
271
272 if let Some(max_tokens) = options.max_tokens {
273 body["max_completion_tokens"] = json!(max_tokens);
274 }
275 if let Some(temperature) = options.temperature {
276 body["temperature"] = json!(temperature);
277 }
278 if let Some(top_p) = options.top_p {
279 body["top_p"] = json!(top_p);
280 }
281
282 let response = self
283 .client
284 .post(&url)
285 .header("api-key", &self.options.api_key)
286 .json(&body)
287 .send()
288 .await
289 .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
290
291 let body: serde_json::Value = check_http_status("Azure OpenAI", response)?
292 .json()
293 .await
294 .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
295
296 let text = body["choices"][0]["message"]["content"]
297 .as_str()
298 .unwrap_or("")
299 .to_string();
300
301 let usage = body.get("usage").map(|u| TokenUsage {
302 prompt_tokens: u["prompt_tokens"].as_u64().unwrap_or(0) as usize,
303 completion_tokens: u["completion_tokens"].as_u64().unwrap_or(0) as usize,
304 total_tokens: u["total_tokens"].as_u64().unwrap_or(0) as usize,
305 });
306
307 Ok(GenerationResult {
308 text,
309 usage,
310 images: vec![],
311 audio: None,
312 })
313 })
314 .await
315 }
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321 use crate::api::ModelRuntimeKey;
322 use crate::provider::remote_common::RemoteProviderBase;
323 use crate::traits::ModelProvider;
324 use std::time::Duration;
325
326 static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
327
328 fn spec_with_opts(
329 alias: &str,
330 task: ModelTask,
331 model_id: &str,
332 options: serde_json::Value,
333 ) -> ModelAliasSpec {
334 ModelAliasSpec {
335 alias: alias.to_string(),
336 task,
337 provider_id: "remote/azure-openai".to_string(),
338 model_id: model_id.to_string(),
339 revision: None,
340 warmup: crate::api::WarmupPolicy::Lazy,
341 required: false,
342 timeout: None,
343 load_timeout: None,
344 retry: None,
345 options,
346 }
347 }
348
349 fn default_opts() -> serde_json::Value {
350 json!({ "resource_name": "my-resource" })
351 }
352
353 #[tokio::test]
354 async fn breaker_reused_for_same_runtime_key() {
355 let _lock = ENV_LOCK.lock().await;
356 unsafe { std::env::set_var("AZURE_OPENAI_API_KEY", "test-key") };
357
358 let provider = RemoteAzureOpenAIProvider::new();
359 let s1 = spec_with_opts(
360 "embed/a",
361 ModelTask::Embed,
362 "text-embedding-ada-002",
363 default_opts(),
364 );
365 let s2 = spec_with_opts(
366 "embed/b",
367 ModelTask::Embed,
368 "text-embedding-ada-002",
369 default_opts(),
370 );
371
372 let _ = provider.load(&s1).await.unwrap();
373 let _ = provider.load(&s2).await.unwrap();
374
375 assert_eq!(provider.breaker_count(), 1);
376
377 unsafe { std::env::remove_var("AZURE_OPENAI_API_KEY") };
378 }
379
380 #[tokio::test]
381 async fn breaker_cleanup_evicts_stale_entries() {
382 let _lock = ENV_LOCK.lock().await;
383 unsafe { std::env::set_var("AZURE_OPENAI_API_KEY", "test-key") };
384
385 let provider = RemoteAzureOpenAIProvider::new();
386 let stale = spec_with_opts(
387 "embed/stale",
388 ModelTask::Embed,
389 "text-embedding-ada-002",
390 default_opts(),
391 );
392 let fresh = spec_with_opts("chat/fresh", ModelTask::Generate, "gpt-4o", default_opts());
393 provider.insert_test_breaker(
394 ModelRuntimeKey::new(&stale),
395 RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
396 );
397 provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
398 assert_eq!(provider.breaker_count(), 2);
399
400 provider.force_cleanup_now_for_test();
401 let _ = provider.load(&fresh).await.unwrap();
402
403 assert_eq!(provider.breaker_count(), 1);
404
405 unsafe { std::env::remove_var("AZURE_OPENAI_API_KEY") };
406 }
407
408 #[tokio::test]
409 async fn load_fails_without_resource_name() {
410 let _lock = ENV_LOCK.lock().await;
411 unsafe { std::env::set_var("AZURE_OPENAI_API_KEY", "test-key") };
412
413 let provider = RemoteAzureOpenAIProvider::new();
414 let s = spec_with_opts(
415 "embed/a",
416 ModelTask::Embed,
417 "text-embedding-ada-002",
418 serde_json::Value::Null,
419 );
420 let result = provider.load(&s).await;
421 assert!(result.is_err());
422 assert!(result.unwrap_err().to_string().contains("resource_name"));
423
424 unsafe { std::env::remove_var("AZURE_OPENAI_API_KEY") };
425 }
426
427 #[tokio::test]
428 async fn rerank_capability_mismatch() {
429 let _lock = ENV_LOCK.lock().await;
430 unsafe { std::env::set_var("AZURE_OPENAI_API_KEY", "test-key") };
431
432 let provider = RemoteAzureOpenAIProvider::new();
433 let s = spec_with_opts(
434 "rerank/a",
435 ModelTask::Rerank,
436 "text-embedding-ada-002",
437 default_opts(),
438 );
439 let result = provider.load(&s).await;
440 assert!(result.is_err());
441 assert!(
442 result
443 .unwrap_err()
444 .to_string()
445 .contains("does not support task")
446 );
447
448 unsafe { std::env::remove_var("AZURE_OPENAI_API_KEY") };
449 }
450
451 #[test]
452 fn azure_url_construction() {
453 let opts = AzureResolvedOptions {
454 api_key: "key".to_string(),
455 resource_name: "my-resource".to_string(),
456 api_version: "2024-10-21".to_string(),
457 embedding_dimensions: None,
458 };
459
460 assert_eq!(
461 opts.embed_url("text-embedding-ada-002"),
462 "https://my-resource.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-10-21"
463 );
464
465 assert_eq!(
466 opts.chat_url("gpt-4o"),
467 "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21"
468 );
469 }
470
471 #[tokio::test]
472 async fn default_embedding_dimensions() {
473 let _lock = ENV_LOCK.lock().await;
474 unsafe { std::env::set_var("AZURE_OPENAI_API_KEY", "test-key") };
475
476 let provider = RemoteAzureOpenAIProvider::new();
477 let s = spec_with_opts(
478 "embed/dim",
479 ModelTask::Embed,
480 "text-embedding-ada-002",
481 default_opts(),
482 );
483 let handle = provider.load(&s).await.unwrap();
484 let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
485 assert_eq!(model.dimensions(), 1536);
486
487 unsafe { std::env::remove_var("AZURE_OPENAI_API_KEY") };
488 }
489
490 #[tokio::test]
491 async fn custom_embedding_dimensions() {
492 let _lock = ENV_LOCK.lock().await;
493 unsafe { std::env::set_var("AZURE_OPENAI_API_KEY", "test-key") };
494
495 let provider = RemoteAzureOpenAIProvider::new();
496 let s = spec_with_opts(
497 "embed/dim-custom",
498 ModelTask::Embed,
499 "text-embedding-ada-002",
500 json!({ "resource_name": "my-resource", "embedding_dimensions": 256 }),
501 );
502 let handle = provider.load(&s).await.unwrap();
503 let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
504 assert_eq!(model.dimensions(), 256);
505
506 unsafe { std::env::remove_var("AZURE_OPENAI_API_KEY") };
507 }
508}