1#[path = "openai_compat_sse.rs"]
2mod openai_compat_sse;
3#[path = "openai_compat_stream.rs"]
4mod openai_compat_stream;
5
6use std::collections::VecDeque;
7
8use serde::{Deserialize, Deserializer};
9use serde_json::{json, Value};
10
11use crate::error::ApiError;
12use crate::providers::{parse_custom_provider_prefix, RetryPolicy};
13use crate::types::{
14 InputContentBlock, InputMessage, MessageRequest, MessageResponse, OutputContentBlock,
15 ToolChoice, ToolDefinition, ToolResultContentBlock, Usage,
16};
17
18use openai_compat_sse::{first_non_empty_field, OpenAiSseParser};
19use openai_compat_stream::StreamState;
20
21pub use openai_compat_stream::MessageStream;
22
23pub const DEFAULT_XAI_BASE_URL: &str = "https://api.x.ai/v1";
24pub const DEFAULT_OPENAI_BASE_URL: &str = "https://api.openai.com/v1";
25const REQUEST_ID_HEADER: &str = "request-id";
26const ALT_REQUEST_ID_HEADER: &str = "x-request-id";
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct OpenAiCompatConfig {
30 pub provider_name: &'static str,
31 pub api_key_env: &'static str,
32 pub base_url_env: &'static str,
33 pub default_base_url: &'static str,
34}
35
36const XAI_ENV_VARS: &[&str] = &["XAI_API_KEY"];
37const OPENAI_ENV_VARS: &[&str] = &["OPENAI_API_KEY"];
38
39impl OpenAiCompatConfig {
40 #[must_use]
41 pub const fn xai() -> Self {
42 Self {
43 provider_name: "xAI",
44 api_key_env: "XAI_API_KEY",
45 base_url_env: "XAI_BASE_URL",
46 default_base_url: DEFAULT_XAI_BASE_URL,
47 }
48 }
49
50 #[must_use]
51 pub const fn openai() -> Self {
52 Self {
53 provider_name: "OpenAI",
54 api_key_env: "OPENAI_API_KEY",
55 base_url_env: "OPENAI_BASE_URL",
56 default_base_url: DEFAULT_OPENAI_BASE_URL,
57 }
58 }
59 #[must_use]
60 pub fn credential_env_vars(self) -> &'static [&'static str] {
61 match self.api_key_env {
62 "XAI_API_KEY" => XAI_ENV_VARS,
63 "OPENAI_API_KEY" => OPENAI_ENV_VARS,
64 _ => &[],
65 }
66 }
67}
68
69#[derive(Clone)]
70pub struct OpenAiCompatClient {
71 http: reqwest::Client,
72 api_key: String,
73 base_url: String,
74 endpoint_query: Option<String>,
75 retry: RetryPolicy,
76}
77
78impl std::fmt::Debug for OpenAiCompatClient {
79 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80 f.debug_struct("OpenAiCompatClient")
81 .field("base_url", &self.base_url)
82 .field("endpoint_query", &self.endpoint_query)
83 .field("api_key", &"***")
84 .finish()
85 }
86}
87
88impl OpenAiCompatClient {
89 #[must_use]
90 pub fn new(api_key: impl Into<String>, config: OpenAiCompatConfig) -> Self {
91 Self {
92 http: crate::default_http_client(),
93 api_key: api_key.into(),
94 base_url: read_base_url(config),
95 endpoint_query: None,
96 retry: RetryPolicy::default(),
97 }
98 }
99
100 #[must_use]
101 pub fn new_custom(base_url: impl Into<String>, api_key: impl Into<String>) -> Self {
102 Self {
103 http: crate::default_http_client(),
104 api_key: api_key.into(),
105 base_url: base_url.into(),
106 endpoint_query: None,
107 retry: RetryPolicy::default(),
108 }
109 }
110
111 #[must_use]
112 pub fn with_endpoint_query(mut self, endpoint_query: Option<String>) -> Self {
113 self.endpoint_query = endpoint_query
114 .map(|s| s.trim().to_string())
115 .filter(|s| !s.is_empty());
116 self
117 }
118
119 pub fn from_env(config: OpenAiCompatConfig) -> Result<Self, ApiError> {
120 let Some(api_key) = read_env_non_empty(config.api_key_env)? else {
121 return Err(ApiError::missing_credentials(
122 config.provider_name,
123 config.credential_env_vars(),
124 ));
125 };
126 Ok(Self::new(api_key, config))
127 }
128
129 #[must_use]
130 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
131 self.base_url = base_url.into();
132 self
133 }
134
135 #[must_use]
136 pub fn with_retry_policy(mut self, retry: RetryPolicy) -> Self {
137 self.retry = retry;
138 self
139 }
140
141 pub async fn send_message(
142 &self,
143 request: &MessageRequest,
144 ) -> Result<MessageResponse, ApiError> {
145 let request = MessageRequest {
146 stream: false,
147 ..request.clone()
148 };
149 let response = self.send_with_retry(&request).await?;
150 let request_id = request_id_from_headers(response.headers());
151 let payload = response.json::<ChatCompletionResponse>().await?;
152 let mut normalized = normalize_response(&request.model, payload)?;
153 if normalized.request_id.is_none() {
154 normalized.request_id = request_id;
155 }
156 Ok(normalized)
157 }
158
159 pub async fn stream_message(
160 &self,
161 request: &MessageRequest,
162 ) -> Result<MessageStream, ApiError> {
163 let response = self
164 .send_with_retry(&request.clone().with_streaming())
165 .await?;
166 Ok(MessageStream {
167 request_id: request_id_from_headers(response.headers()),
168 response,
169 parser: OpenAiSseParser::new(),
170 pending: VecDeque::new(),
171 done: false,
172 state: StreamState::new(request.model.clone()),
173 })
174 }
175
176 async fn send_with_retry(
177 &self,
178 request: &MessageRequest,
179 ) -> Result<reqwest::Response, ApiError> {
180 let mut attempts = 0;
181
182 let last_error = loop {
183 attempts += 1;
184 let retryable_error = match self.send_raw_request(request).await {
185 Ok(response) => match expect_success(response).await {
186 Ok(response) => return Ok(response),
187 Err(error)
188 if error.is_retryable() && attempts <= self.retry.max_retries + 1 =>
189 {
190 error
191 }
192 Err(error) => return Err(error),
193 },
194 Err(error) if error.is_retryable() && attempts <= self.retry.max_retries + 1 => {
195 error
196 }
197 Err(error) => return Err(error),
198 };
199
200 if attempts > self.retry.max_retries {
201 break retryable_error;
202 }
203
204 tokio::time::sleep(self.backoff_for_attempt(attempts)?).await;
205 };
206
207 Err(ApiError::RetriesExhausted {
208 attempts,
209 last_error: Box::new(last_error),
210 })
211 }
212
213 async fn send_raw_request(
214 &self,
215 request: &MessageRequest,
216 ) -> Result<reqwest::Response, ApiError> {
217 let request_url = chat_completions_endpoint(&self.base_url, self.endpoint_query.as_deref());
218 let mut req = self
219 .http
220 .post(&request_url)
221 .header("content-type", "application/json");
222 if !self.api_key.is_empty() {
223 req = req.bearer_auth(&self.api_key);
224 }
225 req.json(&build_chat_completion_request(request))
226 .send()
227 .await
228 .map_err(ApiError::from)
229 }
230
231 fn backoff_for_attempt(&self, attempt: u32) -> Result<std::time::Duration, ApiError> {
232 let Some(multiplier) = 1_u32.checked_shl(attempt.saturating_sub(1)) else {
233 return Err(ApiError::BackoffOverflow {
234 attempt,
235 base_delay: self.retry.initial_backoff,
236 });
237 };
238 Ok(self
239 .retry
240 .initial_backoff
241 .checked_mul(multiplier)
242 .map_or(self.retry.max_backoff, |delay| {
243 delay.min(self.retry.max_backoff)
244 }))
245 }
246}
247
248#[derive(Debug, Deserialize)]
253struct ChatCompletionResponse {
254 id: String,
255 model: String,
256 choices: Vec<ChatChoice>,
257 #[serde(default)]
258 usage: Option<OpenAiUsage>,
259}
260
261#[derive(Debug, Deserialize)]
262struct ChatChoice {
263 message: ChatMessage,
264 #[serde(default)]
265 finish_reason: Option<String>,
266}
267
268#[derive(Debug, Deserialize)]
269struct ChatMessage {
270 role: String,
271 #[serde(default, deserialize_with = "deserialize_openai_text_content")]
272 content: Option<String>,
273 #[serde(default)]
274 reasoning_content: Option<String>,
275 #[serde(default)]
276 reasoning: Option<String>,
277 #[serde(default)]
278 thought: Option<String>,
279 #[serde(default)]
280 thinking: Option<String>,
281 #[serde(default)]
282 tool_calls: Vec<ResponseToolCall>,
283}
284
285impl ChatMessage {
286 fn assistant_visible_text(&self) -> Option<String> {
287 first_non_empty_field(&[
288 &self.content,
289 &self.reasoning_content,
290 &self.reasoning,
291 &self.thought,
292 &self.thinking,
293 ])
294 }
295}
296
297#[derive(Debug, Deserialize)]
298struct ResponseToolCall {
299 id: String,
300 function: ResponseToolFunction,
301}
302
303#[derive(Debug, Deserialize)]
304struct ResponseToolFunction {
305 name: String,
306 arguments: String,
307}
308
309#[derive(Debug, Deserialize)]
310pub(super) struct OpenAiUsage {
311 #[serde(default)]
312 pub prompt_tokens: u32,
313 #[serde(default)]
314 pub completion_tokens: u32,
315}
316
317#[derive(Debug, Deserialize)]
318struct ErrorEnvelope {
319 error: ErrorBody,
320}
321
322#[derive(Debug, Deserialize)]
323struct ErrorBody {
324 #[serde(rename = "type")]
325 error_type: Option<String>,
326 message: Option<String>,
327}
328
329fn upstream_openai_model(model: &str) -> String {
334 parse_custom_provider_prefix(model)
335 .map(|(_, rest)| rest.to_string())
336 .unwrap_or_else(|| model.to_string())
337}
338
339fn build_chat_completion_request(request: &MessageRequest) -> Value {
340 let mut messages = Vec::new();
341 if let Some(system) = request.system.as_ref().filter(|value| !value.is_empty()) {
342 messages.push(json!({
343 "role": "system",
344 "content": system,
345 }));
346 }
347 for message in &request.messages {
348 messages.extend(translate_message(message));
349 }
350
351 let upstream_model = upstream_openai_model(&request.model);
352 const MAX_TOKENS_OPENAI_COMPAT_CAP: u32 = 32_768;
353 let max_tokens = request.max_tokens.clamp(1, MAX_TOKENS_OPENAI_COMPAT_CAP);
354 let mut payload = json!({
355 "model": upstream_model,
356 "max_tokens": max_tokens,
357 "messages": messages,
358 "stream": request.stream,
359 });
360
361 if let Some(tools) = &request.tools {
362 payload["tools"] =
363 Value::Array(tools.iter().map(openai_tool_definition).collect::<Vec<_>>());
364 }
365 if let Some(tool_choice) = &request.tool_choice {
366 payload["tool_choice"] = openai_tool_choice(tool_choice);
367 }
368
369 payload
370}
371
372fn translate_message(message: &InputMessage) -> Vec<Value> {
373 match message.role.as_str() {
374 "assistant" => {
375 let mut text = String::new();
376 let mut tool_calls = Vec::new();
377 for block in &message.content {
378 match block {
379 InputContentBlock::Text { text: value } => text.push_str(value),
380 InputContentBlock::ToolUse { id, name, input } => tool_calls.push(json!({
381 "id": id,
382 "type": "function",
383 "function": {
384 "name": name,
385 "arguments": serde_json::to_string(input).unwrap_or_default(),
386 }
387 })),
388 InputContentBlock::ToolResult { .. } => {}
389 }
390 }
391 if text.is_empty() && tool_calls.is_empty() {
392 Vec::new()
393 } else {
394 vec![json!({
395 "role": "assistant",
396 "content": (!text.is_empty()).then_some(text),
397 "tool_calls": tool_calls,
398 })]
399 }
400 }
401 _ => message
402 .content
403 .iter()
404 .filter_map(|block| match block {
405 InputContentBlock::Text { text } => Some(json!({
406 "role": "user",
407 "content": text,
408 })),
409 InputContentBlock::ToolResult {
410 tool_use_id,
411 content,
412 is_error,
413 } => Some(json!({
414 "role": "tool",
415 "tool_call_id": tool_use_id,
416 "content": flatten_tool_result_content(content),
417 "is_error": is_error,
418 })),
419 InputContentBlock::ToolUse { .. } => None,
420 })
421 .collect(),
422 }
423}
424
425fn flatten_tool_result_content(content: &[ToolResultContentBlock]) -> String {
426 content
427 .iter()
428 .map(|block| match block {
429 ToolResultContentBlock::Text { text } => text.clone(),
430 ToolResultContentBlock::Json { value } => value.to_string(),
431 })
432 .collect::<Vec<_>>()
433 .join("\n")
434}
435
436fn openai_tool_definition(tool: &ToolDefinition) -> Value {
437 json!({
438 "type": "function",
439 "function": {
440 "name": tool.name,
441 "description": tool.description,
442 "parameters": tool.input_schema,
443 }
444 })
445}
446
447fn openai_tool_choice(tool_choice: &ToolChoice) -> Value {
448 match tool_choice {
449 ToolChoice::Auto => Value::String("auto".to_string()),
450 ToolChoice::Any => Value::String("required".to_string()),
451 ToolChoice::Tool { name } => json!({
452 "type": "function",
453 "function": { "name": name },
454 }),
455 }
456}
457
458fn normalize_response(
459 model: &str,
460 response: ChatCompletionResponse,
461) -> Result<MessageResponse, ApiError> {
462 let choice = response
463 .choices
464 .into_iter()
465 .next()
466 .ok_or(ApiError::InvalidSseFrame(
467 "chat completion response missing choices",
468 ))?;
469 let mut content = Vec::new();
470 if let Some(text) = choice.message.assistant_visible_text() {
471 content.push(OutputContentBlock::Text { text });
472 }
473 for tool_call in choice.message.tool_calls {
474 content.push(OutputContentBlock::ToolUse {
475 id: tool_call.id,
476 name: tool_call.function.name,
477 input: parse_tool_arguments(&tool_call.function.arguments),
478 });
479 }
480
481 Ok(MessageResponse {
482 id: response.id,
483 kind: "message".to_string(),
484 role: choice.message.role,
485 content,
486 model: response.model.if_empty_then(model.to_string()),
487 stop_reason: choice
488 .finish_reason
489 .map(|value| normalize_finish_reason(&value)),
490 stop_sequence: None,
491 usage: Usage {
492 input_tokens: response
493 .usage
494 .as_ref()
495 .map_or(0, |usage| usage.prompt_tokens),
496 cache_creation_input_tokens: 0,
497 cache_read_input_tokens: 0,
498 output_tokens: response
499 .usage
500 .as_ref()
501 .map_or(0, |usage| usage.completion_tokens),
502 },
503 request_id: None,
504 })
505}
506
507fn parse_tool_arguments(arguments: &str) -> Value {
508 serde_json::from_str(arguments).unwrap_or_else(|_| json!({ "raw": arguments }))
509}
510
511fn deserialize_openai_text_content<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
517where
518 D: Deserializer<'de>,
519{
520 #[derive(Deserialize)]
521 #[serde(untagged)]
522 enum Raw {
523 Str(String),
524 Arr(Vec<Value>),
525 }
526 match Option::<Raw>::deserialize(deserializer)? {
527 None => Ok(None),
528 Some(Raw::Str(s)) if s.is_empty() => Ok(None),
529 Some(Raw::Str(s)) => Ok(Some(s)),
530 Some(Raw::Arr(parts)) => {
531 let mut joined = String::new();
532 for part in parts {
533 match part {
534 Value::Object(map) => {
535 if let Some(text) = map.get("text").and_then(Value::as_str) {
536 joined.push_str(text);
537 } else if let Some(text) = map.get("content").and_then(Value::as_str) {
538 joined.push_str(text);
539 }
540 }
541 Value::String(s) => joined.push_str(&s),
542 _ => {}
543 }
544 }
545 Ok((!joined.is_empty()).then_some(joined))
546 }
547 }
548}
549
550fn read_env_non_empty(key: &str) -> Result<Option<String>, ApiError> {
555 match std::env::var(key) {
556 Ok(value) if !value.is_empty() => Ok(Some(value)),
557 Ok(_) | Err(std::env::VarError::NotPresent) => Ok(None),
558 Err(error) => Err(ApiError::from(error)),
559 }
560}
561
562#[must_use]
563pub fn has_api_key(key: &str) -> bool {
564 read_env_non_empty(key)
565 .ok()
566 .and_then(std::convert::identity)
567 .is_some()
568}
569
570#[must_use]
571pub fn read_base_url(config: OpenAiCompatConfig) -> String {
572 std::env::var(config.base_url_env).unwrap_or_else(|_| config.default_base_url.to_string())
573}
574
575fn chat_completions_endpoint(base_url: &str, extra_query: Option<&str>) -> String {
576 let trimmed = base_url.trim();
577 let (path_part, base_query) = match trimmed.split_once('?') {
578 Some((p, q)) => (p.trim_end_matches('/'), Some(q)),
579 None => (trimmed.trim_end_matches('/'), None),
580 };
581 let path = if path_part.ends_with("/chat/completions") {
582 path_part.to_string()
583 } else {
584 format!("{path_part}/chat/completions")
585 };
586 merge_url_query(&path, base_query, extra_query)
587}
588
589fn merge_url_query(path: &str, base_query: Option<&str>, extra_query: Option<&str>) -> String {
590 let mut segments: Vec<&str> = Vec::new();
591 if let Some(q) = base_query.map(str::trim).filter(|q| !q.is_empty()) {
592 segments.push(q);
593 }
594 if let Some(q) = extra_query.map(str::trim).filter(|q| !q.is_empty()) {
595 segments.push(q);
596 }
597 if segments.is_empty() {
598 path.to_string()
599 } else {
600 format!("{path}?{}", segments.join("&"))
601 }
602}
603
604fn request_id_from_headers(headers: &reqwest::header::HeaderMap) -> Option<String> {
605 headers
606 .get(REQUEST_ID_HEADER)
607 .or_else(|| headers.get(ALT_REQUEST_ID_HEADER))
608 .and_then(|value| value.to_str().ok())
609 .map(ToOwned::to_owned)
610}
611
612async fn expect_success(response: reqwest::Response) -> Result<reqwest::Response, ApiError> {
613 let status = response.status();
614 if status.is_success() {
615 return Ok(response);
616 }
617
618 let body = response.text().await.unwrap_or_default();
619 let parsed_error = serde_json::from_str::<ErrorEnvelope>(&body).ok();
620 let retryable = is_retryable_status(status);
621
622 Err(ApiError::Api {
623 status,
624 error_type: parsed_error
625 .as_ref()
626 .and_then(|error| error.error.error_type.clone()),
627 message: parsed_error
628 .as_ref()
629 .and_then(|error| error.error.message.clone()),
630 body,
631 retryable,
632 })
633}
634
635const fn is_retryable_status(status: reqwest::StatusCode) -> bool {
636 matches!(status.as_u16(), 408 | 409 | 429 | 500 | 502 | 503 | 504)
637}
638
639fn normalize_finish_reason(value: &str) -> String {
640 match value {
641 "stop" => "end_turn",
642 "tool_calls" => "tool_use",
643 other => other,
644 }
645 .to_string()
646}
647
648trait StringExt {
649 fn if_empty_then(self, fallback: String) -> String;
650}
651
652impl StringExt for String {
653 fn if_empty_then(self, fallback: String) -> String {
654 if self.is_empty() {
655 fallback
656 } else {
657 self
658 }
659 }
660}
661
662#[cfg(test)]
663mod openai_compat_inner_tests {
664 use super::*;
665 use crate::types::OutputContentBlock;
666
667 #[test]
668 fn chat_completions_url_appends_api_version() {
669 assert_eq!(
670 chat_completions_endpoint(
671 "https://my.openai.azure.com/openai/deployments/gpt4",
672 Some("api-version=2024-02-15-preview"),
673 ),
674 "https://my.openai.azure.com/openai/deployments/gpt4/chat/completions?api-version=2024-02-15-preview"
675 );
676 }
677
678 #[test]
679 fn chat_completions_url_merges_base_query_and_api_version() {
680 assert_eq!(
681 chat_completions_endpoint(
682 "https://x/v1/chat/completions?existing=1",
683 Some("api-version=2024-02-15-preview"),
684 ),
685 "https://x/v1/chat/completions?existing=1&api-version=2024-02-15-preview"
686 );
687 }
688
689 #[test]
690 fn non_streaming_message_parses_content_array() {
691 let json = r#"{
692 "id":"1",
693 "model":"qwen",
694 "choices":[{
695 "message":{"role":"assistant","content":[{"type":"text","text":"hello"}]},
696 "finish_reason":"stop"
697 }],
698 "usage":{"prompt_tokens":1,"completion_tokens":1}
699 }"#;
700 let resp: ChatCompletionResponse = serde_json::from_str(json).unwrap();
701 let msg = normalize_response("qwen", resp).expect("normalize");
702 assert_eq!(
703 msg.content,
704 vec![OutputContentBlock::Text {
705 text: "hello".to_string()
706 }]
707 );
708 }
709
710 #[test]
711 fn non_streaming_reasoning_only_message() {
712 let json = r#"{
713 "id":"1",
714 "model":"qwen",
715 "choices":[{
716 "message":{"role":"assistant","content":null,"reasoning_content":"think"},
717 "finish_reason":"stop"
718 }],
719 "usage":{"prompt_tokens":1,"completion_tokens":1}
720 }"#;
721 let resp: ChatCompletionResponse = serde_json::from_str(json).unwrap();
722 let msg = normalize_response("qwen", resp).expect("normalize");
723 assert_eq!(
724 msg.content,
725 vec![OutputContentBlock::Text {
726 text: "think".to_string()
727 }]
728 );
729 }
730}
731
732#[cfg(test)]
733#[path = "openai_compat_tests.rs"]
734mod tests;