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, ProviderToolCall, Model, ModelCapabilities};
use crate::storage::Message;
use crate::tools::Tool;
#[derive(Debug, Clone)]
pub struct OpenAIProvider {
client: Client,
api_key: String,
base_url: String,
model: Model,
max_tokens: Option<u32>,
temperature: Option<f32>,
}
#[derive(Debug, Serialize)]
struct ChatCompletionRequest {
model: String,
messages: Vec<OpenAIMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
tools: Option<Vec<OpenAITool>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_choice: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
struct OpenAIMessage {
role: String,
content: String,
#[serde(skip_serializing_if = "Option::is_none")]
tool_calls: Option<Vec<OpenAIToolCall>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_call_id: Option<String>,
}
#[derive(Debug, Serialize)]
struct OpenAITool {
#[serde(rename = "type")]
tool_type: String,
function: OpenAIFunction,
}
#[derive(Debug, Serialize)]
struct OpenAIFunction {
name: String,
description: String,
parameters: serde_json::Value,
}
#[derive(Debug, Serialize, Deserialize)]
struct OpenAIToolCall {
id: String,
#[serde(rename = "type")]
call_type: String,
function: OpenAIFunctionCall,
}
#[derive(Debug, Serialize, Deserialize)]
struct OpenAIFunctionCall {
name: String,
arguments: String,
}
#[derive(Debug, Serialize, Deserialize)]
struct ChatCompletionResponse {
id: String,
object: String,
created: u64,
model: String,
choices: Vec<Choice>,
usage: Option<Usage>,
}
#[derive(Debug, Serialize, Deserialize)]
struct Choice {
index: u32,
message: OpenAIMessage,
finish_reason: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Usage {
prompt_tokens: u32,
completion_tokens: u32,
total_tokens: u32,
}
#[derive(Debug, Deserialize)]
struct StreamChunk {
id: String,
object: String,
created: u64,
model: String,
choices: Vec<StreamChoice>,
}
#[derive(Debug, Deserialize)]
struct StreamChoice {
index: u32,
delta: Delta,
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct Delta {
role: Option<String>,
content: Option<String>,
tool_calls: Option<Vec<OpenAIToolCall>>,
}
impl OpenAIProvider {
pub fn new(
api_key: String,
model: Model,
base_url: Option<String>,
max_tokens: Option<u32>,
temperature: Option<f32>,
) -> Result<Self, ProviderError> {
if api_key.is_empty() {
return Err(ProviderError::Configuration("API key is required".into()));
}
let client = Client::builder()
.timeout(std::time::Duration::from_secs(60))
.build()
.map_err(|e| ProviderError::Configuration(format!("Failed to create HTTP client: {}", e)))?;
Ok(Self {
client,
api_key,
base_url: base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string()),
model,
max_tokens,
temperature,
})
}
fn convert_messages(&self, messages: Vec<Message>) -> Vec<OpenAIMessage> {
messages
.into_iter()
.map(|msg| OpenAIMessage {
role: match msg.role {
crate::storage::MessageRole::User => "user".to_string(),
crate::storage::MessageRole::Assistant => "assistant".to_string(),
crate::storage::MessageRole::System => "system".to_string(),
crate::storage::MessageRole::Tool => "tool".to_string(),
},
content: match msg.content {
crate::storage::MessageContent::Text(text) => text,
crate::storage::MessageContent::Structured(parts) => {
parts.iter().map(|part| match part {
crate::storage::ContentPart::Text { text } => text.clone(),
crate::storage::ContentPart::Code { code, language } => {
format!("```{}\n{}\n```", language, code)
}
crate::storage::ContentPart::ToolCall { name, parameters } => {
format!("Tool call: {} with parameters: {}", name, parameters)
}
crate::storage::ContentPart::ToolResult { result, .. } => result.clone(),
_ => "[Unsupported content]".to_string(),
}).collect::<Vec<_>>().join("\n")
}
crate::storage::MessageContent::ToolCall { name, parameters, .. } => {
format!("Tool call: {} with parameters: {}", name, parameters)
}
crate::storage::MessageContent::ToolResult { result, .. } => result,
},
tool_calls: None, tool_call_id: None,
})
.collect()
}
fn convert_tools(&self, tools: Vec<Box<dyn Tool>>) -> Vec<OpenAITool> {
tools
.into_iter()
.map(|tool| OpenAITool {
tool_type: "function".to_string(),
function: OpenAIFunction {
name: tool.name().to_string(),
description: tool.description().to_string(),
parameters: tool.parameter_schema(),
},
})
.collect()
}
fn convert_tool_calls(&self, tool_calls: &[OpenAIToolCall]) -> Vec<ProviderToolCall> {
tool_calls
.iter()
.map(|call| ProviderToolCall {
id: call.id.clone(),
name: call.function.name.clone(),
parameters: serde_json::from_str(&call.function.arguments)
.unwrap_or(serde_json::Value::Null),
})
.collect()
}
fn parse_stream_line(&self, line: &str) -> Option<Result<ProviderEvent, ProviderError>> {
if line.is_empty() || !line.starts_with("data: ") {
return None;
}
let data = &line[6..];
if data == "[DONE]" {
return Some(Ok(ProviderEvent::Complete { token_usage: None }));
}
match serde_json::from_str::<StreamChunk>(data) {
Ok(chunk) => {
if let Some(choice) = chunk.choices.first() {
if let Some(content) = &choice.delta.content {
return Some(Ok(ProviderEvent::ContentChunk {
content: content.clone(),
}));
}
if let Some(tool_calls) = &choice.delta.tool_calls {
if let Some(first_call) = tool_calls.first() {
return Some(Ok(ProviderEvent::ToolCall {
name: first_call.function.name.clone(),
parameters: serde_json::from_str(&first_call.function.arguments)
.unwrap_or(serde_json::Value::Null),
}));
}
}
}
None
}
Err(e) => Some(Err(ProviderError::ResponseParsing(format!(
"Failed to parse stream chunk: {}",
e
)))),
}
}
}
#[async_trait]
impl Provider for OpenAIProvider {
async fn send_messages(
&self,
messages: Vec<Message>,
tools: Vec<Box<dyn Tool>>,
) -> Result<ProviderResponse, ProviderError> {
let request = ChatCompletionRequest {
model: self.model.id.clone(),
messages: self.convert_messages(messages),
max_tokens: self.max_tokens,
temperature: self.temperature,
stream: false,
tools: if tools.is_empty() {
None
} else {
Some(self.convert_tools(tools))
},
tool_choice: None,
};
let response = self
.client
.post(&format!("{}/chat/completions", self.base_url))
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&request)
.send()
.await
.map_err(|e| ProviderError::NetworkError(e.to_string()))?;
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
return Err(ProviderError::ApiError(format!(
"OpenAI API error {}: {}",
status, error_text
)));
}
let completion: ChatCompletionResponse = response
.json()
.await
.map_err(|e| ProviderError::ResponseParsing(e.to_string()))?;
let choice = completion
.choices
.first()
.ok_or_else(|| ProviderError::ResponseParsing("No choices in response".into()))?;
let token_usage = completion.usage.clone().map(|u| TokenUsage {
input_tokens: u.prompt_tokens,
output_tokens: u.completion_tokens,
total_tokens: u.total_tokens,
cache_creation_tokens: 0,
cache_read_tokens: 0,
});
Ok(ProviderResponse {
content: choice.message.content.clone(),
tool_calls: choice
.message
.tool_calls
.as_ref()
.map(|calls| self.convert_tool_calls(calls))
.unwrap_or_default(),
token_usage,
metadata: serde_json::to_value(&completion).unwrap_or_default(),
})
}
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 = ChatCompletionRequest {
model: self.model.id.clone(),
messages: self.convert_messages(messages),
max_tokens: self.max_tokens,
temperature: self.temperature,
stream: true,
tools: if tools.is_empty() {
None
} else {
Some(self.convert_tools(tools))
},
tool_choice: None,
};
let client = self.client.clone();
let base_url = self.base_url.clone();
let api_key = self.api_key.clone();
tokio::spawn(async move {
let response = match client
.post(&format!("{}/chat/completions", base_url))
.header("Authorization", format!("Bearer {}", api_key))
.header("Content-Type", "application/json")
.json(&request)
.send()
.await
{
Ok(response) => response,
Err(e) => {
let _ = tx.send(Err(ProviderError::NetworkError(e.to_string()))).await;
return;
}
};
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!(
"OpenAI API error {}: {}",
status, error_text
))))
.await;
return;
}
let mut stream = response.bytes_stream();
let mut buffer = String::new();
while let Some(chunk) = stream.next().await {
match chunk {
Ok(bytes) => {
buffer.push_str(&String::from_utf8_lossy(&bytes));
while let Some(line_end) = buffer.find('\n') {
let line = buffer[..line_end].trim().to_string();
buffer = buffer[line_end + 1..].to_string();
if let Some(event) = OpenAIProvider::parse_stream_line(&OpenAIProvider {
client: client.clone(),
api_key: api_key.clone(),
base_url: base_url.clone(),
model: Model {
id: "temp".to_string(),
name: "temp".to_string(),
provider: "openai".to_string(),
context_length: 0,
max_output_tokens: 0,
supports_tools: false,
supports_streaming: false,
supports_vision: false,
cost_per_input_token: 0.0,
cost_per_output_token: 0.0,
capabilities: ModelCapabilities::default(),
},
max_tokens: None,
temperature: None,
}, &line) {
if let Err(_) = tx.send(event).await {
break;
}
}
}
}
Err(e) => {
let _ = tx.send(Err(ProviderError::NetworkError(e.to_string()))).await;
break;
}
}
}
});
Ok(rx)
}
fn model(&self) -> &Model {
&self.model
}
fn name(&self) -> &str {
"openai"
}
async fn is_available(&self) -> bool {
let test_request = ChatCompletionRequest {
model: "gpt-3.5-turbo".to_string(),
messages: vec![OpenAIMessage {
role: "user".to_string(),
content: "test".to_string(),
tool_calls: None,
tool_call_id: None,
}],
max_tokens: Some(1),
temperature: Some(0.0),
stream: false,
tools: None,
tool_choice: None,
};
self.client
.post(&format!("{}/chat/completions", self.base_url))
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&test_request)
.send()
.await
.map(|response| response.status().is_success())
.unwrap_or(false)
}
}