agent_base/llm/
openai.rs

1use async_trait::async_trait;
2use eventsource_stream::Eventsource;
3use futures_core::Stream;
4use futures_util::StreamExt;
5use reqwest::Client;
6use serde_json::{Value, json};
7use std::pin::Pin;
8use std::time::Duration;
9
10use super::{LlmCapabilities, LlmClient, ReasoningConfig, ReasoningEffort, StreamChunk, UsageInfo};
11use crate::types::{
12    AgentError, AgentResult, ChatMessage, ImageAttachment, ImageDetail, ResponseFormat,
13    ToolCallMessage,
14};
15
16#[derive(Clone, Debug)]
17pub struct LlmClientConfig {
18    pub connect_timeout: Duration,
19    pub request_timeout: Duration,
20    pub pool_max_idle_per_host: usize,
21    pub pool_idle_timeout: Duration,
22}
23
24impl Default for LlmClientConfig {
25    fn default() -> Self {
26        Self {
27            connect_timeout: Duration::from_secs(15),
28            request_timeout: Duration::from_secs(120),
29            pool_max_idle_per_host: 10,
30            pool_idle_timeout: Duration::from_secs(90),
31        }
32    }
33}
34
35pub struct OpenAiClient {
36    api_key: String,
37    model: String,
38    base_url: String,
39    client: Client,
40}
41
42impl OpenAiClient {
43    pub fn new(api_key: String, model: String, base_url: Option<String>) -> Self {
44        Self::new_with_config(api_key, model, base_url, LlmClientConfig::default())
45    }
46
47    pub fn new_with_config(
48        api_key: String,
49        model: String,
50        base_url: Option<String>,
51        config: LlmClientConfig,
52    ) -> Self {
53        let client = Client::builder()
54            .connect_timeout(config.connect_timeout)
55            .timeout(config.request_timeout)
56            .pool_max_idle_per_host(config.pool_max_idle_per_host)
57            .pool_idle_timeout(config.pool_idle_timeout)
58            .build()
59            .unwrap_or_else(|e| {
60                tracing::warn!(error = %e, "Failed to build reqwest client with custom config, falling back to default");
61                Client::new()
62            });
63        Self {
64            api_key,
65            model,
66            base_url: base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string()),
67            client,
68        }
69    }
70
71    /// 使用不同模型的变体。共享底层 HTTP 连接池，零额外成本。
72    ///
73    /// 类似 Claude Code 的 opus/sonnet/haiku — 同一 client，不同 model。
74    pub fn with_model(&self, model: impl Into<String>) -> Self {
75        Self {
76            api_key: self.api_key.clone(),
77            model: model.into(),
78            base_url: self.base_url.clone(),
79            client: self.client.clone(), // reqwest::Client 内部是 Arc
80        }
81    }
82
83    fn is_qwen_model(&self) -> bool {
84        self.model.starts_with("qwen")
85    }
86
87    fn is_deepseek_model(&self) -> bool {
88        self.model.starts_with("deepseek")
89    }
90
91    fn apply_reasoning_config(
92        &self,
93        request_body: &mut Value,
94        reasoning: Option<&ReasoningConfig>,
95    ) {
96        let Some(config) = reasoning else { return };
97
98        if self.is_qwen_model() {
99            // qwen 模型使用 enable_thinking 和 thinking_budget
100            // 对于 OpenAI 兼容接口，直接放在请求体顶层
101            if let Some(enabled) = config.enabled {
102                if let Some(obj) = request_body.as_object_mut() {
103                    obj.insert("enable_thinking".to_string(), json!(enabled));
104                }
105            }
106            if let Some(budget) = config.budget_tokens {
107                if let Some(obj) = request_body.as_object_mut() {
108                    obj.insert("thinking_budget".to_string(), json!(budget));
109                }
110            }
111            // 将 effort 转换为 thinking_budget
112            if let Some(effort) = &config.effort {
113                let budget = match effort {
114                    ReasoningEffort::None => 0,
115                    ReasoningEffort::Low => 500,
116                    ReasoningEffort::Medium => 2000,
117                    ReasoningEffort::High => 5000,
118                    ReasoningEffort::XHigh => 10000,
119                };
120                if let Some(obj) = request_body.as_object_mut() {
121                    obj.insert("thinking_budget".to_string(), json!(budget));
122                    // 对于 low 和 none，禁用 thinking
123                    if matches!(effort, ReasoningEffort::None | ReasoningEffort::Low) {
124                        obj.insert("enable_thinking".to_string(), json!(false));
125                    } else {
126                        obj.insert("enable_thinking".to_string(), json!(true));
127                    }
128                }
129            }
130        } else if self.is_deepseek_model() {
131            if let Some(effort) = &config.effort {
132                let effort_str = match effort {
133                    ReasoningEffort::None => "none",
134                    ReasoningEffort::Low => "low",
135                    ReasoningEffort::Medium => "medium",
136                    ReasoningEffort::High => "high",
137                    ReasoningEffort::XHigh => "high",
138                };
139                if let Some(obj) = request_body.as_object_mut() {
140                    obj.insert("reasoning_effort".to_string(), json!(effort_str));
141                }
142            }
143            if config.enabled == Some(true) || config.budget_tokens.is_some() {
144                let mut extra_body = serde_json::Map::new();
145                if let Some(enabled) = config.enabled {
146                    extra_body.insert(
147                        "thinking".to_string(),
148                        json!({"type": if enabled { "enabled" } else { "disabled" }}),
149                    );
150                }
151                if let Some(budget) = config.budget_tokens {
152                    extra_body.insert("thinking_budget".to_string(), json!(budget));
153                }
154                if !extra_body.is_empty() {
155                    if let Some(obj) = request_body.as_object_mut() {
156                        obj.insert("extra_body".to_string(), Value::Object(extra_body));
157                    }
158                }
159            }
160        } else {
161            if let Some(effort) = &config.effort {
162                let effort_str = match effort {
163                    ReasoningEffort::None => "none",
164                    ReasoningEffort::Low => "low",
165                    ReasoningEffort::Medium => "medium",
166                    ReasoningEffort::High => "high",
167                    ReasoningEffort::XHigh => "high",
168                };
169                if let Some(obj) = request_body.as_object_mut() {
170                    obj.insert("reasoning_effort".to_string(), json!(effort_str));
171                }
172            }
173        }
174    }
175
176    fn chat_message_to_json(msg: &ChatMessage) -> Value {
177        match msg {
178            ChatMessage::System { content, .. } => json!({
179                "role": "system",
180                "content": content,
181            }),
182            ChatMessage::User {
183                content, images, ..
184            } => {
185                if images.is_empty() {
186                    json!({
187                        "role": "user",
188                        "content": content,
189                    })
190                } else {
191                    let mut content_parts: Vec<Value> = Vec::new();
192                    content_parts.push(json!({"type": "text", "text": content}));
193                    for img in images {
194                        content_parts.push(Self::image_to_json(img));
195                    }
196                    json!({
197                        "role": "user",
198                        "content": content_parts,
199                    })
200                }
201            }
202            ChatMessage::Assistant {
203                content,
204                reasoning_content,
205                tool_calls,
206            } => {
207                let mut obj = serde_json::Map::new();
208                obj.insert("role".to_string(), json!("assistant"));
209                obj.insert("content".to_string(), json!(content));
210                if let Some(reasoning) = reasoning_content {
211                    obj.insert("reasoning_content".to_string(), json!(reasoning));
212                }
213                if let Some(tc) = tool_calls {
214                    let tool_calls_json: Vec<Value> =
215                        tc.iter().map(|t| Self::tool_call_to_json(t)).collect();
216                    obj.insert("tool_calls".to_string(), json!(tool_calls_json));
217                }
218                Value::Object(obj)
219            }
220            ChatMessage::Tool {
221                tool_call_id,
222                content,
223            } => json!({
224                "role": "tool",
225                "tool_call_id": tool_call_id,
226                "content": content,
227            }),
228        }
229    }
230
231    fn tool_call_to_json(tc: &ToolCallMessage) -> Value {
232        json!({
233            "id": tc.id,
234            "type": "function",
235            "function": {
236                "name": tc.name,
237                "arguments": tc.arguments,
238            }
239        })
240    }
241
242    fn image_to_json(img: &ImageAttachment) -> Value {
243        match img {
244            ImageAttachment::Url { url, detail } => {
245                let mut obj = serde_json::Map::new();
246                obj.insert("url".to_string(), json!(url));
247                if let Some(d) = detail {
248                    let detail_str = match d {
249                        ImageDetail::Low => "low",
250                        ImageDetail::High => "high",
251                        ImageDetail::Auto => "auto",
252                    };
253                    obj.insert("detail".to_string(), json!(detail_str));
254                }
255                json!({
256                    "type": "image_url",
257                    "image_url": Value::Object(obj),
258                })
259            }
260            ImageAttachment::Base64 {
261                data,
262                media_type,
263                detail,
264            } => {
265                let mime = media_type.as_deref().unwrap_or("image/jpeg");
266                let data_url = format!("data:{mime};base64,{data}");
267                let mut obj = serde_json::Map::new();
268                obj.insert("url".to_string(), json!(data_url));
269                if let Some(d) = detail {
270                    let detail_str = match d {
271                        ImageDetail::Low => "low",
272                        ImageDetail::High => "high",
273                        ImageDetail::Auto => "auto",
274                    };
275                    obj.insert("detail".to_string(), json!(detail_str));
276                }
277                json!({
278                    "type": "image_url",
279                    "image_url": Value::Object(obj),
280                })
281            }
282        }
283    }
284
285    fn messages_to_json(messages: &[ChatMessage]) -> Vec<Value> {
286        messages.iter().map(Self::chat_message_to_json).collect()
287    }
288}
289
290#[async_trait]
291impl LlmClient for OpenAiClient {
292    async fn chat(
293        &self,
294        messages: &[ChatMessage],
295        tools: &[Value],
296        reasoning: Option<&ReasoningConfig>,
297        response_format: Option<&ResponseFormat>,
298    ) -> AgentResult<Value> {
299        let url = format!("{}/chat/completions", self.base_url);
300        let raw_messages = Self::messages_to_json(messages);
301        let mut request_body = json!({
302            "model": self.model,
303            "messages": raw_messages,
304            "tools": tools,
305            "max_tokens": 8192,
306        });
307
308        self.apply_reasoning_config(&mut request_body, reasoning);
309
310        if let Some(rf) = response_format {
311            if let Some(obj) = request_body.as_object_mut() {
312                obj.insert("response_format".to_string(), rf.to_api_value());
313            }
314        }
315
316        tracing::info!(model = %self.model, msg_count = messages.len(), "llm chat request");
317        tracing::debug!(request_body = %serde_json::to_string_pretty(&request_body).unwrap_or_default(), "llm request body");
318
319        let response = self
320            .client
321            .post(&url)
322            .header("Authorization", format!("Bearer {}", self.api_key))
323            .header("Content-Type", "application/json")
324            .json(&request_body)
325            .send()
326            .await
327            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
328
329        let status = response.status();
330        let res_json: Value = response
331            .json()
332            .await
333            .map_err(|e| AgentError::json(format!("Response JSON parse failed: {e}")))?;
334
335        if !status.is_success() {
336            tracing::warn!(%status, "OpenAI API non-success");
337        }
338
339        if let Some(error) = res_json.get("error") {
340            tracing::warn!(?error, "OpenAI API returned error");
341            return Err(AgentError::LlmApi {
342                message: format!("{error:#?}"),
343            });
344        }
345
346        Ok(res_json)
347    }
348
349    async fn chat_stream(
350        &self,
351        messages: &[ChatMessage],
352        tools: &[Value],
353        reasoning: Option<&ReasoningConfig>,
354        response_format: Option<&ResponseFormat>,
355    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
356        let url = format!("{}/chat/completions", self.base_url);
357        let raw_messages = Self::messages_to_json(messages);
358        let mut request_body = json!({
359            "model": self.model,
360            "messages": raw_messages,
361            "tools": tools,
362            "stream": true,
363            "stream_options": { "include_usage": true },
364            "max_tokens": 8192,
365        });
366
367        self.apply_reasoning_config(&mut request_body, reasoning);
368
369        if let Some(rf) = response_format {
370            if let Some(obj) = request_body.as_object_mut() {
371                obj.insert("response_format".to_string(), rf.to_api_value());
372            }
373        }
374
375        tracing::debug!(request_body = %serde_json::to_string_pretty(&request_body).unwrap_or_default(), "llm stream request body");
376
377        let response = self
378            .client
379            .post(&url)
380            .header("Authorization", format!("Bearer {}", self.api_key))
381            .header("Content-Type", "application/json")
382            .json(&request_body)
383            .send()
384            .await
385            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
386
387        if !response.status().is_success() {
388            let status = response.status();
389            let err_text = response
390                .text()
391                .await
392                .map_err(|e| AgentError::llm(format!("Failed to read error response: {e}")))?;
393            tracing::warn!(%status, error = %err_text, "OpenAI API stream non-success");
394            return Err(AgentError::LlmApi { message: err_text });
395        }
396
397        let stream = response
398            .bytes_stream()
399            .eventsource()
400            .map(|event| match event {
401                Ok(event) => {
402                    if event.data == "[DONE]" {
403                        return Ok(StreamChunk::Stop);
404                    }
405
406                    let data: Value = serde_json::from_str(&event.data)
407                        .map_err(|e| AgentError::json(format!("JSON Parse error: {e}")))?;
408
409                    let choices = data.get("choices").and_then(Value::as_array);
410
411                    if choices.is_none() || choices.map_or(true, |c| c.is_empty()) {
412                        if let Some(usage) = data.get("usage") {
413                            return Ok(StreamChunk::Usage(UsageInfo {
414                                prompt_tokens: usage
415                                    .get("prompt_tokens")
416                                    .and_then(Value::as_u64)
417                                    .map(|v| v as u32),
418                                completion_tokens: usage
419                                    .get("completion_tokens")
420                                    .and_then(Value::as_u64)
421                                    .map(|v| v as u32),
422                                total_tokens: usage
423                                    .get("total_tokens")
424                                    .and_then(Value::as_u64)
425                                    .map(|v| v as u32),
426                            }));
427                        }
428                        return Ok(StreamChunk::Text(String::new()));
429                    }
430
431                    let choice = &choices.unwrap()[0];
432                    let delta = &choice["delta"];
433                    let finish_reason = choice["finish_reason"].as_str().unwrap_or("");
434
435                    if finish_reason == "tool_calls" || delta.get("tool_calls").is_some() {
436                        return Ok(StreamChunk::ToolCall(choice.clone()));
437                    }
438
439                    if let Some(reasoning) = delta.get("reasoning_content") {
440                        if let Some(text) = reasoning.as_str() {
441                            return Ok(StreamChunk::Thought(text.to_string()));
442                        }
443                    }
444
445                    if let Some(content) = delta.get("content") {
446                        if let Some(text) = content.as_str() {
447                            return Ok(StreamChunk::Text(text.to_string()));
448                        }
449                    }
450
451                    if finish_reason == "stop" {
452                        return Ok(StreamChunk::Stop);
453                    }
454
455                    Ok(StreamChunk::Text(String::new()))
456                }
457                Err(e) => Err(AgentError::LlmStream(format!("SSE Stream error: {e}"))),
458            });
459
460        Ok(Box::pin(stream))
461    }
462
463    fn capabilities(&self) -> LlmCapabilities {
464        LlmCapabilities {
465            supports_streaming: true,
466            supports_tools: true,
467            supports_vision: true,
468            supports_thinking: true,
469            max_context_tokens: Some(128_000),
470            max_output_tokens: Some(16_384),
471        }
472    }
473
474    fn model_name(&self) -> &str {
475        &self.model
476    }
477}
agent_base/llm/openai.rs

agent_base/llm/
openai.rs
