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