1use bytes::Bytes;
4use futures::{Stream, StreamExt};
5use reqwest::Client;
6use serde::Deserialize;
7use serde_json::Value as JsonValue;
8use serde_json::json;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::Arc;
12
13use crate::{
14 Api, AssistantMessage, ContentBlock, Context, Model, Provider, ProviderEvent, StopReason,
15 StreamOptions, StreamResult, Usage, error::ProviderError,
16};
17
18use super::shared_client;
19
20#[derive(Clone)]
35pub struct AzureProvider {
36 client: &'static Client,
37 api_key: Option<String>,
38 resource_name: Option<String>,
39 deployment_name: Option<String>,
40}
41
42impl AzureProvider {
43 pub fn new() -> Self {
47 Self {
48 client: shared_client(),
49 api_key: None,
50 resource_name: None,
51 deployment_name: None,
52 }
53 }
54
55 #[cfg(test)]
57 pub fn with_config(
58 api_key: impl Into<String>,
59 resource_name: impl Into<String>,
60 deployment_name: impl Into<String>,
61 ) -> Self {
62 Self {
63 client: shared_client(),
64 api_key: Some(api_key.into()),
65 resource_name: Some(resource_name.into()),
66 deployment_name: Some(deployment_name.into()),
67 }
68 }
69
70 fn build_url(&self, model: &Model) -> Result<String, ProviderError> {
72 if !model.base_url.is_empty() && model.base_url != "https://api.openai.com" {
74 return Ok(format!(
76 "{}/chat/completions?api-version=2024-02-15-preview",
77 model.base_url.trim_end_matches('/')
78 ));
79 }
80
81 let resource = self.resource_name.as_ref().ok_or_else(|| {
83 ProviderError::InvalidResponse("AZURE_OPENAI_RESOURCE_NAME not set".into())
84 })?;
85
86 let deployment = self.deployment_name.as_ref().ok_or_else(|| {
87 ProviderError::InvalidResponse("AZURE_OPENAI_DEPLOYMENT_NAME not set".into())
88 })?;
89
90 let url = format!(
91 "https://{}.openai.azure.com/openai/deployments/{}/chat/completions?api-version=2024-02-15-preview",
92 resource, deployment
93 );
94
95 Ok(url)
96 }
97
98 fn get_api_key(&self, options: &Option<StreamOptions>) -> Result<String, ProviderError> {
100 options
101 .as_ref()
102 .and_then(|o| o.api_key.as_ref())
103 .or(self.api_key.as_ref())
104 .cloned()
105 .ok_or_else(|| ProviderError::MissingApiKey)
106 }
107
108 fn build_headers(
110 &self,
111 api_key: &str,
112 options: &Option<StreamOptions>,
113 ) -> Result<reqwest::header::HeaderMap, ProviderError> {
114 let mut headers = reqwest::header::HeaderMap::new();
115
116 headers.insert(
118 "api-key",
119 api_key.parse().map_err(|e| {
120 ProviderError::InvalidResponse(format!("invalid header value: {e}"))
121 })?,
122 );
123 headers.insert(
124 reqwest::header::CONTENT_TYPE,
125 "application/json".parse().map_err(|e| {
126 ProviderError::InvalidResponse(format!("invalid header value: {e}"))
127 })?,
128 );
129
130 if let Some(opts) = options {
132 for (k, v) in &opts.headers {
133 if let (Ok(name), Ok(value)) = (
134 k.parse::<reqwest::header::HeaderName>(),
135 v.parse::<reqwest::header::HeaderValue>(),
136 ) {
137 headers.insert(name, value);
138 }
139 }
140 }
141
142 Ok(headers)
143 }
144}
145
146impl Default for AzureProvider {
147 fn default() -> Self {
148 Self::new()
149 }
150}
151
152impl Provider for AzureProvider {
153 fn stream<'a>(
154 &'a self,
155 model: &'a Model,
156 context: &'a Context,
157 options: Option<StreamOptions>,
158 ) -> Pin<Box<dyn Future<Output = StreamResult> + Send + 'a>> {
159 Box::pin(async move {
160 let url = self.build_url(model)?;
162
163 let api_key = self.get_api_key(&options)?;
165
166 let messages = build_messages(context)?;
168
169 let mut body = serde_json::json!({
171 "messages": messages,
172 "stream": true,
173 });
174
175 if model.id != "default" && model.id != "azure" {
177 body["model"] = serde_json::json!(model.id);
178 }
179
180 if let Some(ref opts) = options {
182 if let Some(temp) = opts.temperature {
183 body["temperature"] = serde_json::json!(temp);
184 }
185
186 if let Some(max) = opts.max_tokens {
187 body["max_tokens"] = serde_json::json!(max);
188 }
189 }
190
191 if !context.tools.is_empty() {
193 body["tools"] = build_tools(&context.tools)?;
194 }
195
196 if let Some(choice) = options
198 .as_ref()
199 .and_then(|o| build_tool_choice(o.tool_choice.as_ref()))
200 {
201 body["tool_choice"] = choice;
202 }
203
204 let headers = self.build_headers(&api_key, &options)?;
206
207 let response = self
209 .client
210 .post(&url)
211 .headers(headers)
212 .json(&body)
213 .send()
214 .await
215 .map_err(ProviderError::RequestFailed)?;
216
217 if !response.status().is_success() {
218 let status = response.status();
219 let body: String = response.text().await.unwrap_or_default();
220 return Err(ProviderError::HttpError(
221 crate::error::HttpErrorDetail::new(status.as_u16(), body),
222 ));
223 }
224
225 let provider_name = model.provider.clone();
227 let model_id = model.id.clone();
228
229 let stream = response
230 .bytes_stream()
231 .scan(
232 Vec::<u8>::new(),
233 move |pending_bytes, chunk: Result<Bytes, reqwest::Error>| {
234 let pn = provider_name.clone();
235 let mid = model_id.clone();
236 let events = match chunk {
241 Ok(bytes) => {
242 let mut combined =
247 Vec::with_capacity(pending_bytes.len() + bytes.len());
248 combined.extend_from_slice(pending_bytes);
249 combined.extend_from_slice(&bytes);
250 let (text, trailing) = super::sse::split_complete_lines(&combined);
251 *pending_bytes = trailing;
252 parse_sse_events(&text, &pn, &mid)
253 }
254 Err(e) => vec![ProviderEvent::Error {
255 reason: StopReason::Error,
256 error: create_error_message(&e.to_string(), &pn, &mid),
257 }],
258 };
259 std::future::ready(Some(futures::stream::iter(events)))
260 },
261 )
262 .flatten();
263
264 Ok(Box::pin(stream) as Pin<Box<dyn Stream<Item = ProviderEvent> + Send>>)
265 })
266 }
267}
268
269fn build_messages(context: &Context) -> Result<Vec<JsonValue>, ProviderError> {
271 let mut messages = Vec::new();
272
273 if let Some(ref prompt) = context.system_prompt {
275 messages.push(serde_json::json!({
276 "role": "system",
277 "content": prompt,
278 }));
279 }
280
281 for msg in &context.messages {
283 match msg {
284 crate::Message::User(u) => {
285 let content: String = match &u.content {
286 crate::MessageContent::Text(s) => s.clone(),
287 crate::MessageContent::Blocks(blocks) => blocks_to_content(blocks)?.to_string(),
288 };
289 messages.push(serde_json::json!({
290 "role": "user",
291 "content": content,
292 }));
293 }
294 crate::Message::Assistant(a) => {
295 let content = blocks_to_content(&a.content)?.to_string();
296 messages.push(serde_json::json!({
297 "role": "assistant",
298 "content": content,
299 }));
300 }
301 crate::Message::ToolResult(t) => {
302 let content = blocks_to_content(&t.content)?.to_string();
303 messages.push(serde_json::json!({
304 "role": "tool",
305 "tool_call_id": t.tool_call_id,
306 "tool_name": t.tool_name,
307 "content": content,
308 }));
309 }
310 }
311 }
312
313 Ok(messages)
314}
315
316fn blocks_to_content(blocks: &[ContentBlock]) -> Result<JsonValue, ProviderError> {
318 if blocks.len() == 1
319 && let Some(text) = blocks[0].as_text()
320 {
321 return Ok(JsonValue::String(text.to_string()));
322 }
323
324 let items: Result<Vec<_>, _> = blocks
325 .iter()
326 .map(|block| match block {
327 ContentBlock::Text(t) => Ok(serde_json::json!({
328 "type": "text",
329 "text": t.text,
330 })),
331 ContentBlock::ToolCall(tc) => Ok(serde_json::json!({
332 "type": "function",
333 "id": tc.id,
334 "function": {
335 "name": tc.name,
336 "arguments": tc.arguments.to_string(),
337 },
338 })),
339 ContentBlock::Thinking(th) => Ok(serde_json::json!({
340 "type": "thinking",
341 "thinking": th.thinking,
342 })),
343 ContentBlock::Image(img) => Ok(serde_json::json!({
344 "type": "image_url",
345 "image_url": {
346 "url": format!("data:{};base64,{}", img.mime_type, img.data),
347 },
348 })),
349 ContentBlock::Unknown(_) => Err(ProviderError::InvalidResponse(
350 "Unknown content block type".into(),
351 )),
352 })
353 .collect();
354
355 Ok(serde_json::json!(items?))
356}
357
358fn build_tool_choice(tool_choice: Option<&crate::tools::ToolChoice>) -> Option<JsonValue> {
360 match tool_choice {
361 None | Some(crate::tools::ToolChoice::Auto) => None,
362 Some(crate::tools::ToolChoice::Named(name)) => {
363 Some(json!({"type": "function", "function": {"name": name}}))
364 }
365 }
366}
367
368fn build_tools(tools: &[crate::Tool]) -> Result<JsonValue, ProviderError> {
370 let items: Vec<_> = tools
371 .iter()
372 .map(|tool| {
373 serde_json::json!({
374 "type": "function",
375 "function": {
376 "name": tool.name,
377 "description": tool.description,
378 "parameters": tool.parameters,
379 },
380 })
381 })
382 .collect();
383
384 Ok(serde_json::json!(items))
385}
386
387fn parse_sse_events(text: &str, provider: &str, model_id: &str) -> Vec<ProviderEvent> {
391 let mut events = Vec::with_capacity(text.len() / 80);
392 let mut partial_message = AssistantMessage::new(Api::OpenAiCompletions, provider, model_id);
393
394 let mut accumulated_usage = Usage::default();
395
396 for line in text.split('\n') {
397 let line = line.trim_end_matches('\r');
398 if line.is_empty() {
399 continue;
400 }
401
402 if !line.starts_with("data: ") {
404 continue;
405 }
406
407 let data = &line[6..]; if data == "[DONE]" {
411 break;
412 }
413
414 if data.is_empty() {
415 continue;
416 }
417
418 let chunk = match serde_json::from_str::<SSEChunk>(data) {
419 Ok(c) => c,
420 Err(_) => continue,
421 };
422
423 let this_chunk_usage = chunk.usage.as_ref();
425
426 for choice in &chunk.choices {
427 if let Some(delta) = &choice.delta {
428 if let Some(content) = &delta.content {
429 let last_text_idx = partial_message
432 .content
433 .iter()
434 .rposition(|b| matches!(b, ContentBlock::Text(_)));
435 if let Some(idx) = last_text_idx
436 && let ContentBlock::Text(t) = &mut partial_message.content[idx]
437 {
438 t.text.push_str(content);
439 } else {
440 partial_message
441 .content
442 .push(ContentBlock::Text(crate::TextContent::new(content.clone())));
443 }
444 events.push(ProviderEvent::TextDelta {
445 content_index: choice.index,
446 delta: content.clone(),
447 partial: Arc::new(partial_message.clone()),
448 });
449 }
450
451 if let Some(tool_calls) = &delta.tool_calls {
452 for tc in tool_calls {
453 if let Some(func) = &tc.function {
454 events.push(ProviderEvent::ToolCallDelta {
455 content_index: choice.index,
456 delta: func.arguments.clone().unwrap_or_default(),
457 partial: Arc::new(partial_message.clone()),
458 });
459 }
460 }
461 }
462 }
463
464 if choice.finish_reason.is_some() {
465 let reason = match choice.finish_reason.as_deref() {
468 Some("stop") => StopReason::Stop,
469 Some("length") => StopReason::Length,
470 Some("tool_calls") => StopReason::ToolUse,
471 _ => StopReason::Stop,
472 };
473
474 let mut done_msg = partial_message.clone();
475
476 if let Some(usage) = this_chunk_usage {
478 done_msg.usage.input = usage.prompt_tokens;
479 done_msg.usage.output = usage.completion_tokens;
480 done_msg.usage.cache_read = usage
481 .prompt_tokens_details
482 .as_ref()
483 .map(|d| d.cached_tokens)
484 .unwrap_or(0);
485 done_msg.usage.total_tokens = usage.total_tokens;
486 } else {
487 done_msg.usage = accumulated_usage.clone();
488 }
489
490 events.push(ProviderEvent::Done {
491 reason,
492 message: done_msg,
493 });
494 }
495 }
496
497 if let Some(usage) = this_chunk_usage {
499 accumulated_usage.input = usage.prompt_tokens;
500 accumulated_usage.output = usage.completion_tokens;
501 accumulated_usage.cache_read = usage
502 .prompt_tokens_details
503 .as_ref()
504 .map(|d| d.cached_tokens)
505 .unwrap_or(0);
506 accumulated_usage.total_tokens = usage.total_tokens;
507 }
508 }
509
510 events
511}
512
513fn create_error_message(msg: &str, provider: &str, model_id: &str) -> AssistantMessage {
515 let mut message = AssistantMessage::new(Api::OpenAiCompletions, provider, model_id);
516 message.stop_reason = StopReason::Error;
517 message.error_message = Some(msg.to_string());
518 message
519}
520
521#[derive(Debug, Deserialize)]
523struct SSEChunk {
524 _id: Option<String>,
525 #[serde(rename = "model")]
526 _model: Option<String>,
527 choices: Vec<Choice>,
528 usage: Option<UsageInfo>,
529}
530
531#[derive(Debug, Deserialize)]
532struct Choice {
533 index: usize,
534 delta: Option<Delta>,
535 finish_reason: Option<String>,
536}
537
538#[derive(Debug, Deserialize)]
539struct Delta {
540 content: Option<String>,
541 tool_calls: Option<Vec<ToolCallDelta>>,
542}
543
544#[derive(Debug, Deserialize)]
545struct ToolCallDelta {
546 _index: Option<usize>,
547 _id: Option<String>,
548 #[serde(rename = "type")]
549 _type_: Option<String>,
550 function: Option<FunctionDelta>,
551}
552
553#[derive(Debug, Deserialize)]
554struct FunctionDelta {
555 _name: Option<String>,
556 arguments: Option<String>,
557}
558
559#[derive(Debug, Deserialize, Clone)]
560struct UsageInfo {
561 prompt_tokens: usize,
562 completion_tokens: usize,
563 total_tokens: usize,
564 #[serde(rename = "prompt_tokens_details")]
565 prompt_tokens_details: Option<PromptTokensDetails>,
566}
567
568#[derive(Debug, Deserialize, Clone)]
569struct PromptTokensDetails {
570 #[serde(rename = "cached_tokens")]
571 cached_tokens: usize,
572}
573
574#[cfg(test)]
575mod tests {
576 use super::*;
577
578 #[test]
579 fn build_tool_choice_maps_named_to_azure_shape() {
580 assert!(build_tool_choice(None).is_none());
581 assert!(build_tool_choice(Some(&crate::tools::ToolChoice::Auto)).is_none());
582 assert_eq!(
583 build_tool_choice(Some(&crate::tools::ToolChoice::Named("todo".into()))),
584 Some(serde_json::json!({"type": "function", "function": {"name": "todo"}}))
585 );
586 }
587
588 fn make_test_model(id: &str, base_url: &str) -> Model {
589 Model::new(id, id, Api::OpenAiCompletions, "azure", base_url)
590 }
591
592 #[test]
593 fn test_build_url_from_base_url() {
594 let provider = AzureProvider::new();
595 let model = make_test_model(
596 "gpt-4o",
597 "https://my-resource.openai.azure.com/openai/deployments/gpt-4o",
598 );
599
600 let url = provider.build_url(&model).unwrap();
601 assert!(url.contains("api-version=2024-02-15-preview"));
602 assert!(url.contains("my-resource"));
603 assert!(url.contains("gpt-4o"));
604 }
605
606 #[test]
607 fn test_build_url_missing_resource() {
608 let provider = AzureProvider {
609 client: shared_client(),
610 api_key: Some("test-key".to_string()),
611 resource_name: None,
612 deployment_name: Some("gpt-4o".to_string()),
613 };
614
615 let model = make_test_model("default", "");
616
617 let result = provider.build_url(&model);
618 assert!(result.is_err());
619 match result.unwrap_err() {
620 ProviderError::InvalidResponse(msg) => {
621 assert!(msg.contains("AZURE_OPENAI_RESOURCE_NAME"));
622 }
623 _ => panic!("Expected InvalidResponse"),
624 }
625 }
626
627 #[test]
628 fn test_build_url_missing_deployment() {
629 let provider = AzureProvider {
630 client: shared_client(),
631 api_key: Some("test-key".to_string()),
632 resource_name: Some("my-resource".to_string()),
633 deployment_name: None,
634 };
635
636 let model = make_test_model("default", "");
637
638 let result = provider.build_url(&model);
639 assert!(result.is_err());
640 match result.unwrap_err() {
641 ProviderError::InvalidResponse(msg) => {
642 assert!(msg.contains("AZURE_OPENAI_DEPLOYMENT_NAME"));
643 }
644 _ => panic!("Expected InvalidResponse"),
645 }
646 }
647
648 #[test]
649 fn test_build_url_from_env_vars() {
650 let provider = AzureProvider {
651 client: shared_client(),
652 api_key: Some("test-key".to_string()),
653 resource_name: Some("my-resource".to_string()),
654 deployment_name: Some("gpt-4o".to_string()),
655 };
656
657 let model = make_test_model("default", "");
658
659 let url = provider.build_url(&model).unwrap();
660 assert_eq!(
661 url,
662 "https://my-resource.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-02-15-preview"
663 );
664 }
665
666 #[test]
667 fn test_parse_sse_events_text() {
668 let sse_data = r#"data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}
669
670data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":"stop"}]}
671
672data: [DONE]"#;
673
674 let events = parse_sse_events(sse_data, "azure", "gpt-4o");
675
676 assert!(events.len() >= 3);
678
679 match &events[0] {
681 ProviderEvent::TextDelta { delta, .. } => assert_eq!(delta, "Hello"),
682 _ => panic!("Expected TextDelta event"),
683 }
684
685 match &events[events.len() - 1] {
687 ProviderEvent::Done { reason, .. } => assert_eq!(*reason, StopReason::Stop),
688 _ => panic!("Expected Done event"),
689 }
690 }
691
692 #[test]
693 fn test_parse_sse_events_with_tool_calls() {
694 let sse_data = r#"data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"id":"call_abc123","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}
695
696data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"location\":"}}]},"finish_reason":null}]}
697
698data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"Boston\"}"}}]},"finish_reason":"tool_calls"}]}
699
700data: [DONE]"#;
701
702 let events = parse_sse_events(sse_data, "azure", "gpt-4o");
703
704 assert!(events.len() >= 4);
706
707 let has_tool_call = events
709 .iter()
710 .any(|e| matches!(e, ProviderEvent::ToolCallDelta { .. }));
711 assert!(
712 has_tool_call,
713 "Should have at least one ToolCallDelta event"
714 );
715
716 match &events[events.len() - 1] {
718 ProviderEvent::Done { reason, .. } => assert_eq!(*reason, StopReason::ToolUse),
719 _ => panic!("Expected Done event with ToolUse reason"),
720 }
721 }
722
723 #[test]
724 fn test_parse_sse_events_usage() {
725 let sse_data = r#"data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"prompt_tokens_details":{"cached_tokens":0}}}
726
727data: [DONE]"#;
728
729 let events = parse_sse_events(sse_data, "azure", "gpt-4o");
730
731 let done_event = events
733 .iter()
734 .find(|e| matches!(e, ProviderEvent::Done { .. }));
735 assert!(done_event.is_some());
736
737 if let ProviderEvent::Done { message, .. } = done_event.unwrap() {
738 assert_eq!(message.usage.input, 10);
739 assert_eq!(message.usage.output, 5);
740 assert_eq!(message.usage.total_tokens, 15);
741 }
742 }
743
744 #[test]
745 fn test_build_headers_includes_api_key() {
746 let provider = AzureProvider::new();
747 let api_key = "test-api-key-12345";
748
749 let headers = provider.build_headers(api_key, &None).unwrap();
750
751 let api_key_header = headers.get("api-key");
753 assert!(api_key_header.is_some());
754 assert_eq!(api_key_header.unwrap().to_str().unwrap(), api_key);
755
756 let content_type = headers.get(reqwest::header::CONTENT_TYPE);
758 assert!(content_type.is_some());
759 }
760
761 #[test]
762 fn test_build_headers_no_bearer_token() {
763 let provider = AzureProvider::new();
764 let api_key = "test-api-key-12345";
765
766 let headers = provider.build_headers(api_key, &None).unwrap();
767
768 let auth_header = headers.get(reqwest::header::AUTHORIZATION);
770 assert!(
771 auth_header.is_none(),
772 "Azure should not use Bearer token authentication"
773 );
774 }
775
776 #[test]
777 fn test_with_config_constructor() {
778 let provider = AzureProvider::with_config("my-api-key", "my-resource", "gpt-4o");
779
780 let model = make_test_model("default", "");
782
783 let url = provider.build_url(&model).unwrap();
784 assert!(url.contains("my-resource"));
785 assert!(url.contains("gpt-4o"));
786 }
787
788 #[test]
789 fn test_azure_endpoint_format() {
790 let provider = AzureProvider {
791 client: shared_client(),
792 api_key: Some("key".to_string()),
793 resource_name: Some("my-resource".to_string()),
794 deployment_name: Some("gpt-4-turbo".to_string()),
795 };
796
797 let model = make_test_model("default", "");
798 let url = provider.build_url(&model).unwrap();
799
800 assert!(url.starts_with("https://"));
802 assert!(url.contains(".openai.azure.com"));
803 assert!(url.contains("/openai/deployments/"));
804 assert!(url.contains("gpt-4-turbo"));
805 assert!(url.contains("chat/completions"));
806 assert!(url.contains("api-version=2024-02-15-preview"));
807 }
808}