112d144
use std::{borrow::Cow, time::Duration};
use config::ai::AiConfig;
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum AgentError {
#[error("ai is not configured: {0}")]
NotConfigured(String),
#[error("ai request failed: {0}")]
Transport(String),
#[error("ai returned status {status}: {body}")]
Status { status: u16, body: String },
#[error("ai rejected the request with status {status}: {body}")]
Rejected { status: u16, body: String },
#[error("ai response could not be parsed: {0}")]
Decode(String),
#[error("ai returned an incomplete response: {0}")]
Incomplete(String),
#[error("invalid ai task: {0}")]
InvalidTask(String),
#[error("ai task execution failed: {0}")]
Execution(String),
}
pub type AgentResult<T> = Result<T, AgentError>;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
fn transport_message(error: &reqwest::Error) -> String {
let kind = if error.is_timeout() {
"timed out"
} else if error.is_connect() {
"connect failed"
} else {
"request failed"
};
let mut message = format!("{kind}: {error}");
let mut cause = std::error::Error::source(error);
while let Some(next) = cause {
message = format!("{message}: {next}");
cause = next.source();
}
message
}
impl AgentError {
#[must_use]
pub fn is_retryable(&self) -> bool {
matches!(
self,
Self::Transport(_) | Self::Status { .. } | Self::Incomplete(_) | Self::Execution(_)
)
}
}
const DEEPSEEK_JSON_INSTRUCTION: &str = "\n\nJSON 输出要求：严格按照上方 JSON 示例返回一个完整的 JSON 对象。只输出 JSON，不要输出 Markdown 代码块、思考过程、解释、前缀或后缀。";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Provider {
DeepSeek,
Compatible,
}
impl Provider {
fn detect(endpoint: &str, model: &str) -> Self {
let deepseek_endpoint = reqwest::Url::parse(endpoint)
.ok()
.and_then(|url| url.host_str().map(str::to_owned))
.is_some_and(|host| {
host.eq_ignore_ascii_case("deepseek.com")
|| host.to_ascii_lowercase().ends_with(".deepseek.com")
});
if deepseek_endpoint || model.trim().to_ascii_lowercase().starts_with("deepseek") {
Self::DeepSeek
} else {
Self::Compatible
}
}
fn name(self) -> &'static str {
match self {
Self::DeepSeek => "deepseek",
Self::Compatible => "openai-compatible",
}
}
}
#[derive(Debug, Clone, Serialize)]
struct ChatRequest {
model: String,
messages: Vec<ChatMessage>,
temperature: f32,
max_tokens: u32,
#[serde(skip_serializing_if = "Option::is_none")]
response_format: Option<ResponseFormat>,
#[serde(skip_serializing_if = "Option::is_none")]
thinking_budget: Option<u32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChatMessage {
role: String,
content: String,
}
#[derive(Debug, Clone, Serialize)]
struct ResponseFormat {
#[serde(rename = "type")]
kind: String,
}
#[derive(Debug, Clone, Deserialize)]
struct ChatResponse {
#[serde(default)]
choices: Vec<ChatChoice>,
#[serde(default)]
usage: Option<ChatUsage>,
}
#[derive(Debug, Clone, Deserialize)]
struct ChatChoice {
#[serde(default)]
finish_reason: Option<String>,
message: Option<ChatChoiceMessage>,
}
#[derive(Debug, Clone, Deserialize)]
struct ChatChoiceMessage {
#[serde(default)]
content: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
struct ChatUsage {
#[serde(default)]
prompt_tokens: Option<i32>,
#[serde(default)]
completion_tokens: Option<i32>,
}
/// 一次模型调用的完整产物：正文、用量与原始响应体。
#[derive(Debug, Clone)]
pub struct AgentCompletion {
pub content: String,
pub prompt_tokens: Option<i32>,
pub completion_tokens: Option<i32>,
pub raw: serde_json::Value,
}
/// OpenAI 兼容的单一模型客户端。
///
/// 平台内所有 AI 能力都经由本客户端访问 `[ai].model`，不存在按功能切换模型的入口。
#[derive(Clone)]
pub struct AgentClient {
http: reqwest::Client,
endpoint: String,
provider: Provider,
api_key: String,
model: String,
temperature: f32,
max_output_tokens: u32,
max_attempts: i32,
thinking_budget: Option<u32>,
}
impl AgentClient {
pub fn new(cfg: &AiConfig) -> AgentResult<Self> {
if cfg.base_url.trim().is_empty() || cfg.model.trim().is_empty() {
return Err(AgentError::NotConfigured(
"base_url and model are required".to_string(),
));
}
let http = reqwest::Client::builder()
.timeout(Duration::from_secs(cfg.request_timeout_secs.max(1)))
.connect_timeout(CONNECT_TIMEOUT)
.build()
.map_err(|e| AgentError::Transport(e.to_string()))?;
let endpoint = cfg.endpoint();
let model = cfg.model.trim().to_string();
Ok(Self {
http,
provider: Provider::detect(&endpoint, &model),
endpoint,
api_key: cfg.api_key.trim().to_string(),
model,
temperature: cfg.temperature,
max_output_tokens: cfg.max_output_tokens,
max_attempts: cfg.max_attempts.max(1),
thinking_budget: (cfg.thinking_budget > 0).then_some(cfg.thinking_budget),
})
}
pub fn model(&self) -> &str {
&self.model
}
/// 发送一次 chat completion 请求，失败时按指数退避重试。
pub async fn complete(
&self,
system_prompt: &str,
user_prompt: &str,
json_mode: bool,
) -> AgentResult<AgentCompletion> {
let system_prompt = self.system_prompt(system_prompt, json_mode);
let body = ChatRequest {
model: self.model.clone(),
messages: vec![
ChatMessage {
role: "system".to_string(),
content: system_prompt.into_owned(),
},
ChatMessage {
role: "user".to_string(),
content: user_prompt.to_string(),
},
],
temperature: self.temperature,
max_tokens: self.max_output_tokens,
response_format: json_mode.then(|| ResponseFormat {
kind: "json_object".to_string(),
}),
thinking_budget: self.thinking_budget,
};
let mut last_error: Option<AgentError> = None;
for attempt in 0..self.max_attempts {
if attempt > 0 {
let backoff = Duration::from_millis(500 * (1u64 << attempt.min(5)));
tokio::time::sleep(backoff).await;
}
match self.send_once(&body, json_mode).await {
Ok(completion) => return Ok(completion),
Err(error) => {
let retryable = self.should_retry(&error);
tracing::warn!(
%error,
provider = self.provider.name(),
model = %self.model,
attempt = attempt + 1,
"ai completion attempt failed"
);
last_error = Some(error);
if !retryable {
break;
}
}
}
}
Err(last_error.unwrap_or_else(|| AgentError::Transport("no attempt made".to_string())))
}
fn system_prompt<'a>(&self, prompt: &'a str, json_mode: bool) -> Cow<'a, str> {
if self.provider == Provider::DeepSeek && json_mode {
Cow::Owned(format!("{prompt}{DEEPSEEK_JSON_INSTRUCTION}"))
} else {
Cow::Borrowed(prompt)
}
}
fn should_retry(&self, error: &AgentError) -> bool {
if self.provider != Provider::DeepSeek {
return error.is_retryable();
}
match error {
AgentError::Status { status, .. } => deepseek_retryable_status(*status),
_ => error.is_retryable(),
}
}
async fn send_once(&self, body: &ChatRequest, json_mode: bool) -> AgentResult<AgentCompletion> {
let mut request = self
.http
.post(&self.endpoint)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.json(body);
if !self.api_key.is_empty() {
request = request.bearer_auth(&self.api_key);
}
let span = export::telemetry::http_client_span("POST", &self.endpoint);
let response = export::telemetry::send(span, request)
.await
.map_err(|e| AgentError::Transport(transport_message(&e)))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|e| AgentError::Transport(transport_message(&e)))?;
if !status.is_success() {
let status = status.as_u16();
let body = truncate(&text, 2048);
if self.provider == Provider::DeepSeek && !deepseek_retryable_status(status) {
return Err(AgentError::Rejected { status, body });
}
return Err(AgentError::Status { status, body });
}
let parsed: ChatResponse = serde_json::from_str(&text)
.map_err(|e| AgentError::Decode(format!("{e}: {}", truncate(&text, 512))))?;
let choice = parsed.choices.first().ok_or_else(|| {
if self.provider == Provider::DeepSeek && json_mode {
AgentError::Incomplete("deepseek JSON mode returned no choices".to_string())
} else {
AgentError::Decode("response has no choices".to_string())
}
})?;
let content = choice
.message
.as_ref()
.and_then(|message| message.content.as_deref());
if self.provider == Provider::DeepSeek && json_mode {
validate_deepseek_json(choice.finish_reason.as_deref(), content)?;
}
let content = content
.map(str::to_owned)
.ok_or_else(|| AgentError::Decode("response has no message content".to_string()))?;
let raw: serde_json::Value = serde_json::from_str(&text).unwrap_or(serde_json::Value::Null);
Ok(AgentCompletion {
content,
prompt_tokens: parsed.usage.as_ref().and_then(|u| u.prompt_tokens),
completion_tokens: parsed.usage.as_ref().and_then(|u| u.completion_tokens),
raw,
})
}
}
fn deepseek_retryable_status(status: u16) -> bool {
matches!(status, 408 | 429) || (500..=599).contains(&status)
}
fn validate_deepseek_json(finish_reason: Option<&str>, content: Option<&str>) -> AgentResult<()> {
match finish_reason {
Some("length") => {
return Err(AgentError::Decode(
"deepseek JSON output was truncated; increase ai.max_output_tokens or reduce the input"
.to_string(),
));
}
Some("insufficient_system_resource" | "aborted") => {
return Err(AgentError::Incomplete(format!(
"deepseek stopped with finish_reason={}",
finish_reason.unwrap_or_default()
)));
}
Some("content_filter" | "tool_calls") => {
return Err(AgentError::Decode(format!(
"deepseek JSON output stopped with finish_reason={}",
finish_reason.unwrap_or_default()
)));
}
_ => {}
}
let content = content
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| {
AgentError::Incomplete("deepseek JSON mode returned empty content".to_string())
})?;
let value: serde_json::Value = serde_json::from_str(content).map_err(|error| {
AgentError::Incomplete(format!("deepseek JSON mode returned invalid JSON: {error}"))
})?;
if !value.is_object() {
return Err(AgentError::Incomplete(
"deepseek JSON mode returned a non-object value".to_string(),
));
}
Ok(())
}
pub(crate) fn truncate(value: &str, max: usize) -> String {
if value.chars().count() <= max {
return value.to_string();
}
let mut out: String = value.chars().take(max).collect();
out.push_str("…(truncated)");
out
}
/// 从模型输出中提取 JSON 对象；容忍围栏代码块与前后说明文字。
pub fn extract_json(content: &str) -> AgentResult<serde_json::Value> {
let direct = content.trim();
if let Ok(value) = serde_json::from_str::<serde_json::Value>(direct) {
return Ok(value);
}
if let Some(inner) = fenced_json(direct)
&& let Ok(value) = serde_json::from_str::<serde_json::Value>(&inner)
{
return Ok(value);
}
if let (Some(start), Some(end)) = (direct.find('{'), direct.rfind('}'))
&& start < end
&& let Ok(value) = serde_json::from_str::<serde_json::Value>(&direct[start..=end])
{
return Ok(value);
}
Err(AgentError::Decode(
"model output did not contain JSON".to_string(),
))
}
fn fenced_json(value: &str) -> Option<String> {
let start = value.find("```")?;
let rest = &value[start + 3..];
let rest = rest.strip_prefix("json").unwrap_or(rest);
let end = rest.find("```")?;
Some(rest[..end].trim().to_string())
}