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}
515
516#[cfg(test)]
517mod tests {
518    use super::*;
519    use crate::types::{ChatMessage, ImageAttachment, ImageDetail, ToolCallMessage};
520    use futures_util::TryStreamExt;
521    use wiremock::matchers::{method, path};
522    use wiremock::{Mock, MockServer, ResponseTemplate};
523
524    fn client(model: &str) -> OpenAiClient {
525        OpenAiClient::new("test-key".into(), model.into(), None)
526    }
527
528    // ---- pure helpers ----
529
530    #[test]
531    fn chat_message_to_json_all_variants() {
532        let v = OpenAiClient::chat_message_to_json(&ChatMessage::system("be helpful"));
533        assert_eq!(
534            v,
535            serde_json::json!({"role": "system", "content": "be helpful"})
536        );
537
538        let v = OpenAiClient::chat_message_to_json(&ChatMessage::user("hi"));
539        assert_eq!(v, serde_json::json!({"role": "user", "content": "hi"}));
540
541        let v = OpenAiClient::chat_message_to_json(&ChatMessage::user_with_images(
542            "look",
543            vec![ImageAttachment::Url {
544                url: "http://x/a.png".into(),
545                detail: None,
546            }],
547        ));
548        assert_eq!(v["role"], "user");
549        assert_eq!(
550            v["content"][0],
551            serde_json::json!({"type": "text", "text": "look"})
552        );
553        assert_eq!(v["content"][1]["type"], "image_url");
554
555        let v = OpenAiClient::chat_message_to_json(&ChatMessage::assistant("hello"));
556        assert_eq!(v["role"], "assistant");
557        assert_eq!(v["content"], "hello");
558
559        let v = OpenAiClient::chat_message_to_json(&ChatMessage::assistant_with_reasoning(
560            "answer",
561            "let me think",
562        ));
563        assert_eq!(v["reasoning_content"], "let me think");
564
565        let v = OpenAiClient::chat_message_to_json(&ChatMessage::assistant_tool_call(
566            "call_1", "echo", "{}",
567        ));
568        assert_eq!(v["tool_calls"][0]["function"]["name"], "echo");
569        assert_eq!(v["tool_calls"][0]["type"], "function");
570
571        let v = OpenAiClient::chat_message_to_json(&ChatMessage::tool("tid", "result"));
572        assert_eq!(
573            v,
574            serde_json::json!({"role": "tool", "tool_call_id": "tid", "content": "result"})
575        );
576
577        let v = OpenAiClient::chat_message_to_json(&ChatMessage::Custom {
578            role: "artifact".into(),
579            data: serde_json::json!({"x": 1}),
580        });
581        assert_eq!(v["role"], "artifact");
582        assert_eq!(v["content"], "{\"x\":1}");
583    }
584
585    #[test]
586    fn image_to_json_url_and_base64() {
587        let v = OpenAiClient::image_to_json(&ImageAttachment::Url {
588            url: "http://x/a.png".into(),
589            detail: None,
590        });
591        assert_eq!(
592            v,
593            serde_json::json!({"type": "image_url", "image_url": {"url": "http://x/a.png"}})
594        );
595
596        let v = OpenAiClient::image_to_json(&ImageAttachment::Url {
597            url: "http://x/a.png".into(),
598            detail: Some(ImageDetail::High),
599        });
600        assert_eq!(v["image_url"]["detail"], "high");
601
602        let v = OpenAiClient::image_to_json(&ImageAttachment::Base64 {
603            data: "abc".into(),
604            media_type: Some("image/png".into()),
605            detail: None,
606        });
607        assert_eq!(v["image_url"]["url"], "data:image/png;base64,abc");
608
609        let v = OpenAiClient::image_to_json(&ImageAttachment::Base64 {
610            data: "abc".into(),
611            media_type: None,
612            detail: Some(ImageDetail::Low),
613        });
614        assert_eq!(v["image_url"]["url"], "data:image/jpeg;base64,abc");
615        assert_eq!(v["image_url"]["detail"], "low");
616    }
617
618    #[test]
619    fn tool_call_to_json_shape() {
620        let tc = ToolCallMessage {
621            id: "call_1".into(),
622            name: "echo".into(),
623            arguments: "{\"x\":1}".into(),
624        };
625        let v = OpenAiClient::tool_call_to_json(&tc);
626        assert_eq!(v["id"], "call_1");
627        assert_eq!(v["type"], "function");
628        assert_eq!(v["function"]["name"], "echo");
629        assert_eq!(v["function"]["arguments"], "{\"x\":1}");
630    }
631
632    #[test]
633    fn apply_reasoning_config_none_is_noop() {
634        let c = client("gpt-4o");
635        let mut body = serde_json::json!({"model": "gpt-4o"});
636        c.apply_reasoning_config(&mut body, None);
637        assert_eq!(body, serde_json::json!({"model": "gpt-4o"}));
638    }
639
640    #[test]
641    fn apply_reasoning_config_qwen_flags() {
642        let c = client("qwen-max");
643        let mut body = serde_json::json!({"model": "qwen-max"});
644        let rc = ReasoningConfig {
645            enabled: Some(true),
646            budget_tokens: Some(1000),
647            effort: None,
648        };
649        c.apply_reasoning_config(&mut body, Some(&rc));
650        assert_eq!(body["enable_thinking"], serde_json::json!(true));
651        assert_eq!(body["thinking_budget"], serde_json::json!(1000));
652    }
653
654    #[test]
655    fn apply_reasoning_config_qwen_effort_maps_budget() {
656        let c = client("qwen-max");
657        let mut body = serde_json::json!({"model": "qwen-max"});
658        let rc = ReasoningConfig {
659            enabled: None,
660            budget_tokens: None,
661            effort: Some(ReasoningEffort::High),
662        };
663        c.apply_reasoning_config(&mut body, Some(&rc));
664        assert_eq!(body["thinking_budget"], serde_json::json!(5000));
665        assert_eq!(body["enable_thinking"], serde_json::json!(true));
666
667        let mut body = serde_json::json!({"model": "qwen-max"});
668        let rc = ReasoningConfig {
669            enabled: None,
670            budget_tokens: None,
671            effort: Some(ReasoningEffort::Low),
672        };
673        c.apply_reasoning_config(&mut body, Some(&rc));
674        assert_eq!(body["enable_thinking"], serde_json::json!(false));
675    }
676
677    #[test]
678    fn apply_reasoning_config_deepseek_effort() {
679        let c = client("deepseek-chat");
680        let mut body = serde_json::json!({"model": "deepseek-chat"});
681        let rc = ReasoningConfig {
682            enabled: None,
683            budget_tokens: None,
684            effort: Some(ReasoningEffort::Medium),
685        };
686        c.apply_reasoning_config(&mut body, Some(&rc));
687        assert_eq!(body["reasoning_effort"], serde_json::json!("medium"));
688    }
689
690    #[test]
691    fn apply_reasoning_config_deepseek_thinking_extra_body() {
692        let c = client("deepseek-chat");
693        let mut body = serde_json::json!({"model": "deepseek-chat"});
694        let rc = ReasoningConfig {
695            enabled: Some(true),
696            budget_tokens: Some(2000),
697            effort: None,
698        };
699        c.apply_reasoning_config(&mut body, Some(&rc));
700        assert_eq!(
701            body["extra_body"]["thinking"]["type"],
702            serde_json::json!("enabled")
703        );
704        assert_eq!(
705            body["extra_body"]["thinking_budget"],
706            serde_json::json!(2000)
707        );
708    }
709
710    #[test]
711    fn apply_reasoning_config_other_effort() {
712        let c = client("gpt-4o");
713        let mut body = serde_json::json!({"model": "gpt-4o"});
714        let rc = ReasoningConfig {
715            enabled: None,
716            budget_tokens: None,
717            effort: Some(ReasoningEffort::XHigh),
718        };
719        c.apply_reasoning_config(&mut body, Some(&rc));
720        assert_eq!(body["reasoning_effort"], serde_json::json!("high"));
721    }
722
723    #[test]
724    fn with_model_and_capabilities() {
725        let c = OpenAiClient::new("k".into(), "gpt-4o".into(), Some("http://x".into()));
726        let c2 = c.with_model("gpt-4o-mini");
727        assert_eq!(c2.model_name(), "gpt-4o-mini");
728        assert_eq!(c.model_name(), "gpt-4o");
729
730        let caps = c.capabilities();
731        assert!(caps.supports_streaming);
732        assert!(caps.supports_tools);
733        assert!(caps.supports_vision);
734        assert!(caps.supports_thinking);
735        assert_eq!(caps.max_context_tokens, Some(128_000));
736        assert_eq!(caps.max_output_tokens, Some(16_384));
737    }
738
739    // ---- mock HTTP ----
740
741    #[tokio::test]
742    async fn chat_posts_and_parses_response() {
743        let server = MockServer::start().await;
744        Mock::given(method("POST"))
745            .and(path("/chat/completions"))
746            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
747                "id": "cmpl-1",
748                "choices": [{"message": {"role": "assistant", "content": "hi there"}}],
749            })))
750            .mount(&server)
751            .await;
752
753        let client = OpenAiClient::new("test-key".into(), "gpt-4o".into(), Some(server.uri()));
754        let resp = client
755            .chat(&[ChatMessage::user("hello")], &[], None, None)
756            .await
757            .unwrap();
758        assert_eq!(resp["choices"][0]["message"]["content"], "hi there");
759    }
760
761    #[tokio::test]
762    async fn chat_returns_llm_api_error() {
763        let server = MockServer::start().await;
764        Mock::given(method("POST"))
765            .and(path("/chat/completions"))
766            .respond_with(ResponseTemplate::new(400).set_body_json(serde_json::json!({
767                "error": {"message": "invalid api key", "type": "invalid_request_error"},
768            })))
769            .mount(&server)
770            .await;
771
772        let client = OpenAiClient::new("bad-key".into(), "gpt-4o".into(), Some(server.uri()));
773        let resp = client
774            .chat(&[ChatMessage::user("hello")], &[], None, None)
775            .await;
776        assert!(resp.is_err());
777    }
778
779    #[tokio::test]
780    async fn chat_stream_parses_text_thought_toolcall_and_stop() {
781        let server = MockServer::start().await;
782        let sse = concat!(
783            "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":\"\"}]}\n\n",
784            "data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"hmm\"},\"finish_reason\":\"\"}]}\n\n",
785            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"name\":\"echo\",\"arguments\":\"\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n",
786            "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
787            "data: [DONE]\n\n",
788        );
789        Mock::given(method("POST"))
790            .and(path("/chat/completions"))
791            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
792            .mount(&server)
793            .await;
794
795        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
796        let stream = client
797            .chat_stream(&[ChatMessage::user("hi")], &[], None, None)
798            .await
799            .unwrap();
800        let chunks: Vec<StreamChunk> = stream.try_collect().await.unwrap();
801
802        assert!(matches!(&chunks[0], StreamChunk::Text(t) if t == "Hello"));
803        assert!(matches!(&chunks[1], StreamChunk::Thought(t) if t == "hmm"));
804        assert!(matches!(&chunks[2], StreamChunk::ToolCall(_)));
805        assert!(matches!(&chunks[3], StreamChunk::Stop { finish_reason: Some(r) } if r == "stop"));
806        assert!(matches!(
807            &chunks[4],
808            StreamChunk::Stop {
809                finish_reason: None
810            }
811        ));
812    }
813
814    #[tokio::test]
815    async fn chat_stream_parses_usage_chunk() {
816        let server = MockServer::start().await;
817        let sse = concat!(
818            "data: {\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30}}\n\n",
819            "data: [DONE]\n\n",
820        );
821        Mock::given(method("POST"))
822            .and(path("/chat/completions"))
823            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
824            .mount(&server)
825            .await;
826
827        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
828        let stream = client
829            .chat_stream(&[ChatMessage::user("hi")], &[], None, None)
830            .await
831            .unwrap();
832        let chunks: Vec<StreamChunk> = stream.try_collect().await.unwrap();
833
834        assert!(
835            matches!(&chunks[0], StreamChunk::Usage(u) if u.prompt_tokens == Some(10) && u.completion_tokens == Some(20) && u.total_tokens == Some(30))
836        );
837        assert!(matches!(
838            &chunks[1],
839            StreamChunk::Stop {
840                finish_reason: None
841            }
842        ));
843    }
844
845    #[tokio::test]
846    async fn chat_stream_returns_error_on_non_success() {
847        let server = MockServer::start().await;
848        Mock::given(method("POST"))
849            .and(path("/chat/completions"))
850            .respond_with(ResponseTemplate::new(401).set_body_string("unauthorized"))
851            .mount(&server)
852            .await;
853
854        let client = OpenAiClient::new("bad".into(), "gpt-4o".into(), Some(server.uri()));
855        let resp = client
856            .chat_stream(&[ChatMessage::user("hi")], &[], None, None)
857            .await;
858        assert!(resp.is_err());
859    }
860
861    #[tokio::test]
862    async fn chat_stream_errors_on_invalid_json() {
863        let server = MockServer::start().await;
864        let sse = "data: not-json\n\n";
865        Mock::given(method("POST"))
866            .and(path("/chat/completions"))
867            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
868            .mount(&server)
869            .await;
870
871        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
872        let stream = client
873            .chat_stream(&[ChatMessage::user("hi")], &[], None, None)
874            .await
875            .unwrap();
876        let result: Result<Vec<_>, _> = stream.try_collect().await;
877        assert!(result.is_err());
878    }
879
880    // ── B2: remaining enum arms + adapter delegation + HTTP edges ────────
881
882    #[test]
883    fn image_to_json_detail_variants() {
884        // Url + Low / Auto
885        let v = OpenAiClient::image_to_json(&ImageAttachment::Url {
886            url: "http://x/a.png".into(),
887            detail: Some(ImageDetail::Low),
888        });
889        assert_eq!(v["image_url"]["detail"], "low");
890
891        let v = OpenAiClient::image_to_json(&ImageAttachment::Url {
892            url: "http://x/a.png".into(),
893            detail: Some(ImageDetail::Auto),
894        });
895        assert_eq!(v["image_url"]["detail"], "auto");
896
897        // Base64 + High / Auto
898        let v = OpenAiClient::image_to_json(&ImageAttachment::Base64 {
899            data: "abc".into(),
900            media_type: None,
901            detail: Some(ImageDetail::High),
902        });
903        assert_eq!(v["image_url"]["detail"], "high");
904
905        let v = OpenAiClient::image_to_json(&ImageAttachment::Base64 {
906            data: "abc".into(),
907            media_type: None,
908            detail: Some(ImageDetail::Auto),
909        });
910        assert_eq!(v["image_url"]["detail"], "auto");
911    }
912
913    #[test]
914    fn apply_reasoning_config_qwen_effort_all_levels() {
915        // None/Low already covered; exercise None→0, Medium→2000, XHigh→10000.
916        for (effort, budget, enabled) in [
917            (ReasoningEffort::None, 0, false),
918            (ReasoningEffort::Medium, 2000, true),
919            (ReasoningEffort::XHigh, 10000, true),
920        ] {
921            let c = client("qwen-max");
922            let mut body = serde_json::json!({"model": "qwen-max"});
923            let rc = ReasoningConfig {
924                enabled: None,
925                budget_tokens: None,
926                effort: Some(effort),
927            };
928            c.apply_reasoning_config(&mut body, Some(&rc));
929            assert_eq!(body["thinking_budget"], serde_json::json!(budget));
930            assert_eq!(body["enable_thinking"], serde_json::json!(enabled));
931        }
932    }
933
934    #[test]
935    fn apply_reasoning_config_deepseek_effort_all_levels() {
936        // Medium already covered; exercise None/Low/High/XHigh.
937        for (effort, expected) in [
938            (ReasoningEffort::None, "none"),
939            (ReasoningEffort::Low, "low"),
940            (ReasoningEffort::High, "high"),
941            (ReasoningEffort::XHigh, "high"),
942        ] {
943            let c = client("deepseek-chat");
944            let mut body = serde_json::json!({"model": "deepseek-chat"});
945            let rc = ReasoningConfig {
946                enabled: None,
947                budget_tokens: None,
948                effort: Some(effort),
949            };
950            c.apply_reasoning_config(&mut body, Some(&rc));
951            assert_eq!(body["reasoning_effort"], serde_json::json!(expected));
952        }
953    }
954
955    #[test]
956    fn apply_reasoning_config_other_effort_all_levels() {
957        // XHigh already covered; exercise None/Low/Medium/High.
958        for (effort, expected) in [
959            (ReasoningEffort::None, "none"),
960            (ReasoningEffort::Low, "low"),
961            (ReasoningEffort::Medium, "medium"),
962            (ReasoningEffort::High, "high"),
963        ] {
964            let c = client("gpt-4o");
965            let mut body = serde_json::json!({"model": "gpt-4o"});
966            let rc = ReasoningConfig {
967                enabled: None,
968                budget_tokens: None,
969                effort: Some(effort),
970            };
971            c.apply_reasoning_config(&mut body, Some(&rc));
972            assert_eq!(body["reasoning_effort"], serde_json::json!(expected));
973        }
974    }
975
976    #[test]
977    fn stream_client_delegates_capabilities_and_model_name() {
978        let c = OpenAiClient::new("k".into(), "gpt-4o".into(), None);
979        let caps = <OpenAiClient as crate::llm::StreamClient>::capabilities(&c);
980        assert!(caps.supports_streaming);
981        assert_eq!(
982            <OpenAiClient as crate::llm::StreamClient>::model_name(&c),
983            "gpt-4o"
984        );
985    }
986
987    #[tokio::test]
988    async fn stream_client_stream_delegates_to_chat_stream() {
989        let server = MockServer::start().await;
990        let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"},\"finish_reason\":\"\"}]}\n\ndata: [DONE]\n\n";
991        Mock::given(method("POST"))
992            .and(path("/chat/completions"))
993            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
994            .mount(&server)
995            .await;
996
997        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
998        let stream = <OpenAiClient as crate::llm::StreamClient>::stream(
999            &client,
1000            &[ChatMessage::user("hi")],
1001            &[],
1002            None,
1003            None,
1004        )
1005        .await
1006        .unwrap();
1007        let chunks: Vec<StreamChunk> = stream.try_collect().await.unwrap();
1008        assert!(matches!(&chunks[0], StreamChunk::Text(t) if t == "Hi"));
1009    }
1010
1011    #[tokio::test]
1012    async fn chat_errors_on_non_json_response() {
1013        let server = MockServer::start().await;
1014        Mock::given(method("POST"))
1015            .and(path("/chat/completions"))
1016            .respond_with(ResponseTemplate::new(200).set_body_string("not-json"))
1017            .mount(&server)
1018            .await;
1019
1020        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
1021        let resp = client
1022            .chat(&[ChatMessage::user("hi")], &[], None, None)
1023            .await;
1024        assert!(resp.is_err());
1025    }
1026
1027    #[tokio::test]
1028    async fn chat_stream_emits_empty_text_for_no_choices_or_usage() {
1029        let server = MockServer::start().await;
1030        let sse = "data: {\"id\":\"x\"}\n\n";
1031        Mock::given(method("POST"))
1032            .and(path("/chat/completions"))
1033            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
1034            .mount(&server)
1035            .await;
1036
1037        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
1038        let stream = client
1039            .chat_stream(&[ChatMessage::user("hi")], &[], None, None)
1040            .await
1041            .unwrap();
1042        let chunks: Vec<StreamChunk> = stream.try_collect().await.unwrap();
1043        assert!(matches!(&chunks[0], StreamChunk::Text(t) if t.is_empty()));
1044    }
1045
1046    #[tokio::test]
1047    async fn chat_stream_emits_empty_text_for_empty_delta() {
1048        let server = MockServer::start().await;
1049        let sse = "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"\"}]}\n\n";
1050        Mock::given(method("POST"))
1051            .and(path("/chat/completions"))
1052            .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
1053            .mount(&server)
1054            .await;
1055
1056        let client = OpenAiClient::new("k".into(), "gpt-4o".into(), Some(server.uri()));
1057        let stream = client
1058            .chat_stream(&[ChatMessage::user("hi")], &[], None, None)
1059            .await
1060            .unwrap();
1061        let chunks: Vec<StreamChunk> = stream.try_collect().await.unwrap();
1062        assert!(matches!(&chunks[0], StreamChunk::Text(t) if t.is_empty()));
1063    }
1064}
agent_base/llm/openai.rs

agent_base/llm/
openai.rs
