a3s_code_core/llm/
openai.rs

1//! OpenAI-compatible LLM client
2
3use super::http::{default_http_client, normalize_base_url, HttpClient};
4use super::structured;
5use super::types::*;
6use super::LlmClient;
7use crate::llm::types::{ToolResultContent, ToolResultContentField};
8use crate::retry::{AttemptOutcome, RetryConfig};
9use anyhow::{Context, Result};
10use async_trait::async_trait;
11use futures::StreamExt;
12use serde::Deserialize;
13use std::collections::HashMap;
14use std::sync::Arc;
15use std::time::Instant;
16use tokio::sync::mpsc;
17
18/// OpenAI client
19pub struct OpenAiClient {
20    pub(crate) provider_name: String,
21    pub(crate) api_key: SecretString,
22    pub(crate) model: String,
23    pub(crate) base_url: String,
24    pub(crate) chat_completions_path: String,
25    pub(crate) headers: HashMap<String, String>,
26    pub(crate) temperature: Option<f32>,
27    pub(crate) max_tokens: Option<usize>,
28    pub(crate) logprobs: bool,
29    pub(crate) top_logprobs: Option<usize>,
30    pub(crate) http: Arc<dyn HttpClient>,
31    pub(crate) retry_config: RetryConfig,
32    pub(crate) native_structured_support: structured::NativeStructuredSupport,
33}
34
35impl OpenAiClient {
36    pub(crate) fn parse_tool_arguments(tool_name: &str, arguments: &str) -> serde_json::Value {
37        if arguments.trim().is_empty() {
38            return serde_json::Value::Object(Default::default());
39        }
40
41        serde_json::from_str(arguments).unwrap_or_else(|e| {
42            tracing::warn!(
43                "Failed to parse tool arguments JSON for tool '{}': {}",
44                tool_name,
45                e
46            );
47            serde_json::json!({
48                "__parse_error": format!(
49                    "Malformed tool arguments: {}. Raw input: {}",
50                    e, arguments
51                )
52            })
53        })
54    }
55
56    fn merge_stream_text(text_content: &mut String, incoming: &str) -> Option<String> {
57        if incoming.is_empty() {
58            return None;
59        }
60        if text_content.is_empty() {
61            text_content.push_str(incoming);
62            return Some(incoming.to_string());
63        }
64        if incoming == text_content.as_str() || text_content.ends_with(incoming) {
65            return None;
66        }
67        // If incoming contains text_content as a prefix (incoming is the full content),
68        // replace text_content instead of appending (avoids duplicate full content)
69        if incoming.starts_with(text_content.as_str()) && incoming.len() > text_content.len() {
70            let suffix = &incoming[text_content.len()..];
71            if !suffix.is_empty() {
72                *text_content = incoming.to_string();
73                return Some(suffix.to_string());
74            }
75            return None;
76        }
77        if let Some(suffix) = incoming.strip_prefix(text_content.as_str()) {
78            if suffix.is_empty() {
79                return None;
80            }
81            text_content.push_str(suffix);
82            return Some(suffix.to_string());
83        }
84        text_content.push_str(incoming);
85        Some(incoming.to_string())
86    }
87
88    pub fn new(api_key: String, model: String) -> Self {
89        Self {
90            provider_name: "openai".to_string(),
91            api_key: SecretString::new(api_key),
92            model,
93            base_url: "https://api.openai.com".to_string(),
94            chat_completions_path: "/v1/chat/completions".to_string(),
95            headers: HashMap::new(),
96            temperature: None,
97            max_tokens: None,
98            logprobs: false,
99            top_logprobs: None,
100            http: default_http_client(),
101            retry_config: RetryConfig::default(),
102            native_structured_support: structured::NativeStructuredSupport::JsonSchema,
103        }
104    }
105
106    pub fn with_base_url(mut self, base_url: String) -> Self {
107        self.base_url = normalize_base_url(&base_url);
108        self
109    }
110
111    pub fn with_provider_name(mut self, provider_name: impl Into<String>) -> Self {
112        self.provider_name = provider_name.into();
113        self
114    }
115
116    pub fn with_chat_completions_path(mut self, path: impl Into<String>) -> Self {
117        let path = path.into();
118        self.chat_completions_path = if path.starts_with('/') {
119            path
120        } else {
121            format!("/{}", path)
122        };
123        self
124    }
125
126    pub fn with_temperature(mut self, temperature: f32) -> Self {
127        self.temperature = Some(temperature);
128        self
129    }
130
131    pub fn with_headers(mut self, headers: HashMap<String, String>) -> Self {
132        self.headers = headers;
133        self
134    }
135
136    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
137        self.max_tokens = Some(max_tokens);
138        self
139    }
140
141    pub fn with_logprobs(mut self, enabled: bool) -> Self {
142        self.logprobs = enabled;
143        self
144    }
145
146    pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self {
147        self.logprobs = true;
148        self.top_logprobs = Some(top_logprobs);
149        self
150    }
151
152    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
153        self.retry_config = retry_config;
154        self
155    }
156
157    pub fn with_native_structured_support(
158        mut self,
159        support: structured::NativeStructuredSupport,
160    ) -> Self {
161        self.native_structured_support = support;
162        self
163    }
164
165    pub fn with_http_client(mut self, http: Arc<dyn HttpClient>) -> Self {
166        self.http = http;
167        self
168    }
169
170    pub(crate) fn request_headers(&self) -> Vec<(String, String)> {
171        let mut headers = Vec::with_capacity(self.headers.len() + 1);
172        let has_authorization = self
173            .headers
174            .keys()
175            .any(|key| key.eq_ignore_ascii_case("authorization"));
176        if !has_authorization {
177            headers.push((
178                "Authorization".to_string(),
179                format!("Bearer {}", self.api_key.expose()),
180            ));
181        }
182        headers.extend(
183            self.headers
184                .iter()
185                .map(|(key, value)| (key.clone(), value.clone())),
186        );
187        headers
188    }
189
190    pub(crate) fn convert_messages(&self, messages: &[Message]) -> Vec<serde_json::Value> {
191        messages
192            .iter()
193            .map(|msg| {
194                let content: serde_json::Value = if msg.content.len() == 1 {
195                    match &msg.content[0] {
196                        ContentBlock::Text { text } => serde_json::json!(text),
197                        ContentBlock::ToolResult {
198                            tool_use_id,
199                            content,
200                            ..
201                        } => {
202                            let content_str = match content {
203                                ToolResultContentField::Text(s) => s.clone(),
204                                ToolResultContentField::Blocks(blocks) => blocks
205                                    .iter()
206                                    .filter_map(|b| {
207                                        if let ToolResultContent::Text { text } = b {
208                                            Some(text.clone())
209                                        } else {
210                                            None
211                                        }
212                                    })
213                                    .collect::<Vec<_>>()
214                                    .join("\n"),
215                            };
216                            return serde_json::json!({
217                                "role": "tool",
218                                "tool_call_id": tool_use_id,
219                                "content": content_str,
220                            });
221                        }
222                        _ => serde_json::json!(""),
223                    }
224                } else {
225                    serde_json::json!(msg
226                        .content
227                        .iter()
228                        .map(|block| {
229                            match block {
230                                ContentBlock::Text { text } => serde_json::json!({
231                                    "type": "text",
232                                    "text": text,
233                                }),
234                                ContentBlock::Image { source } => serde_json::json!({
235                                    "type": "image_url",
236                                    "image_url": {
237                                        "url": format!(
238                                            "data:{};base64,{}",
239                                            source.media_type, source.data
240                                        ),
241                                    }
242                                }),
243                                ContentBlock::ToolUse { id, name, input } => serde_json::json!({
244                                    "type": "function",
245                                    "id": id,
246                                    "function": {
247                                        "name": name,
248                                        "arguments": input.to_string(),
249                                    }
250                                }),
251                                _ => serde_json::json!({}),
252                            }
253                        })
254                        .collect::<Vec<_>>())
255                };
256
257                // Handle assistant messages — kimi-k2.5 requires reasoning_content
258                // on all assistant messages when thinking mode is enabled
259                if msg.role == "assistant" {
260                    let rc = msg.reasoning_content.as_deref().unwrap_or("");
261                    let tool_calls: Vec<_> = msg.tool_calls();
262                    if !tool_calls.is_empty() {
263                        return serde_json::json!({
264                            "role": "assistant",
265                            "content": msg.text(),
266                            "reasoning_content": rc,
267                            "tool_calls": tool_calls.iter().map(|tc| {
268                                serde_json::json!({
269                                    "id": tc.id,
270                                    "type": "function",
271                                    "function": {
272                                        "name": tc.name,
273                                        "arguments": tc.args.to_string(),
274                                    }
275                                })
276                            }).collect::<Vec<_>>(),
277                        });
278                    }
279                    return serde_json::json!({
280                        "role": "assistant",
281                        "content": content,
282                        "reasoning_content": rc,
283                    });
284                }
285
286                serde_json::json!({
287                    "role": msg.role,
288                    "content": content,
289                })
290            })
291            .collect()
292    }
293
294    pub(crate) fn convert_tools(&self, tools: &[ToolDefinition]) -> Vec<serde_json::Value> {
295        tools
296            .iter()
297            .map(|t| {
298                serde_json::json!({
299                    "type": "function",
300                    "function": {
301                        "name": t.name,
302                        "description": t.description,
303                        "parameters": t.parameters,
304                    }
305                })
306            })
307            .collect()
308    }
309}
310
311impl OpenAiClient {
312    /// Apply a structured-output directive to an OpenAI-compatible request.
313    ///
314    /// OpenAI-compatible APIs support both forced function `tool_choice` and
315    /// native `response_format` (`json_object` / `json_schema` + `strict`).
316    fn apply_directive(
317        request: &mut serde_json::Value,
318        directive: &structured::StructuredDirective,
319    ) {
320        if let Some(tool) = &directive.force_tool {
321            request["tool_choice"] = serde_json::json!({
322                "type": "function",
323                "function": { "name": tool }
324            });
325        }
326        if let Some(rf) = &directive.response_format {
327            request["response_format"] = match rf {
328                structured::ResponseFormat::JsonObject => {
329                    serde_json::json!({ "type": "json_object" })
330                }
331                structured::ResponseFormat::JsonSchema { name, schema } => serde_json::json!({
332                    "type": "json_schema",
333                    "json_schema": { "name": name, "schema": schema, "strict": true }
334                }),
335            };
336        }
337    }
338
339    /// Build a chat-completions request body, optionally applying a directive.
340    fn build_chat_request(
341        &self,
342        messages: &[Message],
343        system: Option<&str>,
344        tools: &[ToolDefinition],
345        directive: Option<&structured::StructuredDirective>,
346    ) -> serde_json::Value {
347        let mut openai_messages = Vec::new();
348
349        if let Some(sys) = system {
350            openai_messages.push(serde_json::json!({
351                "role": "system",
352                "content": sys,
353            }));
354        }
355
356        openai_messages.extend(self.convert_messages(messages));
357
358        let mut request = serde_json::json!({
359            "model": self.model,
360            "messages": openai_messages,
361        });
362
363        if let Some(temp) = self.temperature {
364            request["temperature"] = serde_json::json!(temp);
365        }
366        if let Some(max) = self.max_tokens {
367            request["max_tokens"] = serde_json::json!(max);
368        }
369        if self.logprobs {
370            request["logprobs"] = serde_json::json!(true);
371            if let Some(top_logprobs) = self.top_logprobs {
372                request["top_logprobs"] = serde_json::json!(top_logprobs);
373            }
374        }
375
376        if !tools.is_empty() {
377            request["tools"] = serde_json::json!(self.convert_tools(tools));
378        }
379
380        if let Some(directive) = directive {
381            Self::apply_directive(&mut request, directive);
382        }
383
384        request
385    }
386
387    /// Execute a fully-built (non-streaming) chat-completions request.
388    async fn send_request(&self, request: serde_json::Value) -> Result<LlmResponse> {
389        {
390            let request_started_at = Instant::now();
391            let url = format!("{}{}", self.base_url, self.chat_completions_path);
392            let request_headers = self.request_headers();
393
394            let response = crate::retry::with_retry(&self.retry_config, |_attempt| {
395                let http = &self.http;
396                let url = &url;
397                let request_headers = request_headers.clone();
398                let request = &request;
399                async move {
400                    let headers = request_headers
401                        .iter()
402                        .map(|(key, value)| (key.as_str(), value.as_str()))
403                        .collect::<Vec<_>>();
404                    // Non-streaming: use a non-cancelled token for now
405                    let cancel_token = tokio_util::sync::CancellationToken::new();
406                    match http.post(url, headers, request, cancel_token).await {
407                        Ok(resp) => {
408                            let status = reqwest::StatusCode::from_u16(resp.status)
409                                .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR);
410                            if status.is_success() {
411                                AttemptOutcome::Success(resp.body)
412                            } else if self.retry_config.is_retryable_status(status) {
413                                AttemptOutcome::Retryable {
414                                    status,
415                                    body: resp.body,
416                                    retry_after: None,
417                                }
418                            } else {
419                                AttemptOutcome::Fatal(anyhow::anyhow!(
420                                    "OpenAI API error at {} ({}): {}",
421                                    url,
422                                    status,
423                                    resp.body
424                                ))
425                            }
426                        }
427                        Err(e) => {
428                            tracing::error!("HTTP error: {e:?}");
429                            if crate::llm::http::is_retryable_http_failure(&e) {
430                                AttemptOutcome::Retryable {
431                                    status: reqwest::StatusCode::SERVICE_UNAVAILABLE,
432                                    body: format!("network error: {e}"),
433                                    retry_after: None,
434                                }
435                            } else {
436                                AttemptOutcome::Fatal(e)
437                            }
438                        }
439                    }
440                }
441            })
442            .await?;
443
444            let parsed: OpenAiResponse =
445                serde_json::from_str(&response).context("Failed to parse OpenAI response")?;
446
447            let choice = parsed.choices.into_iter().next().context("No choices")?;
448            let token_logprobs = choice
449                .logprobs
450                .as_ref()
451                .map(openai_logprobs_to_token_logprobs)
452                .unwrap_or_default();
453
454            let mut content = vec![];
455
456            let reasoning_content = choice.message.reasoning_content;
457
458            let text_content = choice.message.content;
459
460            if let Some(text) = text_content {
461                if !text.is_empty() {
462                    content.push(ContentBlock::Text { text });
463                }
464            }
465
466            if let Some(tool_calls) = choice.message.tool_calls {
467                for tc in tool_calls {
468                    content.push(ContentBlock::ToolUse {
469                        id: tc.id,
470                        name: tc.function.name.clone(),
471                        input: Self::parse_tool_arguments(
472                            &tc.function.name,
473                            &tc.function.arguments,
474                        ),
475                    });
476                }
477            }
478
479            let llm_response = LlmResponse {
480                message: Message {
481                    role: "assistant".to_string(),
482                    content,
483                    reasoning_content,
484                },
485                usage: TokenUsage {
486                    prompt_tokens: parsed.usage.prompt_tokens,
487                    completion_tokens: parsed.usage.completion_tokens,
488                    total_tokens: {
489                        let t = parsed.usage.total_tokens;
490                        // MiniMax: fall back to total_characters when total_tokens is 0.
491                        if t == 0 {
492                            parsed.usage.total_characters.unwrap_or(0)
493                        } else {
494                            t
495                        }
496                    },
497                    cache_read_tokens: parsed
498                        .usage
499                        .prompt_tokens_details
500                        .as_ref()
501                        .and_then(|d| d.cached_tokens),
502                    cache_write_tokens: None,
503                },
504                stop_reason: choice.finish_reason,
505                token_logprobs,
506                meta: Some(LlmResponseMeta {
507                    provider: Some(self.provider_name.clone()),
508                    request_model: Some(self.model.clone()),
509                    request_url: Some(url.clone()),
510                    response_id: parsed.id,
511                    response_model: parsed.model,
512                    response_object: parsed.object,
513                    first_token_ms: None,
514                    duration_ms: Some(request_started_at.elapsed().as_millis() as u64),
515                }),
516            };
517
518            crate::telemetry::record_llm_usage(
519                llm_response.usage.prompt_tokens,
520                llm_response.usage.completion_tokens,
521                llm_response.usage.total_tokens,
522                llm_response.stop_reason.as_deref(),
523            );
524
525            Ok(llm_response)
526        }
527    }
528}
529
530#[async_trait]
531impl LlmClient for OpenAiClient {
532    async fn complete(
533        &self,
534        messages: &[Message],
535        system: Option<&str>,
536        tools: &[ToolDefinition],
537    ) -> Result<LlmResponse> {
538        self.send_request(self.build_chat_request(messages, system, tools, None))
539            .await
540    }
541
542    async fn complete_structured(
543        &self,
544        messages: &[Message],
545        system: Option<&str>,
546        tools: &[ToolDefinition],
547        directive: &structured::StructuredDirective,
548    ) -> Result<LlmResponse> {
549        self.send_request(self.build_chat_request(messages, system, tools, Some(directive)))
550            .await
551    }
552
553    fn native_structured_support(&self) -> structured::NativeStructuredSupport {
554        self.native_structured_support
555    }
556
557    async fn complete_streaming(
558        &self,
559        messages: &[Message],
560        system: Option<&str>,
561        tools: &[ToolDefinition],
562        cancel_token: tokio_util::sync::CancellationToken,
563    ) -> Result<mpsc::Receiver<StreamEvent>> {
564        self.send_streaming(
565            self.build_chat_request(messages, system, tools, None),
566            cancel_token,
567        )
568        .await
569    }
570
571    async fn complete_streaming_structured(
572        &self,
573        messages: &[Message],
574        system: Option<&str>,
575        tools: &[ToolDefinition],
576        directive: &structured::StructuredDirective,
577        cancel_token: tokio_util::sync::CancellationToken,
578    ) -> Result<mpsc::Receiver<StreamEvent>> {
579        self.send_streaming(
580            self.build_chat_request(messages, system, tools, Some(directive)),
581            cancel_token,
582        )
583        .await
584    }
585}
586
587#[path = "openai/streaming.rs"]
588mod streaming;
589use streaming::openai_logprobs_to_token_logprobs;
590
591// OpenAI API response types (private)
592#[derive(Debug, Deserialize)]
593pub(crate) struct OpenAiResponse {
594    #[serde(default)]
595    pub(crate) id: Option<String>,
596    #[serde(default)]
597    pub(crate) object: Option<String>,
598    #[serde(default)]
599    pub(crate) model: Option<String>,
600    pub(crate) choices: Vec<OpenAiChoice>,
601    pub(crate) usage: OpenAiUsage,
602}
603
604#[derive(Debug, Deserialize)]
605pub(crate) struct OpenAiChoice {
606    pub(crate) message: OpenAiMessage,
607    pub(crate) finish_reason: Option<String>,
608    #[serde(default)]
609    pub(crate) logprobs: Option<OpenAiChoiceLogprobs>,
610}
611
612#[derive(Debug, Deserialize)]
613pub(crate) struct OpenAiChoiceLogprobs {
614    #[serde(default)]
615    pub(crate) content: Option<Vec<OpenAiTokenLogprob>>,
616}
617
618#[derive(Debug, Deserialize)]
619pub(crate) struct OpenAiTokenLogprob {
620    pub(crate) token: String,
621    pub(crate) logprob: f64,
622    #[serde(default)]
623    pub(crate) bytes: Option<Vec<u8>>,
624    #[serde(default)]
625    pub(crate) top_logprobs: Vec<OpenAiTopLogprob>,
626}
627
628#[derive(Debug, Deserialize)]
629pub(crate) struct OpenAiTopLogprob {
630    pub(crate) token: String,
631    pub(crate) logprob: f64,
632    #[serde(default)]
633    pub(crate) bytes: Option<Vec<u8>>,
634}
635
636#[derive(Debug, Deserialize)]
637pub(crate) struct OpenAiMessage {
638    // glm5.1 (and other GLM/zhipu reasoning models) stream/return reasoning under
639    // `reasoning`, not the `reasoning_content` kimi/deepseek use. Without this alias the
640    // reasoning phase yields zero recognized deltas → no ReasoningDelta events → the
641    // stream-stall watchdog kills long reasoning mid-flight (asset-diagnose "未返回结构化输出").
642    #[serde(alias = "reasoning")]
643    pub(crate) reasoning_content: Option<String>,
644    pub(crate) content: Option<String>,
645    pub(crate) tool_calls: Option<Vec<OpenAiToolCall>>,
646}
647
648#[derive(Debug, Deserialize)]
649pub(crate) struct OpenAiToolCall {
650    pub(crate) id: String,
651    pub(crate) function: OpenAiFunction,
652}
653
654#[derive(Debug, Deserialize)]
655pub(crate) struct OpenAiFunction {
656    pub(crate) name: String,
657    pub(crate) arguments: String,
658}
659
660#[derive(Debug, Deserialize)]
661pub(crate) struct OpenAiUsage {
662    #[serde(default)]
663    pub(crate) prompt_tokens: usize,
664    #[serde(default)]
665    pub(crate) completion_tokens: usize,
666    #[serde(default)]
667    pub(crate) total_tokens: usize,
668    /// MiniMax uses `total_characters` instead of token counts.
669    #[serde(default)]
670    pub(crate) total_characters: Option<usize>,
671    /// OpenAI returns cached token count in `prompt_tokens_details.cached_tokens`
672    #[serde(default)]
673    pub(crate) prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
674}
675
676#[derive(Debug, Deserialize)]
677pub(crate) struct OpenAiPromptTokensDetails {
678    #[serde(default)]
679    pub(crate) cached_tokens: Option<usize>,
680}
681
682// OpenAI streaming types
683#[derive(Debug, Deserialize)]
684pub(crate) struct OpenAiStreamChunk {
685    #[serde(default)]
686    pub(crate) id: Option<String>,
687    #[serde(default)]
688    pub(crate) object: Option<String>,
689    #[serde(default)]
690    pub(crate) model: Option<String>,
691    pub(crate) choices: Vec<OpenAiStreamChoice>,
692    pub(crate) usage: Option<OpenAiUsage>,
693}
694
695#[derive(Debug, Deserialize)]
696pub(crate) struct OpenAiStreamChoice {
697    pub(crate) message: Option<OpenAiMessage>,
698    pub(crate) delta: Option<OpenAiDelta>,
699    pub(crate) finish_reason: Option<String>,
700    #[serde(default)]
701    pub(crate) logprobs: Option<OpenAiChoiceLogprobs>,
702}
703
704#[derive(Debug, Deserialize)]
705pub(crate) struct OpenAiDelta {
706    // glm5.1 (and other GLM/zhipu reasoning models) stream/return reasoning under
707    // `reasoning`, not the `reasoning_content` kimi/deepseek use. Without this alias the
708    // reasoning phase yields zero recognized deltas → no ReasoningDelta events → the
709    // stream-stall watchdog kills long reasoning mid-flight (asset-diagnose "未返回结构化输出").
710    #[serde(alias = "reasoning")]
711    pub(crate) reasoning_content: Option<String>,
712    pub(crate) content: Option<String>,
713    pub(crate) tool_calls: Option<Vec<OpenAiToolCallDelta>>,
714}
715
716#[derive(Debug, Deserialize)]
717pub(crate) struct OpenAiToolCallDelta {
718    pub(crate) index: usize,
719    pub(crate) id: Option<String>,
720    pub(crate) function: Option<OpenAiFunctionDelta>,
721}
722
723#[derive(Debug, Deserialize)]
724pub(crate) struct OpenAiFunctionDelta {
725    pub(crate) name: Option<String>,
726    pub(crate) arguments: Option<String>,
727}
728
729// ============================================================================
730// Tests
731// ============================================================================
732
733#[cfg(test)]
734#[path = "openai/tests.rs"]
735mod tests;
a3s_code_core/llm/openai.rs

a3s_code_core/llm/
openai.rs
