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