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 XaiProvider {
client: Client,
api_key: String,
model: Model,
max_tokens: Option<u32>,
temperature: Option<f32>,
}
#[derive(Debug, Serialize)]
struct XaiRequest {
model: String,
messages: Vec<XaiMessage>,
#[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 XaiMessage {
role: String,
content: String,
}
#[derive(Debug, Deserialize)]
struct XaiResponse {
choices: Vec<XaiChoice>,
usage: Option<XaiUsage>,
model: Option<String>,
}
#[derive(Debug, Deserialize)]
struct XaiChoice {
message: XaiResponseMessage,
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct XaiResponseMessage {
role: String,
content: Option<String>,
tool_calls: Option<Vec<serde_json::Value>>,
}
#[derive(Debug, Deserialize)]
struct XaiUsage {
prompt_tokens: u32,
completion_tokens: u32,
total_tokens: u32,
}
#[derive(Debug, Deserialize)]
struct XaiStreamChunk {
choices: Vec<XaiStreamChoice>,
model: Option<String>,
}
#[derive(Debug, Deserialize)]
struct XaiStreamChoice {
delta: XaiDelta,
finish_reason: Option<String>,
}
#[derive(Debug, Deserialize)]
struct XaiDelta {
role: Option<String>,
content: Option<String>,
}
impl XaiProvider {
pub fn new(
api_key: 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)))?;
Ok(Self {
client,
api_key,
model,
max_tokens,
temperature,
})
}
fn get_api_url(&self) -> &str {
"https://api.x.ai/v1/chat/completions"
}
fn convert_messages(&self, messages: Vec<Message>) -> Vec<XaiMessage> {
messages
.into_iter()
.map(|msg| XaiMessage {
role: msg.role.to_string(),
content: msg.content.as_text(),
})
.collect()
}
fn create_request(&self, messages: Vec<Message>, stream: bool) -> XaiRequest {
XaiRequest {
model: self.model.id.clone(),
messages: self.convert_messages(messages),
max_tokens: self.max_tokens,
temperature: self.temperature,
tools: None, tool_choice: None,
stream,
}
}
fn parse_response(&self, response: XaiResponse) -> 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": response.model.unwrap_or_else(|| self.model.id.clone()),
"provider": "xai"
}),
})
}
}
#[async_trait]
impl Provider for XaiProvider {
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("Authorization", format!("Bearer {}", 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!(
"xAI API error {}: {}", status, error_text
)));
}
let xai_response: XaiResponse = response.json().await
.map_err(|e| ProviderError::ResponseParsing(format!("Failed to parse xAI response: {}", e)))?;
self.parse_response(xai_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().to_string();
let api_key = self.api_key.clone();
tokio::spawn(async move {
let response_result = client
.post(&url)
.header("Authorization", format!("Bearer {}", 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!(
"xAI 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::<XaiStreamChunk>(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 {
"xai"
}
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("Authorization", format!("Bearer {}", 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_xai_provider_creation() {
let model = Model {
id: "grok-beta".to_string(),
name: "Grok Beta".to_string(),
provider: "xai".to_string(),
context_length: 131072,
max_output_tokens: 4096,
supports_tools: false,
supports_streaming: true,
supports_vision: false,
cost_per_input_token: 0.005,
cost_per_output_token: 0.015,
capabilities: ModelCapabilities::default(),
};
let provider = XaiProvider::new(
"test-key".to_string(),
model,
Some(1000),
Some(0.7),
).unwrap();
assert_eq!(provider.name(), "xai");
assert_eq!(provider.model.id, "grok-beta");
}
#[test]
fn test_message_conversion() {
let model = Model {
id: "grok-beta".to_string(),
name: "Grok Beta".to_string(),
provider: "xai".to_string(),
context_length: 131072,
max_output_tokens: 4096,
supports_tools: false,
supports_streaming: true,
supports_vision: false,
cost_per_input_token: 0.005,
cost_per_output_token: 0.015,
capabilities: ModelCapabilities::default(),
};
let provider = XaiProvider::new(
"test-key".to_string(),
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 Grok, a witty AI 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("Tell me a joke!".to_string()),
timestamp: chrono::Utc::now(),
metadata: serde_json::Value::Null,
},
];
let xai_messages = provider.convert_messages(messages);
assert_eq!(xai_messages.len(), 2);
assert_eq!(xai_messages[0].role, "system");
assert_eq!(xai_messages[0].content, "You are Grok, a witty AI assistant.");
assert_eq!(xai_messages[1].role, "user");
assert_eq!(xai_messages[1].content, "Tell me a joke!");
}
#[test]
fn test_api_url() {
let model = Model {
id: "grok-beta".to_string(),
name: "Grok Beta".to_string(),
provider: "xai".to_string(),
context_length: 131072,
max_output_tokens: 4096,
supports_tools: false,
supports_streaming: true,
supports_vision: false,
cost_per_input_token: 0.005,
cost_per_output_token: 0.015,
capabilities: ModelCapabilities::default(),
};
let provider = XaiProvider::new(
"test-key".to_string(),
model,
None,
None,
).unwrap();
assert_eq!(provider.get_api_url(), "https://api.x.ai/v1/chat/completions");
}
}