563bb99
use std::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 response could not be parsed: {0}")]
Decode(String),
#[error("invalid ai task: {0}")]
InvalidTask(String),
/// 任务本身合法，但执行期间依赖（git / 数据库 / 仓库快照）不可用。
///
/// 与 [`Self::InvalidTask`] 区分开：这类失败应当重试，而不是直接进 DLQ。
#[error("ai task execution failed: {0}")]
Execution(String),
}
pub type AgentResult<T> = Result<T, AgentError>;
/// 连接阶段单独设限：端点可能解析出多个地址，坏地址不该吃掉整个请求预算。
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
/// reqwest 的 Display 只剩 "error sending request for url(...)"，真正的原因
/// （超时 / DNS / connect）在 source 链里，这里拼出来，否则线上只剩超时现象。
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::Execution(_)
)
}
}
#[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 {
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,
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()))?;
Ok(Self {
http,
endpoint: cfg.endpoint(),
api_key: cfg.api_key.trim().to_string(),
model: cfg.model.trim().to_string(),
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 body = ChatRequest {
model: self.model.clone(),
messages: vec![
ChatMessage {
role: "system".to_string(),
content: system_prompt.to_string(),
},
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).await {
Ok(completion) => return Ok(completion),
Err(error) => {
let retryable = error.is_retryable();
tracing::warn!(%error, 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())))
}
async fn send_once(&self, body: &ChatRequest) -> 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 response = request
.send()
.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() {
return Err(AgentError::Status {
status: status.as_u16(),
body: truncate(&text, 2048),
});
}
let parsed: ChatResponse = serde_json::from_str(&text)
.map_err(|e| AgentError::Decode(format!("{e}: {}", truncate(&text, 512))))?;
let content = parsed
.choices
.first()
.and_then(|choice| choice.message.as_ref())
.and_then(|message| message.content.clone())
.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,
})
}
}
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())
}