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                && let Some(obj) = request_body.as_object_mut()
103            {
104                obj.insert("enable_thinking".to_string(), json!(enabled));
105            }
106            if let Some(budget) = config.budget_tokens
107                && let Some(obj) = request_body.as_object_mut()
108            {
109                obj.insert("thinking_budget".to_string(), json!(budget));
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                    && let Some(obj) = request_body.as_object_mut()
156                {
157                    obj.insert("extra_body".to_string(), Value::Object(extra_body));
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(Self::tool_call_to_json).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            ChatMessage::Custom { role, data } => json!({
229                "role": role,
230                "content": data.to_string(),
231            }),
232        }
233    }
234
235    fn tool_call_to_json(tc: &ToolCallMessage) -> Value {
236        json!({
237            "id": tc.id,
238            "type": "function",
239            "function": {
240                "name": tc.name,
241                "arguments": tc.arguments,
242            }
243        })
244    }
245
246    fn image_to_json(img: &ImageAttachment) -> Value {
247        match img {
248            ImageAttachment::Url { url, detail } => {
249                let mut obj = serde_json::Map::new();
250                obj.insert("url".to_string(), json!(url));
251                if let Some(d) = detail {
252                    let detail_str = match d {
253                        ImageDetail::Low => "low",
254                        ImageDetail::High => "high",
255                        ImageDetail::Auto => "auto",
256                    };
257                    obj.insert("detail".to_string(), json!(detail_str));
258                }
259                json!({
260                    "type": "image_url",
261                    "image_url": Value::Object(obj),
262                })
263            }
264            ImageAttachment::Base64 {
265                data,
266                media_type,
267                detail,
268            } => {
269                let mime = media_type.as_deref().unwrap_or("image/jpeg");
270                let data_url = format!("data:{mime};base64,{data}");
271                let mut obj = serde_json::Map::new();
272                obj.insert("url".to_string(), json!(data_url));
273                if let Some(d) = detail {
274                    let detail_str = match d {
275                        ImageDetail::Low => "low",
276                        ImageDetail::High => "high",
277                        ImageDetail::Auto => "auto",
278                    };
279                    obj.insert("detail".to_string(), json!(detail_str));
280                }
281                json!({
282                    "type": "image_url",
283                    "image_url": Value::Object(obj),
284                })
285            }
286        }
287    }
288
289    fn messages_to_json(messages: &[ChatMessage]) -> Vec<Value> {
290        messages.iter().map(Self::chat_message_to_json).collect()
291    }
292}
293
294#[async_trait]
295impl LlmClient for OpenAiClient {
296    async fn chat(
297        &self,
298        messages: &[ChatMessage],
299        tools: &[Value],
300        reasoning: Option<&ReasoningConfig>,
301        response_format: Option<&ResponseFormat>,
302    ) -> AgentResult<Value> {
303        let url = format!("{}/chat/completions", self.base_url);
304        let raw_messages = Self::messages_to_json(messages);
305        let mut request_body = json!({
306            "model": self.model,
307            "messages": raw_messages,
308            "tools": tools,
309            "max_tokens": 8192,
310        });
311
312        self.apply_reasoning_config(&mut request_body, reasoning);
313
314        if let Some(rf) = response_format
315            && let Some(obj) = request_body.as_object_mut()
316        {
317            obj.insert("response_format".to_string(), rf.to_api_value());
318        }
319
320        tracing::info!(model = %self.model, msg_count = messages.len(), "llm chat request");
321        tracing::debug!(request_body = %serde_json::to_string_pretty(&request_body).unwrap_or_default(), "llm request body");
322
323        let response = self
324            .client
325            .post(&url)
326            .header("Authorization", format!("Bearer {}", self.api_key))
327            .header("Content-Type", "application/json")
328            .json(&request_body)
329            .send()
330            .await
331            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
332
333        let status = response.status();
334        let res_json: Value = response
335            .json()
336            .await
337            .map_err(|e| AgentError::json(format!("Response JSON parse failed: {e}")))?;
338
339        if !status.is_success() {
340            tracing::warn!(%status, "OpenAI API non-success");
341        }
342
343        if let Some(error) = res_json.get("error") {
344            tracing::warn!(?error, "OpenAI API returned error");
345            return Err(AgentError::LlmApi {
346                message: format!("{error:#?}"),
347            });
348        }
349
350        Ok(res_json)
351    }
352
353    async fn chat_stream(
354        &self,
355        messages: &[ChatMessage],
356        tools: &[Value],
357        reasoning: Option<&ReasoningConfig>,
358        response_format: Option<&ResponseFormat>,
359    ) -> AgentResult<Pin<Box<dyn Stream<Item = AgentResult<StreamChunk>> + Send>>> {
360        let url = format!("{}/chat/completions", self.base_url);
361        let raw_messages = Self::messages_to_json(messages);
362        let mut request_body = json!({
363            "model": self.model,
364            "messages": raw_messages,
365            "tools": tools,
366            "stream": true,
367            "stream_options": { "include_usage": true },
368            "max_tokens": 8192,
369        });
370
371        self.apply_reasoning_config(&mut request_body, reasoning);
372
373        if let Some(rf) = response_format
374            && let Some(obj) = request_body.as_object_mut()
375        {
376            obj.insert("response_format".to_string(), rf.to_api_value());
377        }
378
379        tracing::debug!(request_body = %serde_json::to_string_pretty(&request_body).unwrap_or_default(), "llm stream request body");
380
381        let response = self
382            .client
383            .post(&url)
384            .header("Authorization", format!("Bearer {}", self.api_key))
385            .header("Content-Type", "application/json")
386            .json(&request_body)
387            .send()
388            .await
389            .map_err(|e| AgentError::llm(format!("HTTP request failed: {e}")))?;
390
391        if !response.status().is_success() {
392            let status = response.status();
393            let err_text = response
394                .text()
395                .await
396                .map_err(|e| AgentError::llm(format!("Failed to read error response: {e}")))?;
397            tracing::warn!(%status, error = %err_text, "OpenAI API stream non-success");
398            return Err(AgentError::LlmApi { message: err_text });
399        }
400
401        let stream = response
402            .bytes_stream()
403            .eventsource()
404            .map(|event| match event {
405                Ok(event) => {
406                    if event.data == "[DONE]" {
407                        return Ok(StreamChunk::Stop {
408                            finish_reason: None,
409                        });
410                    }
411
412                    let data: Value = serde_json::from_str(&event.data)
413                        .map_err(|e| AgentError::json(format!("JSON Parse error: {e}")))?;
414
415                    let choices = data.get("choices").and_then(Value::as_array);
416
417                    if choices.is_none() || choices.is_none_or(|c| c.is_empty()) {
418                        if let Some(usage) = data.get("usage") {
419                            return Ok(StreamChunk::Usage(UsageInfo {
420                                prompt_tokens: usage
421                                    .get("prompt_tokens")
422                                    .and_then(Value::as_u64)
423                                    .map(|v| v as u32),
424                                completion_tokens: usage
425                                    .get("completion_tokens")
426                                    .and_then(Value::as_u64)
427                                    .map(|v| v as u32),
428                                total_tokens: usage
429                                    .get("total_tokens")
430                                    .and_then(Value::as_u64)
431                                    .map(|v| v as u32),
432                            }));
433                        }
434                        return Ok(StreamChunk::Text(String::new()));
435                    }
436
437                    let choice = &choices.unwrap()[0];
438                    let delta = &choice["delta"];
439                    let finish_reason = choice["finish_reason"].as_str().unwrap_or("");
440
441                    if finish_reason == "tool_calls" || delta.get("tool_calls").is_some() {
442                        return Ok(StreamChunk::ToolCall(choice.clone()));
443                    }
444
445                    if let Some(reasoning) = delta.get("reasoning_content")
446                        && let Some(text) = reasoning.as_str()
447                    {
448                        return Ok(StreamChunk::Thought(text.to_string()));
449                    }
450
451                    if let Some(content) = delta.get("content")
452                        && let Some(text) = content.as_str()
453                    {
454                        return Ok(StreamChunk::Text(text.to_string()));
455                    }
456
457                    if finish_reason == "stop" || finish_reason == "length" {
458                        return Ok(StreamChunk::Stop {
459                            finish_reason: Some(finish_reason.to_string()),
460                        });
461                    }
462
463                    Ok(StreamChunk::Text(String::new()))
464                }
465                Err(e) => Err(AgentError::LlmStream(format!("SSE Stream error: {e}"))),
466            });
467
468        Ok(Box::pin(stream))
469    }
470
471    fn capabilities(&self) -> LlmCapabilities {
472        LlmCapabilities {
473            supports_streaming: true,
474            supports_tools: true,
475            supports_vision: true,
476            supports_thinking: true,
477            max_context_tokens: Some(128_000),
478            max_output_tokens: Some(16_384),
479        }
480    }
481
482    fn model_name(&self) -> &str {
483        &self.model
484    }
485}
486
487#[async_trait]
488impl super::StreamClient for OpenAiClient {
489    async fn stream(
490        &self,
491        messages: &[crate::types::ChatMessage],
492        tools: &[serde_json::Value],
493        reasoning: Option<&super::ReasoningConfig>,
494        response_format: Option<&crate::types::ResponseFormat>,
495    ) -> crate::types::AgentResult<
496        std::pin::Pin<
497            Box<
498                dyn futures_core::Stream<Item = crate::types::AgentResult<super::StreamChunk>>
499                    + Send,
500            >,
501        >,
502    > {
503        <Self as super::LlmClient>::chat_stream(self, messages, tools, reasoning, response_format)
504            .await
505    }
506
507    fn capabilities(&self) -> super::LlmCapabilities {
508        <Self as super::LlmClient>::capabilities(self)
509    }
510
511    fn model_name(&self) -> &str {
512        <Self as super::LlmClient>::model_name(self)
513    }
514}
agent_base/llm/openai.rs

agent_base/llm/
openai.rs
