zhenyi

Fork 0

Clone

feat(agent): enhance DeepSeek integration with improved error handling and JSON validation

ZhenYi committed 1 month ago

parent 81f2cbd commit 3d404c6 Browse files

lib/agent/src/client.rs +141 −28 View

@@ -1,4 +1,4 @@

1use std::time::Duration;

1use std::{borrow::Cow, time::Duration};

22

33use config::ai::AiConfig;

44use serde::{Deserialize, Serialize};

@@ -12,24 +12,22 @@

1212 Transport(String),

1313 #[error("ai returned status {status}: {body}")]

1414 Status { status: u16, body: String },

15 #[error("ai rejected the request with status {status}: {body}")]

16 Rejected { status: u16, body: String },

1517 #[error("ai response could not be parsed: {0}")]

1618 Decode(String),

19 #[error("ai returned an incomplete response: {0}")]

20 Incomplete(String),

1721 #[error("invalid ai task: {0}")]

1822 InvalidTask(String),

19 /// 任务本身合法，但执行期间依赖（git / 数据库 / 仓库快照）不可用。

20 ///

21 /// 与 [`Self::InvalidTask`] 区分开：这类失败应当重试，而不是直接进 DLQ。

2223 #[error("ai task execution failed: {0}")]

2324 Execution(String),

2425}

2526

2627pub type AgentResult<T> = Result<T, AgentError>;

2728

28/// 连接阶段单独设限：端点可能解析出多个地址，坏地址不该吃掉整个请求预算。

2929const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);

3030

31/// reqwest 的 Display 只剩 "error sending request for url(...)"，真正的原因

32/// （超时 / DNS / connect）在 source 链里，这里拼出来，否则线上只剩超时现象。

3331fn transport_message(error: &reqwest::Error) -> String {

3432 let kind = if error.is_timeout() {

3533 "timed out"

@@ -48,19 +46,47 @@

4846}

4947

5048impl AgentError {

51 /// 该错误是否值得重投。

52 ///

53 /// 任务描述本身非法（例如目标仓库已被删除）重试多少次结果都一样，应直接失败；

54 /// 端点抖动与依赖暂不可用则应重试。

5549 #[must_use]

5650 pub fn is_retryable(&self) -> bool {

5751 matches!(

5852 self,

59 Self::Transport(_) | Self::Status { .. } | Self::Execution(_)

53 Self::Transport(_) | Self::Status { .. } | Self::Incomplete(_) | Self::Execution(_)

6054 )

6155 }

6256}

6357

58const DEEPSEEK_JSON_INSTRUCTION: &str = "\n\nJSON 输出要求：严格按照上方 JSON 示例返回一个完整的 JSON 对象。只输出 JSON，不要输出 Markdown 代码块、思考过程、解释、前缀或后缀。";

59

60#[derive(Debug, Clone, Copy, PartialEq, Eq)]

61enum Provider {

62 DeepSeek,

63 Compatible,

64}

65

66impl Provider {

67 fn detect(endpoint: &str, model: &str) -> Self {

68 let deepseek_endpoint = reqwest::Url::parse(endpoint)

69 .ok()

70 .and_then(|url| url.host_str().map(str::to_owned))

71 .is_some_and(|host| {

72 host.eq_ignore_ascii_case("deepseek.com")

73 || host.to_ascii_lowercase().ends_with(".deepseek.com")

74 });

75 if deepseek_endpoint || model.trim().to_ascii_lowercase().starts_with("deepseek") {

76 Self::DeepSeek

77 } else {

78 Self::Compatible

79 }

80 }

81

82 fn name(self) -> &'static str {

83 match self {

84 Self::DeepSeek => "deepseek",

85 Self::Compatible => "openai-compatible",

86 }

87 }

88}

89

6490#[derive(Debug, Clone, Serialize)]

