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                            AttemptOutcome::Fatal(e)
430                        }
431                    }
432                }
433            })
434            .await?;
435
436            let parsed: OpenAiResponse =
437                serde_json::from_str(&response).context("Failed to parse OpenAI response")?;
438
439            let choice = parsed.choices.into_iter().next().context("No choices")?;
440            let token_logprobs = choice
441                .logprobs
442                .as_ref()
443                .map(openai_logprobs_to_token_logprobs)
444                .unwrap_or_default();
445
446            let mut content = vec![];
447
448            let reasoning_content = choice.message.reasoning_content;
449
450            let text_content = choice.message.content;
451
452            if let Some(text) = text_content {
453                if !text.is_empty() {
454                    content.push(ContentBlock::Text { text });
455                }
456            }
457
458            if let Some(tool_calls) = choice.message.tool_calls {
459                for tc in tool_calls {
460                    content.push(ContentBlock::ToolUse {
461                        id: tc.id,
462                        name: tc.function.name.clone(),
463                        input: Self::parse_tool_arguments(
464                            &tc.function.name,
465                            &tc.function.arguments,
466                        ),
467                    });
468                }
469            }
470
471            let llm_response = LlmResponse {
472                message: Message {
473                    role: "assistant".to_string(),
474                    content,
475                    reasoning_content,
476                },
477                usage: TokenUsage {
478                    prompt_tokens: parsed.usage.prompt_tokens,
479                    completion_tokens: parsed.usage.completion_tokens,
480                    total_tokens: {
481                        let t = parsed.usage.total_tokens;
482                        // MiniMax: fall back to total_characters when total_tokens is 0.
483                        if t == 0 {
484                            parsed.usage.total_characters.unwrap_or(0)
485                        } else {
486                            t
487                        }
488                    },
489                    cache_read_tokens: parsed
490                        .usage
491                        .prompt_tokens_details
492                        .as_ref()
493                        .and_then(|d| d.cached_tokens),
494                    cache_write_tokens: None,
495                },
496                stop_reason: choice.finish_reason,
497                token_logprobs,
498                meta: Some(LlmResponseMeta {
499                    provider: Some(self.provider_name.clone()),
500                    request_model: Some(self.model.clone()),
501                    request_url: Some(url.clone()),
502                    response_id: parsed.id,
503                    response_model: parsed.model,
504                    response_object: parsed.object,
505                    first_token_ms: None,
506                    duration_ms: Some(request_started_at.elapsed().as_millis() as u64),
507                }),
508            };
509
510            crate::telemetry::record_llm_usage(
511                llm_response.usage.prompt_tokens,
512                llm_response.usage.completion_tokens,
513                llm_response.usage.total_tokens,
514                llm_response.stop_reason.as_deref(),
515            );
516
517            Ok(llm_response)
518        }
519    }
520}
521
522#[async_trait]
523impl LlmClient for OpenAiClient {
524    async fn complete(
525        &self,
526        messages: &[Message],
527        system: Option<&str>,
528        tools: &[ToolDefinition],
529    ) -> Result<LlmResponse> {
530        self.send_request(self.build_chat_request(messages, system, tools, None))
531            .await
532    }
533
534    async fn complete_structured(
535        &self,
536        messages: &[Message],
537        system: Option<&str>,
538        tools: &[ToolDefinition],
539        directive: &structured::StructuredDirective,
540    ) -> Result<LlmResponse> {
541        self.send_request(self.build_chat_request(messages, system, tools, Some(directive)))
542            .await
543    }
544
545    fn native_structured_support(&self) -> structured::NativeStructuredSupport {
546        self.native_structured_support
547    }
548
549    async fn complete_streaming(
550        &self,
551        messages: &[Message],
552        system: Option<&str>,
553        tools: &[ToolDefinition],
554        cancel_token: tokio_util::sync::CancellationToken,
555    ) -> Result<mpsc::Receiver<StreamEvent>> {
556        self.send_streaming(
557            self.build_chat_request(messages, system, tools, None),
558            cancel_token,
559        )
560        .await
561    }
562
563    async fn complete_streaming_structured(
564        &self,
565        messages: &[Message],
566        system: Option<&str>,
567        tools: &[ToolDefinition],
568        directive: &structured::StructuredDirective,
569        cancel_token: tokio_util::sync::CancellationToken,
570    ) -> Result<mpsc::Receiver<StreamEvent>> {
571        self.send_streaming(
572            self.build_chat_request(messages, system, tools, Some(directive)),
573            cancel_token,
574        )
575        .await
576    }
577}
578
579#[path = "openai/streaming.rs"]
580mod streaming;
581use streaming::openai_logprobs_to_token_logprobs;
582
583// OpenAI API response types (private)
584#[derive(Debug, Deserialize)]
585pub(crate) struct OpenAiResponse {
586    #[serde(default)]
587    pub(crate) id: Option<String>,
588    #[serde(default)]
589    pub(crate) object: Option<String>,
590    #[serde(default)]
591    pub(crate) model: Option<String>,
592    pub(crate) choices: Vec<OpenAiChoice>,
593    pub(crate) usage: OpenAiUsage,
594}
595
596#[derive(Debug, Deserialize)]
597pub(crate) struct OpenAiChoice {
598    pub(crate) message: OpenAiMessage,
599    pub(crate) finish_reason: Option<String>,
600    #[serde(default)]
601    pub(crate) logprobs: Option<OpenAiChoiceLogprobs>,
602}
603
604#[derive(Debug, Deserialize)]
605pub(crate) struct OpenAiChoiceLogprobs {
606    #[serde(default)]
607    pub(crate) content: Option<Vec<OpenAiTokenLogprob>>,
608}
609
610#[derive(Debug, Deserialize)]
611pub(crate) struct OpenAiTokenLogprob {
612    pub(crate) token: String,
613    pub(crate) logprob: f64,
614    #[serde(default)]
615    pub(crate) bytes: Option<Vec<u8>>,
616    #[serde(default)]
617    pub(crate) top_logprobs: Vec<OpenAiTopLogprob>,
618}
619
620#[derive(Debug, Deserialize)]
621pub(crate) struct OpenAiTopLogprob {
622    pub(crate) token: String,
623    pub(crate) logprob: f64,
624    #[serde(default)]
625    pub(crate) bytes: Option<Vec<u8>>,
626}
627
628#[derive(Debug, Deserialize)]
629pub(crate) struct OpenAiMessage {
630    // glm5.1 (and other GLM/zhipu reasoning models) stream/return reasoning under
631    // `reasoning`, not the `reasoning_content` kimi/deepseek use. Without this alias the
632    // reasoning phase yields zero recognized deltas → no ReasoningDelta events → the
633    // stream-stall watchdog kills long reasoning mid-flight (asset-diagnose "未返回结构化输出").
634    #[serde(alias = "reasoning")]
635    pub(crate) reasoning_content: Option<String>,
636    pub(crate) content: Option<String>,
637    pub(crate) tool_calls: Option<Vec<OpenAiToolCall>>,
638}
639
640#[derive(Debug, Deserialize)]
641pub(crate) struct OpenAiToolCall {
642    pub(crate) id: String,
643    pub(crate) function: OpenAiFunction,
644}
645
646#[derive(Debug, Deserialize)]
647pub(crate) struct OpenAiFunction {
648    pub(crate) name: String,
649    pub(crate) arguments: String,
650}
651
652#[derive(Debug, Deserialize)]
653pub(crate) struct OpenAiUsage {
654    #[serde(default)]
655    pub(crate) prompt_tokens: usize,
656    #[serde(default)]
657    pub(crate) completion_tokens: usize,
658    #[serde(default)]
659    pub(crate) total_tokens: usize,
660    /// MiniMax uses `total_characters` instead of token counts.
661    #[serde(default)]
662    pub(crate) total_characters: Option<usize>,
663    /// OpenAI returns cached token count in `prompt_tokens_details.cached_tokens`
664    #[serde(default)]
665    pub(crate) prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
666}
667
668#[derive(Debug, Deserialize)]
669pub(crate) struct OpenAiPromptTokensDetails {
670    #[serde(default)]
671    pub(crate) cached_tokens: Option<usize>,
672}
673
674// OpenAI streaming types
675#[derive(Debug, Deserialize)]
676pub(crate) struct OpenAiStreamChunk {
677    #[serde(default)]
678    pub(crate) id: Option<String>,
679    #[serde(default)]
680    pub(crate) object: Option<String>,
681    #[serde(default)]
682    pub(crate) model: Option<String>,
683    pub(crate) choices: Vec<OpenAiStreamChoice>,
684    pub(crate) usage: Option<OpenAiUsage>,
685}
686
687#[derive(Debug, Deserialize)]
688pub(crate) struct OpenAiStreamChoice {
689    pub(crate) message: Option<OpenAiMessage>,
690    pub(crate) delta: Option<OpenAiDelta>,
691    pub(crate) finish_reason: Option<String>,
692    #[serde(default)]
693    pub(crate) logprobs: Option<OpenAiChoiceLogprobs>,
694}
695
696#[derive(Debug, Deserialize)]
697pub(crate) struct OpenAiDelta {
698    // glm5.1 (and other GLM/zhipu reasoning models) stream/return reasoning under
699    // `reasoning`, not the `reasoning_content` kimi/deepseek use. Without this alias the
700    // reasoning phase yields zero recognized deltas → no ReasoningDelta events → the
701    // stream-stall watchdog kills long reasoning mid-flight (asset-diagnose "未返回结构化输出").
702    #[serde(alias = "reasoning")]
703    pub(crate) reasoning_content: Option<String>,
704    pub(crate) content: Option<String>,
705    pub(crate) tool_calls: Option<Vec<OpenAiToolCallDelta>>,
706}
707
708#[derive(Debug, Deserialize)]
709pub(crate) struct OpenAiToolCallDelta {
710    pub(crate) index: usize,
711    pub(crate) id: Option<String>,
712    pub(crate) function: Option<OpenAiFunctionDelta>,
713}
714
715#[derive(Debug, Deserialize)]
716pub(crate) struct OpenAiFunctionDelta {
717    pub(crate) name: Option<String>,
718    pub(crate) arguments: Option<String>,
719}
720
721// ============================================================================
722// Tests
723// ============================================================================
724
725#[cfg(test)]
726#[path = "openai/tests.rs"]
727mod tests;
a3s_code_core/llm/openai.rs

a3s_code_core/llm/
openai.rs
