1use super::http::{default_http_client, join_chat_completions_url, normalize_base_url, HttpClient};
4use super::structured;
5use super::types::*;
6use super::{LlmClient, ModelGenerationPool};
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
18pub 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.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_value = match content {
203 ToolResultContentField::Text(s) => serde_json::json!(s),
204 ToolResultContentField::Blocks(blocks) => {
205 let mut parts = Vec::new();
206 let mut has_image = false;
207 for block in blocks {
208 match block {
209 ToolResultContent::Text { text } => {
210 parts.push(serde_json::json!({
211 "type": "text",
212 "text": text,
213 }));
214 }
215 ToolResultContent::Image { source } => {
216 has_image = true;
217 parts.push(serde_json::json!({
218 "type": "image_url",
219 "image_url": {
220 "url": format!(
221 "data:{};base64,{}",
222 source.media_type, source.data
223 ),
224 }
225 }));
226 }
227 }
228 }
229 if has_image {
230 serde_json::json!(parts)
231 } else {
232 let text = parts
233 .iter()
234 .filter_map(|part| part.get("text")?.as_str())
235 .collect::<Vec<_>>()
236 .join("\n");
237 serde_json::json!(text)
238 }
239 }
240 };
241 return serde_json::json!({
242 "role": "tool",
243 "tool_call_id": tool_use_id,
244 "content": content_value,
245 });
246 }
247 _ => serde_json::json!(""),
248 }
249 } else {
250 serde_json::json!(msg
251 .content
252 .iter()
253 .map(|block| {
254 match block {
255 ContentBlock::Text { text } => serde_json::json!({
256 "type": "text",
257 "text": text,
258 }),
259 ContentBlock::Image { source } => serde_json::json!({
260 "type": "image_url",
261 "image_url": {
262 "url": format!(
263 "data:{};base64,{}",
264 source.media_type, source.data
265 ),
266 }
267 }),
268 ContentBlock::ToolUse { id, name, input } => serde_json::json!({
269 "type": "function",
270 "id": id,
271 "function": {
272 "name": name,
273 "arguments": input.to_string(),
274 }
275 }),
276 _ => serde_json::json!({}),
277 }
278 })
279 .collect::<Vec<_>>())
280 };
281
282 if msg.role == "assistant" {
285 let rc = msg.reasoning_content.as_deref().unwrap_or("");
286 let tool_calls: Vec<_> = msg.tool_calls();
287 if !tool_calls.is_empty() {
288 return serde_json::json!({
289 "role": "assistant",
290 "content": msg.text(),
291 "reasoning_content": rc,
292 "tool_calls": tool_calls.iter().map(|tc| {
293 serde_json::json!({
294 "id": tc.id,
295 "type": "function",
296 "function": {
297 "name": tc.name,
298 "arguments": tc.args.to_string(),
299 }
300 })
301 }).collect::<Vec<_>>(),
302 });
303 }
304 return serde_json::json!({
305 "role": "assistant",
306 "content": content,
307 "reasoning_content": rc,
308 });
309 }
310
311 serde_json::json!({
312 "role": msg.role,
313 "content": content,
314 })
315 })
316 .collect()
317 }
318
319 pub(crate) fn convert_tools(&self, tools: &[ToolDefinition]) -> Vec<serde_json::Value> {
320 tools
321 .iter()
322 .map(|t| {
323 serde_json::json!({
324 "type": "function",
325 "function": {
326 "name": t.name,
327 "description": t.description,
328 "parameters": t.parameters,
329 }
330 })
331 })
332 .collect()
333 }
334}
335
336impl OpenAiClient {
337 fn apply_directive(
342 request: &mut serde_json::Value,
343 directive: &structured::StructuredDirective,
344 ) {
345 if let Some(tool) = &directive.force_tool {
346 request["tool_choice"] = serde_json::json!({
347 "type": "function",
348 "function": { "name": tool }
349 });
350 }
351 if let Some(rf) = &directive.response_format {
352 request["response_format"] = match rf {
353 structured::ResponseFormat::JsonObject => {
354 serde_json::json!({ "type": "json_object" })
355 }
356 structured::ResponseFormat::JsonSchema { name, schema } => serde_json::json!({
357 "type": "json_schema",
358 "json_schema": { "name": name, "schema": schema, "strict": true }
359 }),
360 };
361 }
362 }
363
364 fn build_chat_request(
366 &self,
367 messages: &[Message],
368 system: Option<&str>,
369 tools: &[ToolDefinition],
370 directive: Option<&structured::StructuredDirective>,
371 ) -> serde_json::Value {
372 let mut openai_messages = Vec::new();
373
374 if let Some(sys) = system {
375 openai_messages.push(serde_json::json!({
376 "role": "system",
377 "content": sys,
378 }));
379 }
380
381 openai_messages.extend(self.convert_messages(messages));
382
383 let mut request = serde_json::json!({
384 "model": self.model,
385 "messages": openai_messages,
386 });
387
388 if let Some(temp) = self.temperature {
389 request["temperature"] = serde_json::json!(temp);
390 }
391 if let Some(max) = self.max_tokens {
392 request["max_tokens"] = serde_json::json!(max);
393 }
394 if self.logprobs {
395 request["logprobs"] = serde_json::json!(true);
396 if let Some(top_logprobs) = self.top_logprobs {
397 request["top_logprobs"] = serde_json::json!(top_logprobs);
398 }
399 }
400
401 if !tools.is_empty() {
402 request["tools"] = serde_json::json!(self.convert_tools(tools));
403 }
404
405 if let Some(directive) = directive {
406 Self::apply_directive(&mut request, directive);
407 }
408
409 request
410 }
411
412 async fn send_request(&self, request: serde_json::Value) -> Result<LlmResponse> {
414 {
415 let request_started_at = Instant::now();
416 let url = join_chat_completions_url(&self.base_url, &self.chat_completions_path);
417 let request_headers = self.request_headers();
418
419 let response = crate::retry::with_retry(&self.retry_config, |_attempt| {
420 let http = &self.http;
421 let url = &url;
422 let request_headers = request_headers.clone();
423 let request = &request;
424 async move {
425 let headers = request_headers
426 .iter()
427 .map(|(key, value)| (key.as_str(), value.as_str()))
428 .collect::<Vec<_>>();
429 let cancel_token = tokio_util::sync::CancellationToken::new();
431 match http.post(url, headers, request, cancel_token).await {
432 Ok(resp) => {
433 let status = reqwest::StatusCode::from_u16(resp.status)
434 .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR);
435 if status.is_success() {
436 AttemptOutcome::Success(resp.body)
437 } else if self.retry_config.is_retryable_status(status) {
438 AttemptOutcome::Retryable {
439 status,
440 body: resp.body,
441 retry_after: None,
442 }
443 } else {
444 AttemptOutcome::Fatal(anyhow::Error::new(
445 crate::llm::NonRetryableLlmError::from_status(
446 &self.provider_name,
447 status.as_u16(),
448 format!("at {url}: {}", resp.body),
449 ),
450 ))
451 }
452 }
453 Err(e) => {
454 tracing::error!("HTTP error: {e:?}");
455 if crate::llm::http::is_retryable_http_failure(&e) {
456 AttemptOutcome::Retryable {
457 status: reqwest::StatusCode::SERVICE_UNAVAILABLE,
458 body: format!("network error: {e}"),
459 retry_after: None,
460 }
461 } else {
462 AttemptOutcome::Fatal(e)
463 }
464 }
465 }
466 }
467 })
468 .await?;
469
470 let parsed: OpenAiResponse =
471 serde_json::from_str(&response).context("Failed to parse OpenAI response")?;
472
473 let choice = parsed.choices.into_iter().next().context("No choices")?;
474 let token_logprobs = choice
475 .logprobs
476 .as_ref()
477 .map(openai_logprobs_to_token_logprobs)
478 .unwrap_or_default();
479
480 let mut content = vec![];
481
482 let reasoning_content = choice.message.reasoning_content;
483
484 let text_content = choice.message.content;
485
486 if let Some(text) = text_content {
487 if !text.is_empty() {
488 content.push(ContentBlock::Text { text });
489 }
490 }
491
492 if let Some(tool_calls) = choice.message.tool_calls {
493 for (index, tc) in tool_calls.into_iter().enumerate() {
494 let id = if tc.id.trim().is_empty() {
495 format!("call_{index}")
496 } else {
497 tc.id
498 };
499 content.push(ContentBlock::ToolUse {
500 id,
501 name: tc.function.name.clone(),
502 input: Self::parse_tool_arguments(
503 &tc.function.name,
504 &tc.function.arguments,
505 ),
506 });
507 }
508 }
509
510 let llm_response = LlmResponse {
511 message: Message {
512 role: "assistant".to_string(),
513 content,
514 reasoning_content,
515 transcript_text: None,
516 transcript_visibility: Default::default(),
517 },
518 usage: TokenUsage {
519 prompt_tokens: parsed.usage.prompt_tokens,
520 completion_tokens: parsed.usage.completion_tokens,
521 total_tokens: {
522 let t = parsed.usage.total_tokens;
523 if t == 0 {
525 parsed.usage.total_characters.unwrap_or(0)
526 } else {
527 t
528 }
529 },
530 cache_read_tokens: parsed
531 .usage
532 .prompt_tokens_details
533 .as_ref()
534 .and_then(|d| d.cached_tokens),
535 cache_write_tokens: None,
536 },
537 stop_reason: choice.finish_reason,
538 token_logprobs,
539 meta: Some(LlmResponseMeta {
540 provider: Some(self.provider_name.clone()),
541 request_model: Some(self.model.clone()),
542 request_url: Some(url.clone()),
543 response_id: parsed.id,
544 response_model: parsed.model,
545 response_object: parsed.object,
546 first_token_ms: None,
547 duration_ms: Some(request_started_at.elapsed().as_millis() as u64),
548 }),
549 };
550
551 crate::telemetry::record_llm_usage(
552 llm_response.usage.prompt_tokens,
553 llm_response.usage.completion_tokens,
554 llm_response.usage.total_tokens,
555 llm_response.stop_reason.as_deref(),
556 );
557
558 Ok(llm_response)
559 }
560 }
561}
562
563#[async_trait]
564impl LlmClient for OpenAiClient {
565 fn model_generation_pool(&self) -> Option<ModelGenerationPool> {
566 ModelGenerationPool::for_endpoint(
567 &self.provider_name,
568 &self.model,
569 &self.base_url,
570 self.model_generation_concurrency(),
571 )
572 .ok()
573 }
574
575 async fn complete(
576 &self,
577 messages: &[Message],
578 system: Option<&str>,
579 tools: &[ToolDefinition],
580 ) -> Result<LlmResponse> {
581 self.send_request(self.build_chat_request(messages, system, tools, None))
582 .await
583 }
584
585 async fn complete_structured(
586 &self,
587 messages: &[Message],
588 system: Option<&str>,
589 tools: &[ToolDefinition],
590 directive: &structured::StructuredDirective,
591 ) -> Result<LlmResponse> {
592 self.send_request(self.build_chat_request(messages, system, tools, Some(directive)))
593 .await
594 }
595
596 fn native_structured_support(&self) -> structured::NativeStructuredSupport {
597 self.native_structured_support
598 }
599
600 fn has_distinct_non_streaming_transport(&self) -> bool {
601 true
602 }
603
604 async fn complete_streaming(
605 &self,
606 messages: &[Message],
607 system: Option<&str>,
608 tools: &[ToolDefinition],
609 cancel_token: tokio_util::sync::CancellationToken,
610 ) -> Result<mpsc::Receiver<StreamEvent>> {
611 self.send_streaming(
612 self.build_chat_request(messages, system, tools, None),
613 cancel_token,
614 )
615 .await
616 }
617
618 async fn complete_streaming_structured(
619 &self,
620 messages: &[Message],
621 system: Option<&str>,
622 tools: &[ToolDefinition],
623 directive: &structured::StructuredDirective,
624 cancel_token: tokio_util::sync::CancellationToken,
625 ) -> Result<mpsc::Receiver<StreamEvent>> {
626 self.send_streaming(
627 self.build_chat_request(messages, system, tools, Some(directive)),
628 cancel_token,
629 )
630 .await
631 }
632}
633
634#[path = "openai/streaming.rs"]
635mod streaming;
636use streaming::openai_logprobs_to_token_logprobs;
637
638#[derive(Debug, Deserialize)]
640pub(crate) struct OpenAiResponse {
641 #[serde(default)]
642 pub(crate) id: Option<String>,
643 #[serde(default)]
644 pub(crate) object: Option<String>,
645 #[serde(default)]
646 pub(crate) model: Option<String>,
647 pub(crate) choices: Vec<OpenAiChoice>,
648 pub(crate) usage: OpenAiUsage,
649}
650
651#[derive(Debug, Deserialize)]
652pub(crate) struct OpenAiChoice {
653 pub(crate) message: OpenAiMessage,
654 pub(crate) finish_reason: Option<String>,
655 #[serde(default)]
656 pub(crate) logprobs: Option<OpenAiChoiceLogprobs>,
657}
658
659#[derive(Debug, Deserialize)]
660pub(crate) struct OpenAiChoiceLogprobs {
661 #[serde(default)]
662 pub(crate) content: Option<Vec<OpenAiTokenLogprob>>,
663}
664
665#[derive(Debug, Deserialize)]
666pub(crate) struct OpenAiTokenLogprob {
667 pub(crate) token: String,
668 pub(crate) logprob: f64,
669 #[serde(default)]
670 pub(crate) bytes: Option<Vec<u8>>,
671 #[serde(default)]
672 pub(crate) top_logprobs: Vec<OpenAiTopLogprob>,
673}
674
675#[derive(Debug, Deserialize)]
676pub(crate) struct OpenAiTopLogprob {
677 pub(crate) token: String,
678 pub(crate) logprob: f64,
679 #[serde(default)]
680 pub(crate) bytes: Option<Vec<u8>>,
681}
682
683#[derive(Debug, Deserialize)]
684pub(crate) struct OpenAiMessage {
685 #[serde(alias = "reasoning")]
690 pub(crate) reasoning_content: Option<String>,
691 pub(crate) content: Option<String>,
692 pub(crate) tool_calls: Option<Vec<OpenAiToolCall>>,
693}
694
695#[derive(Debug, Deserialize)]
696pub(crate) struct OpenAiToolCall {
697 pub(crate) id: String,
698 pub(crate) function: OpenAiFunction,
699}
700
701#[derive(Debug, Deserialize)]
702pub(crate) struct OpenAiFunction {
703 pub(crate) name: String,
704 pub(crate) arguments: String,
705}
706
707#[derive(Debug, Deserialize)]
708pub(crate) struct OpenAiUsage {
709 #[serde(default)]
710 pub(crate) prompt_tokens: usize,
711 #[serde(default)]
712 pub(crate) completion_tokens: usize,
713 #[serde(default)]
714 pub(crate) total_tokens: usize,
715 #[serde(default)]
717 pub(crate) total_characters: Option<usize>,
718 #[serde(default)]
720 pub(crate) prompt_tokens_details: Option<OpenAiPromptTokensDetails>,
721}
722
723#[derive(Debug, Deserialize)]
724pub(crate) struct OpenAiPromptTokensDetails {
725 #[serde(default)]
726 pub(crate) cached_tokens: Option<usize>,
727}
728
729#[derive(Debug, Deserialize)]
731pub(crate) struct OpenAiStreamChunk {
732 #[serde(default)]
733 pub(crate) id: Option<String>,
734 #[serde(default)]
735 pub(crate) object: Option<String>,
736 #[serde(default)]
737 pub(crate) model: Option<String>,
738 pub(crate) choices: Vec<OpenAiStreamChoice>,
739 pub(crate) usage: Option<OpenAiUsage>,
740}
741
742#[derive(Debug, Deserialize)]
743pub(crate) struct OpenAiStreamChoice {
744 pub(crate) message: Option<OpenAiMessage>,
745 pub(crate) delta: Option<OpenAiDelta>,
746 pub(crate) finish_reason: Option<String>,
747 #[serde(default)]
748 pub(crate) logprobs: Option<OpenAiChoiceLogprobs>,
749}
750
751#[derive(Debug, Deserialize)]
752pub(crate) struct OpenAiDelta {
753 #[serde(alias = "reasoning")]
758 pub(crate) reasoning_content: Option<String>,
759 pub(crate) content: Option<String>,
760 pub(crate) tool_calls: Option<Vec<OpenAiToolCallDelta>>,
761}
762
763#[derive(Debug, Deserialize)]
764pub(crate) struct OpenAiToolCallDelta {
765 pub(crate) index: usize,
766 pub(crate) id: Option<String>,
767 pub(crate) function: Option<OpenAiFunctionDelta>,
768}
769
770#[derive(Debug, Deserialize)]
771pub(crate) struct OpenAiFunctionDelta {
772 pub(crate) name: Option<String>,
773 pub(crate) arguments: Option<String>,
774}
775
776#[cfg(test)]
781#[path = "openai/tests.rs"]
782mod tests;