use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use thiserror::Error;
use tracing::{debug, info, warn};
use uuid::Uuid;
use mr_common::LlmConfig;
use super::types::{ChatCompletionRequest, ChatCompletionResponse, LlmMessage};
#[derive(Debug, Error)]
pub enum LlmError {
#[error("LLM HTTP request failed: {0}")]
Http(String),
#[error("LLM HTTP status {status}: {body}")]
HttpStatus { status: u16, body: String },
#[error("LLM response parse failed: {0}")]
Parse(String),
#[error("LLM response is empty")]
EmptyResponse,
#[error("LLM request cancelled")]
Cancelled,
}
#[async_trait]
pub trait LlmClient: Send + Sync {
async fn chat(&self, messages: &[LlmMessage]) -> Result<String, LlmError>;
fn configured(&self) -> bool;
}
pub struct OpenAiCompatibleClient {
config: LlmConfig,
http: reqwest::Client,
}
impl OpenAiCompatibleClient {
pub fn new(config: LlmConfig) -> Self {
let timeout = Duration::from_secs(config.timeout_secs.max(1));
let http = reqwest::Client::builder()
.timeout(timeout)
.connect_timeout(Duration::from_secs(15))
.build()
.expect("failed to build HTTP client");
Self { config, http }
}
fn endpoint(&self) -> String {
format!("{}/chat/completions", self.config.url.trim_end_matches('/'))
}
fn build_request(&self, messages: &[LlmMessage]) -> ChatCompletionRequest {
let mut extra = serde_json::Map::new();
if let Some(serde_json::Value::Object(map)) = self
.config
.think_level()
.to_provider_param(&self.config.provider)
{
extra.extend(map);
}
ChatCompletionRequest {
model: self.config.model.clone(),
messages: messages.to_vec(),
max_tokens: Some(self.config.max_tokens),
temperature: Some(self.config.temperature),
extra: serde_json::Value::Object(extra),
}
}
}
#[async_trait]
impl LlmClient for OpenAiCompatibleClient {
async fn chat(&self, messages: &[LlmMessage]) -> Result<String, LlmError> {
if !self.configured() {
return Err(LlmError::Http(
"LLM api_key is empty, check [llm] config".to_string(),
));
}
let request = self.build_request(messages);
let request_id = Uuid::new_v4();
debug!(
request_id = %request_id,
provider = %self.config.provider,
model = %self.config.model,
"LLM chat request"
);
let response = self
.http
.post(self.endpoint())
.bearer_auth(&self.config.api_key)
.json(&request)
.send()
.await
.map_err(|e| LlmError::Http(e.to_string()))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|e| LlmError::Http(e.to_string()))?;
if !status.is_success() {
return Err(LlmError::HttpStatus {
status: status.as_u16(),
body: body.chars().take(500).collect(),
});
}
let parsed: ChatCompletionResponse =
serde_json::from_str(&body).map_err(|e| LlmError::Parse(e.to_string()))?;
let text = parsed
.choices
.first()
.map(|c| c.message.text())
.unwrap_or_default();
if let Some(usage) = parsed.usage {
info!(
request_id = %request_id,
prompt_tokens = usage.prompt_tokens,
completion_tokens = usage.completion_tokens,
"LLM chat completed"
);
}
if text.is_empty() {
warn!(request_id = %request_id, "LLM response has no content");
return Err(LlmError::EmptyResponse);
}
Ok(text)
}
fn configured(&self) -> bool {
!self.config.api_key.trim().is_empty()
}
}
#[derive(Debug, Clone)]
pub struct MockLlmClient {
response: String,
configured: bool,
call_count: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}
impl MockLlmClient {
pub fn new(response: impl Into<String>) -> Self {
Self {
response: response.into(),
configured: true,
call_count: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
pub fn unconfigured() -> Self {
Self {
response: String::new(),
configured: false,
call_count: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
pub fn call_count(&self) -> usize {
self.call_count.load(std::sync::atomic::Ordering::Relaxed)
}
}
#[async_trait]
impl LlmClient for MockLlmClient {
async fn chat(&self, _messages: &[LlmMessage]) -> Result<String, LlmError> {
self.call_count
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
if !self.configured {
return Err(LlmError::Http("not configured".to_string()));
}
Ok(self.response.clone())
}
fn configured(&self) -> bool {
self.configured
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
#[test]
fn test_endpoint_joins_url() {
let client = OpenAiCompatibleClient::new(LlmConfig::default());
assert_eq!(
client.endpoint(),
"https://api.deepseek.com/v1/chat/completions"
);
}
#[test]
fn test_build_request_deepseek_thinking() {
let mut config = LlmConfig::default();
config.provider = "deepseek".to_string();
config.think_level = "xhigh".to_string();
let client = OpenAiCompatibleClient::new(config);
let req = client.build_request(&[LlmMessage::user("hello")]);
assert_eq!(req.extra["thinking"]["effort"], "max");
assert_eq!(req.model, "deepseek-chat");
assert_eq!(req.max_tokens, Some(4096));
}
#[test]
fn test_build_request_openai_reasoning() {
let mut config = LlmConfig::default();
config.provider = "openai".to_string();
config.model = "gpt-4o-mini".to_string();
config.think_level = "medium".to_string();
let client = OpenAiCompatibleClient::new(config);
let req = client.build_request(&[]);
assert_eq!(req.extra["reasoning_effort"], "medium");
}
#[test]
fn test_build_request_off_no_thinking() {
let mut config = LlmConfig::default();
config.think_level = "off".to_string();
let client = OpenAiCompatibleClient::new(config);
let req = client.build_request(&[]);
assert!(req.extra.as_object().unwrap().is_empty());
}
#[test]
fn test_configured_check() {
let client = OpenAiCompatibleClient::new(LlmConfig::default());
assert!(!client.configured());
let mut config = LlmConfig::default();
config.api_key = "sk-test".to_string();
let client = OpenAiCompatibleClient::new(config);
assert!(client.configured());
}
#[test]
fn test_arc_client_trait_object() {
let client: Arc<dyn LlmClient> = Arc::new(MockLlmClient::new("ok"));
assert!(client.configured());
}
}