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

agent_base/llm/
openai.rs
