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

agent_base/llm/
openai.rs
