use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;
use tokio_stream::StreamExt;
use crate::core::error::ProviderError;
use crate::core::TokenUsage;
use crate::llm::{Provider, ProviderEvent, ProviderResponse, Model};
use crate::storage::{Message, MessageRole};
use crate::tools::Tool;
pub struct AzureProvider {
client: Client,
api_key: String,
endpoint: String,
deployment_name: String,
api_version: String,
model: Model,
max_tokens: Option<u32>,
temperature: Option<f32>,
}
#[derive(Debug, Serialize)]
struct AzureRequest {
messages: Vec<AzureMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<serde_json::Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_choice: Option<String>,
stream: bool,
}
#[derive(Debug, Serialize)]
struct AzureMessage {
role: String,
content: String,
}
#[derive(Debug, Deserialize)]
struct AzureResponse {
choices: Vec<AzureChoice>,
usage: Option<AzureUsage>,
}
#[derive(Debug, Deserialize)]
struct AzureChoice {
message: AzureResponseMessage,
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct AzureResponseMessage {
role: String,
content: Option<String>,
tool_calls: Option<Vec<serde_json::Value>>,
}
#[derive(Debug, Deserialize)]
struct AzureUsage {
prompt_tokens: u32,
completion_tokens: u32,
total_tokens: u32,
}
#[derive(Debug, Deserialize)]
struct AzureStreamChunk {
choices: Vec<AzureStreamChoice>,
}
#[derive(Debug, Deserialize)]
struct AzureStreamChoice {
delta: AzureDelta,
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct AzureDelta {
role: Option<String>,
content: Option<String>,
}
impl AzureProvider {
pub fn new(
api_key: String,
endpoint: String,
deployment_name: String,
api_version: Option<String>,
model: Model,
max_tokens: Option<u32>,
temperature: Option<f32>,
) -> Result<Self, ProviderError> {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| ProviderError::Configuration(format!("Failed to create HTTP client: {}", e)))?;
let endpoint = endpoint.trim_end_matches('/').to_string();
Ok(Self {
client,
api_key,
endpoint,
deployment_name,
api_version: api_version.unwrap_or_else(|| "2024-02-15-preview".to_string()),
model,
max_tokens,
temperature,
})
}
fn get_api_url(&self) -> String {
format!(
"{}/openai/deployments/{}/chat/completions?api-version={}",
self.endpoint, self.deployment_name, self.api_version
)
}
fn convert_messages(&self, messages: Vec<Message>) -> Vec<AzureMessage> {
messages
.into_iter()
.map(|msg| AzureMessage {
role: msg.role.to_string(),
content: msg.content.as_text(),
})
.collect()
}
fn create_request(&self, messages: Vec<Message>, stream: bool) -> AzureRequest {
AzureRequest {
messages: self.convert_messages(messages),
max_tokens: self.max_tokens,
temperature: self.temperature,
tools: None, tool_choice: None,
stream,
}
}
fn parse_response(&self, response: AzureResponse) -> Result<ProviderResponse, ProviderError> {
let choice = response.choices
.into_iter()
.next()
.ok_or_else(|| ProviderError::ResponseParsing("No choices in response".to_string()))?;
let content = choice.message.content
.unwrap_or_default();
let token_usage = response.usage.map(|usage| TokenUsage {
input_tokens: usage.prompt_tokens,
output_tokens: usage.completion_tokens,
total_tokens: usage.prompt_tokens + usage.completion_tokens,
cache_creation_tokens: 0,
cache_read_tokens: 0,
});
Ok(ProviderResponse {
content,
tool_calls: Vec::new(), token_usage,
metadata: serde_json::json!({
"finish_reason": choice.finish_reason,
"model": self.model.id,
"provider": "azure",
"deployment": self.deployment_name
}),
})
}
}
#[async_trait]
impl Provider for AzureProvider {
async fn send_messages(
&self,
messages: Vec<Message>,
_tools: Vec<Box<dyn Tool>>,
) -> Result<ProviderResponse, ProviderError> {
let request = self.create_request(messages, false);
let response = self.client
.post(&self.get_api_url())
.header("api-key", &self.api_key)
.header("Content-Type", "application/json")
.json(&request)
.send()
.await
.map_err(|e| ProviderError::Http(e))?;
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
return Err(ProviderError::ApiError(format!(
"Azure API error {}: {}", status, error_text
)));
}
let azure_response: AzureResponse = response.json().await
.map_err(|e| ProviderError::ResponseParsing(format!("Failed to parse Azure response: {}", e)))?;
self.parse_response(azure_response)
}
async fn stream_response(
&self,
messages: Vec<Message>,
_tools: Vec<Box<dyn Tool>>,
) -> Result<mpsc::Receiver<Result<ProviderEvent, ProviderError>>, ProviderError> {
let (tx, rx) = mpsc::channel(100);
let request = self.create_request(messages, true);
let client = self.client.clone();
let url = self.get_api_url();
let api_key = self.api_key.clone();
tokio::spawn(async move {
let response_result = client
.post(&url)
.header("api-key", &api_key)
.header("Content-Type", "application/json")
.json(&request)
.send()
.await;
match response_result {
Ok(response) => {
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
let _ = tx.send(Err(ProviderError::ApiError(format!(
"Azure API error {}: {}", status, error_text
)))).await;
return;
}
let mut stream = response.bytes_stream();
let mut buffer = String::new();
while let Some(chunk_result) = stream.next().await {
match chunk_result {
Ok(chunk) => {
let chunk_str = String::from_utf8_lossy(&chunk);
buffer.push_str(&chunk_str);
while let Some(line_end) = buffer.find('\n') {
let line = buffer[..line_end].trim().to_string();
buffer = buffer[line_end + 1..].to_string();
if line.starts_with("data: ") {
let data = &line[6..];
if data == "[DONE]" {
let _ = tx.send(Ok(ProviderEvent::Complete { token_usage: None })).await;
return;
}
if let Ok(chunk) = serde_json::from_str::<AzureStreamChunk>(data) {
for choice in chunk.choices {
if let Some(content) = choice.delta.content {
let _ = tx.send(Ok(ProviderEvent::ContentChunk { content })).await;
}
if choice.finish_reason.is_some() {
let _ = tx.send(Ok(ProviderEvent::Complete { token_usage: None })).await;
return;
}
}
}
}
}
}
Err(e) => {
let _ = tx.send(Err(ProviderError::Streaming(format!("Stream error: {}", e)))).await;
break;
}
}
}
}
Err(e) => {
let _ = tx.send(Err(ProviderError::Http(e))).await;
}
}
});
Ok(rx)
}
fn model(&self) -> &Model {
&self.model
}
fn name(&self) -> &str {
"azure"
}
async fn is_available(&self) -> bool {
let test_messages = vec![Message {
id: "test".to_string(),
session_id: "test".to_string(),
role: MessageRole::User,
content: crate::storage::MessageContent::Text("test".to_string()),
timestamp: chrono::Utc::now(),
metadata: serde_json::Value::Null,
}];
let request = self.create_request(test_messages, false);
self.client
.post(&self.get_api_url())
.header("api-key", &self.api_key)
.header("Content-Type", "application/json")
.json(&request)
.send()
.await
.map(|response| response.status().is_success())
.unwrap_or(false)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::llm::ModelCapabilities;
#[test]
fn test_azure_provider_creation() {
let model = Model {
id: "gpt-4".to_string(),
name: "GPT-4".to_string(),
provider: "azure".to_string(),
context_length: 8192,
max_output_tokens: 4000,
supports_tools: true,
supports_streaming: true,
supports_vision: false,
cost_per_input_token: 0.03,
cost_per_output_token: 0.06,
capabilities: ModelCapabilities::default(),
};
let provider = AzureProvider::new(
"test-key".to_string(),
"https://test.openai.azure.com".to_string(),
"gpt-4-deployment".to_string(),
None,
model,
Some(1000),
Some(0.7),
).unwrap();
assert_eq!(provider.name(), "azure");
assert_eq!(provider.deployment_name, "gpt-4-deployment");
assert_eq!(provider.api_version, "2024-02-15-preview");
}
#[test]
fn test_api_url_generation() {
let model = Model {
id: "gpt-4".to_string(),
name: "GPT-4".to_string(),
provider: "azure".to_string(),
context_length: 8192,
max_output_tokens: 4000,
supports_tools: true,
supports_streaming: true,
supports_vision: false,
cost_per_input_token: 0.03,
cost_per_output_token: 0.06,
capabilities: ModelCapabilities::default(),
};
let provider = AzureProvider::new(
"test-key".to_string(),
"https://test.openai.azure.com/".to_string(), "my-deployment".to_string(),
Some("2023-12-01-preview".to_string()),
model,
None,
None,
).unwrap();
let expected_url = "https://test.openai.azure.com/openai/deployments/my-deployment/chat/completions?api-version=2023-12-01-preview";
assert_eq!(provider.get_api_url(), expected_url);
}
#[test]
fn test_message_conversion() {
let model = Model {
id: "gpt-4".to_string(),
name: "GPT-4".to_string(),
provider: "azure".to_string(),
context_length: 8192,
max_output_tokens: 4000,
supports_tools: true,
supports_streaming: true,
supports_vision: false,
cost_per_input_token: 0.03,
cost_per_output_token: 0.06,
capabilities: ModelCapabilities::default(),
};
let provider = AzureProvider::new(
"test-key".to_string(),
"https://test.openai.azure.com".to_string(),
"gpt-4-deployment".to_string(),
None,
model,
None,
None,
).unwrap();
let messages = vec![
Message {
id: "1".to_string(),
session_id: "test".to_string(),
role: MessageRole::System,
content: crate::storage::MessageContent::Text("You are a helpful assistant.".to_string()),
timestamp: chrono::Utc::now(),
metadata: serde_json::Value::Null,
},
Message {
id: "2".to_string(),
session_id: "test".to_string(),
role: MessageRole::User,
content: crate::storage::MessageContent::Text("Hello!".to_string()),
timestamp: chrono::Utc::now(),
metadata: serde_json::Value::Null,
},
];
let azure_messages = provider.convert_messages(messages);
assert_eq!(azure_messages.len(), 2);
assert_eq!(azure_messages[0].role, "system");
assert_eq!(azure_messages[0].content, "You are a helpful assistant.");
assert_eq!(azure_messages[1].role, "user");
assert_eq!(azure_messages[1].content, "Hello!");
}
}