6591struct ChatRequest {

6692 model: String,

@@ -95,6 +121,8 @@

95121

96122#[derive(Debug, Clone, Deserialize)]

97123struct ChatChoice {

124 #[serde(default)]

125 finish_reason: Option<String>,

98126 message: Option<ChatChoiceMessage>,

99127}

100128

@@ -128,6 +156,7 @@

128156pub struct AgentClient {

129157 http: reqwest::Client,

130158 endpoint: String,

159 provider: Provider,

131160 api_key: String,

132161 model: String,

133162 temperature: f32,

@@ -148,11 +177,14 @@

148177 .connect_timeout(CONNECT_TIMEOUT)

149178 .build()

150179 .map_err(|e| AgentError::Transport(e.to_string()))?;

180 let endpoint = cfg.endpoint();

181 let model = cfg.model.trim().to_string();

151182 Ok(Self {

152183 http,

153 endpoint: cfg.endpoint(),

184 provider: Provider::detect(&endpoint, &model),

185 endpoint,

154186 api_key: cfg.api_key.trim().to_string(),

155 model: cfg.model.trim().to_string(),

187 model,

156188 temperature: cfg.temperature,

157189 max_output_tokens: cfg.max_output_tokens,

158190 max_attempts: cfg.max_attempts.max(1),

@@ -171,12 +203,13 @@

171203 user_prompt: &str,

172204 json_mode: bool,

173205 ) -> AgentResult<AgentCompletion> {

206 let system_prompt = self.system_prompt(system_prompt, json_mode);

174207 let body = ChatRequest {

175208 model: self.model.clone(),

176209 messages: vec![

177210 ChatMessage {

178211 role: "system".to_string(),

179 content: system_prompt.to_string(),

212 content: system_prompt.into_owned(),

180213 },

181214 ChatMessage {

182215 role: "user".to_string(),

@@ -197,11 +230,17 @@

197230 let backoff = Duration::from_millis(500 * (1u64 << attempt.min(5)));

198231 tokio::time::sleep(backoff).await;

199232 }

200 match self.send_once(&body).await {

233 match self.send_once(&body, json_mode).await {

201234 Ok(completion) => return Ok(completion),

202235 Err(error) => {

203 let retryable = error.is_retryable();

204 tracing::warn!(%error, attempt = attempt + 1, "ai completion attempt failed");

236 let retryable = self.should_retry(&error);

237 tracing::warn!(

238 %error,

239 provider = self.provider.name(),

240 model = %self.model,

241 attempt = attempt + 1,

242 "ai completion attempt failed"

243 );

205244 last_error = Some(error);

206245 if !retryable {

207246 break;

@@ -212,7 +251,25 @@

212251 Err(last_error.unwrap_or_else(|| AgentError::Transport("no attempt made".to_string())))

213252 }

214253

215 async fn send_once(&self, body: &ChatRequest) -> AgentResult<AgentCompletion> {

254 fn system_prompt<'a>(&self, prompt: &'a str, json_mode: bool) -> Cow<'a, str> {

255 if self.provider == Provider::DeepSeek && json_mode {

256 Cow::Owned(format!("{prompt}{DEEPSEEK_JSON_INSTRUCTION}"))

257 } else {

258 Cow::Borrowed(prompt)

259 }

260 }

261

262 fn should_retry(&self, error: &AgentError) -> bool {

263 if self.provider != Provider::DeepSeek {

264 return error.is_retryable();

265 }

266 match error {

267 AgentError::Status { status, .. } => deepseek_retryable_status(*status),

268 _ => error.is_retryable(),

269 }

270 }

271

272 async fn send_once(&self, body: &ChatRequest, json_mode: bool) -> AgentResult<AgentCompletion> {

216273 let mut request = self

217274 .http

218275 .post(&self.endpoint)

@@ -231,18 +288,31 @@

231288 .await

232289 .map_err(|e| AgentError::Transport(transport_message(&e)))?;

233290 if !status.is_success() {

234 return Err(AgentError::Status {

235 status: status.as_u16(),

236 body: truncate(&text, 2048),

237 });

291 let status = status.as_u16();

292 let body = truncate(&text, 2048);

293 if self.provider == Provider::DeepSeek && !deepseek_retryable_status(status) {

294 return Err(AgentError::Rejected { status, body });

295 }

296 return Err(AgentError::Status { status, body });

238297 }

239298 let parsed: ChatResponse = serde_json::from_str(&text)

240299 .map_err(|e| AgentError::Decode(format!("{e}: {}", truncate(&text, 512))))?;

241 let content = parsed

242 .choices

243 .first()

244 .and_then(|choice| choice.message.as_ref())

245 .and_then(|message| message.content.clone())

300 let choice = parsed.choices.first().ok_or_else(|| {

301 if self.provider == Provider::DeepSeek && json_mode {

302 AgentError::Incomplete("deepseek JSON mode returned no choices".to_string())

303 } else {

304 AgentError::Decode("response has no choices".to_string())

305 }

306 })?;

307 let content = choice

308 .message

309 .as_ref()

310 .and_then(|message| message.content.as_deref());

311 if self.provider == Provider::DeepSeek && json_mode {

312 validate_deepseek_json(choice.finish_reason.as_deref(), content)?;

313 }

314 let content = content

315 .map(str::to_owned)

246316 .ok_or_else(|| AgentError::Decode("response has no message content".to_string()))?;

247317 let raw: serde_json::Value = serde_json::from_str(&text).unwrap_or(serde_json::Value::Null);

248318 Ok(AgentCompletion {

@@ -254,6 +324,49 @@

254324 }

255325}

256326

327fn deepseek_retryable_status(status: u16) -> bool {

328 matches!(status, 408 | 429) || (500..=599).contains(&status)

329}

330

331fn validate_deepseek_json(finish_reason: Option<&str>, content: Option<&str>) -> AgentResult<()> {

332 match finish_reason {

333 Some("length") => {

334 return Err(AgentError::Decode(

335 "deepseek JSON output was truncated; increase ai.max_output_tokens or reduce the input"

336 .to_string(),

337 ));

338 }

339 Some("insufficient_system_resource" | "aborted") => {

340 return Err(AgentError::Incomplete(format!(

341 "deepseek stopped with finish_reason={}",

342 finish_reason.unwrap_or_default()

343 )));

344 }

345 Some("content_filter" | "tool_calls") => {

346 return Err(AgentError::Decode(format!(

347 "deepseek JSON output stopped with finish_reason={}",

348 finish_reason.unwrap_or_default()

349 )));

350 }

351 _ => {}

352 }

353 let content = content

354 .map(str::trim)

355 .filter(|value| !value.is_empty())

356 .ok_or_else(|| {

357 AgentError::Incomplete("deepseek JSON mode returned empty content".to_string())

358 })?;

359 let value: serde_json::Value = serde_json::from_str(content).map_err(|error| {

360 AgentError::Incomplete(format!("deepseek JSON mode returned invalid JSON: {error}"))

361 })?;

362 if !value.is_object() {

363 return Err(AgentError::Incomplete(

364 "deepseek JSON mode returned a non-object value".to_string(),

365 ));

366 }

367 Ok(())

368}

369

257370pub(crate) fn truncate(value: &str, max: usize) -> String {

258371 if value.chars().count() <= max {

259372 return value.to_string();