tenant.openagents/omega

1pub mod batches;
2pub mod completion;
3pub mod responses;
4
5use anyhow::{Context as _, Result, anyhow};
6use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream};
7use http_client::{
8    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt,
9    StatusCode,
10    http::{HeaderMap, HeaderValue},
11};
12pub use language_model_core::ReasoningEffort;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use std::{convert::TryFrom, future::Future};
16use strum::EnumIter;
17use thiserror::Error;
18
19pub const OPEN_AI_API_URL: &str = "https://api.openai.com/v1";
20
21fn is_none_or_empty<T: AsRef<[U]>, U>(opt: &Option<T>) -> bool {
22    opt.as_ref().is_none_or(|v| v.as_ref().is_empty())
23}
24
25#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
26#[serde(rename_all = "lowercase")]
27pub enum Role {
28    User,
29    Assistant,
30    System,
31    Tool,
32}
33
34impl TryFrom<String> for Role {
35    type Error = anyhow::Error;
36
37    fn try_from(value: String) -> Result<Self> {
38        match value.as_str() {
39            "user" => Ok(Self::User),
40            "assistant" => Ok(Self::Assistant),
41            "system" => Ok(Self::System),
42            "tool" => Ok(Self::Tool),
43            _ => anyhow::bail!("invalid role '{value}'"),
44        }
45    }
46}
47
48impl From<Role> for String {
49    fn from(val: Role) -> Self {
50        match val {
51            Role::User => "user".to_owned(),
52            Role::Assistant => "assistant".to_owned(),
53            Role::System => "system".to_owned(),
54            Role::Tool => "tool".to_owned(),
55        }
56    }
57}
58
59#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
60#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
61pub enum Model {
62    #[serde(rename = "gpt-4")]
63    Four,
64    #[serde(rename = "gpt-4o-mini")]
65    FourOmniMini,
66    #[serde(rename = "o3")]
67    O3,
68    #[serde(rename = "gpt-5")]
69    Five,
70    #[serde(rename = "gpt-5-mini")]
71    FiveMini,
72    #[serde(rename = "gpt-5-nano")]
73    FiveNano,
74    #[serde(rename = "gpt-5.1")]
75    FivePointOne,
76    #[serde(rename = "gpt-5.2")]
77    FivePointTwo,
78    #[serde(rename = "gpt-5.3-codex")]
79    FivePointThreeCodex,
80    #[serde(rename = "gpt-5.4-nano")]
81    FivePointFourNano,
82    #[serde(rename = "gpt-5.4-mini")]
83    FivePointFourMini,
84    #[serde(rename = "gpt-5.4")]
85    FivePointFour,
86    #[serde(rename = "gpt-5.4-pro")]
87    FivePointFourPro,
88    #[serde(rename = "gpt-5.5")]
89    FivePointFive,
90    #[serde(rename = "gpt-5.5-pro")]
91    FivePointFivePro,
92    #[serde(rename = "gpt-5.6-sol")]
93    #[default]
94    FivePointSixSol,
95    #[serde(rename = "gpt-5.6-terra")]
96    FivePointSixTerra,
97    #[serde(rename = "gpt-5.6-luna")]
98    FivePointSixLuna,
99    #[serde(rename = "custom")]
100    Custom {
101        name: String,
102        /// The name displayed in the UI, such as in the agent panel model dropdown menu.
103        display_name: Option<String>,
104        max_tokens: u64,
105        max_output_tokens: Option<u64>,
106        max_completion_tokens: Option<u64>,
107        reasoning_effort: Option<ReasoningEffort>,
108        #[serde(default = "default_supports_chat_completions")]
109        supports_chat_completions: bool,
110        #[serde(default = "default_supports_images")]
111        supports_images: bool,
112    },
113}
114
115const fn default_supports_chat_completions() -> bool {
116    true
117}
118
119const fn default_supports_images() -> bool {
120    true
121}
122
123impl Model {
124    pub fn default_fast() -> Self {
125        Self::FivePointSixLuna
126    }
127
128    pub fn from_id(id: &str) -> Result<Self> {
129        match id {
130            "gpt-4" => Ok(Self::Four),
131            "gpt-4o-mini" => Ok(Self::FourOmniMini),
132            "o3" => Ok(Self::O3),
133            "gpt-5" => Ok(Self::Five),
134            "gpt-5-mini" => Ok(Self::FiveMini),
135            "gpt-5-nano" => Ok(Self::FiveNano),
136            "gpt-5.1" => Ok(Self::FivePointOne),
137            "gpt-5.2" => Ok(Self::FivePointTwo),
138            "gpt-5.3-codex" => Ok(Self::FivePointThreeCodex),
139            "gpt-5.4-nano" => Ok(Self::FivePointFourNano),
140            "gpt-5.4-mini" => Ok(Self::FivePointFourMini),
141            "gpt-5.4" => Ok(Self::FivePointFour),
142            "gpt-5.4-pro" => Ok(Self::FivePointFourPro),
143            "gpt-5.5" => Ok(Self::FivePointFive),
144            "gpt-5.5-pro" => Ok(Self::FivePointFivePro),
145            "gpt-5.6-sol" => Ok(Self::FivePointSixSol),
146            "gpt-5.6-terra" => Ok(Self::FivePointSixTerra),
147            "gpt-5.6-luna" => Ok(Self::FivePointSixLuna),
148            invalid_id => anyhow::bail!("invalid model id '{invalid_id}'"),
149        }
150    }
151
152    pub fn id(&self) -> &str {
153        match self {
154            Self::Four => "gpt-4",
155            Self::FourOmniMini => "gpt-4o-mini",
156            Self::O3 => "o3",
157            Self::Five => "gpt-5",
158            Self::FiveMini => "gpt-5-mini",
159            Self::FiveNano => "gpt-5-nano",
160            Self::FivePointOne => "gpt-5.1",
161            Self::FivePointTwo => "gpt-5.2",
162            Self::FivePointThreeCodex => "gpt-5.3-codex",
163            Self::FivePointFourNano => "gpt-5.4-nano",
164            Self::FivePointFourMini => "gpt-5.4-mini",
165            Self::FivePointFour => "gpt-5.4",
166            Self::FivePointFourPro => "gpt-5.4-pro",
167            Self::FivePointFive => "gpt-5.5",
168            Self::FivePointFivePro => "gpt-5.5-pro",
169            Self::FivePointSixSol => "gpt-5.6-sol",
170            Self::FivePointSixTerra => "gpt-5.6-terra",
171            Self::FivePointSixLuna => "gpt-5.6-luna",
172            Self::Custom { name, .. } => name,
173        }
174    }
175
176    pub fn display_name(&self) -> &str {
177        match self {
178            Self::Four => "gpt-4",
179            Self::FourOmniMini => "gpt-4o-mini",
180            Self::O3 => "o3",
181            Self::Five => "gpt-5",
182            Self::FiveMini => "gpt-5-mini",
183            Self::FiveNano => "gpt-5-nano",
184            Self::FivePointOne => "gpt-5.1",
185            Self::FivePointTwo => "gpt-5.2",
186            Self::FivePointThreeCodex => "gpt-5.3-codex",
187            Self::FivePointFourNano => "gpt-5.4-nano",
188            Self::FivePointFourMini => "gpt-5.4-mini",
189            Self::FivePointFour => "gpt-5.4",
190            Self::FivePointFourPro => "gpt-5.4-pro",
191            Self::FivePointFive => "gpt-5.5",
192            Self::FivePointFivePro => "gpt-5.5-pro",
193            Self::FivePointSixSol => "gpt-5.6-sol",
194            Self::FivePointSixTerra => "gpt-5.6-terra",
195            Self::FivePointSixLuna => "gpt-5.6-luna",
196            Self::Custom { display_name, .. } => display_name.as_deref().unwrap_or(&self.id()),
197        }
198    }
199
200    pub fn max_token_count(&self) -> u64 {
201        match self {
202            Self::Four => 8_192,
203            Self::FourOmniMini => 128_000,
204            Self::O3 => 200_000,
205            Self::Five => 272_000,
206            Self::FiveMini => 400_000,
207            Self::FiveNano => 400_000,
208            Self::FivePointOne => 400_000,
209            Self::FivePointTwo => 400_000,
210            Self::FivePointThreeCodex => 400_000,
211            Self::FivePointFourNano => 400_000,
212            Self::FivePointFourMini => 400_000,
213            Self::FivePointFour => 1_050_000,
214            Self::FivePointFourPro => 1_050_000,
215            Self::FivePointFive => 1_050_000,
216            Self::FivePointFivePro => 1_050_000,
217            Self::FivePointSixSol => 1_050_000,
218            Self::FivePointSixTerra => 1_050_000,
219            Self::FivePointSixLuna => 1_050_000,
220            Self::Custom { max_tokens, .. } => *max_tokens,
221        }
222    }
223
224    pub fn max_output_tokens(&self) -> Option<u64> {
225        match self {
226            Self::Custom {
227                max_output_tokens, ..
228            } => *max_output_tokens,
229            Self::Four => Some(8_192),
230            Self::FourOmniMini => Some(16_384),
231            Self::O3 => Some(100_000),
232            Self::Five => Some(128_000),
233            Self::FiveMini => Some(128_000),
234            Self::FiveNano => Some(128_000),
235            Self::FivePointOne => Some(128_000),
236            Self::FivePointTwo => Some(128_000),
237            Self::FivePointThreeCodex => Some(128_000),
238            Self::FivePointFourNano => Some(128_000),
239            Self::FivePointFourMini => Some(128_000),
240            Self::FivePointFour => Some(128_000),
241            Self::FivePointFourPro => Some(128_000),
242            Self::FivePointFive => Some(128_000),
243            Self::FivePointFivePro => Some(128_000),
244            Self::FivePointSixSol => Some(128_000),
245            Self::FivePointSixTerra => Some(128_000),
246            Self::FivePointSixLuna => Some(128_000),
247        }
248    }
249
250    pub fn reasoning_effort(&self) -> Option<ReasoningEffort> {
251        match self {
252            Self::Custom {
253                reasoning_effort, ..
254            } => reasoning_effort.to_owned(),
255            Self::FivePointOne
256            | Self::FivePointTwo
257            | Self::FivePointFour
258            | Self::FivePointFourMini
259            | Self::FivePointFourNano => Some(ReasoningEffort::None),
260            Self::FivePointSixSol => Some(ReasoningEffort::Low),
261            Self::O3
262            | Self::Five
263            | Self::FiveMini
264            | Self::FiveNano
265            | Self::FivePointThreeCodex
266            | Self::FivePointFourPro
267            | Self::FivePointFive
268            | Self::FivePointFivePro
269            | Self::FivePointSixTerra
270            | Self::FivePointSixLuna => Some(ReasoningEffort::Medium),
271            _ => None,
272        }
273    }
274
275    pub fn supported_reasoning_efforts(&self) -> &'static [ReasoningEffort] {
276        match self {
277            Self::Custom {
278                reasoning_effort: Some(effort),
279                ..
280            } => match effort {
281                ReasoningEffort::None => &[ReasoningEffort::None],
282                ReasoningEffort::Minimal => &[ReasoningEffort::Minimal],
283                ReasoningEffort::Low => &[ReasoningEffort::Low],
284                ReasoningEffort::Medium => &[ReasoningEffort::Medium],
285                ReasoningEffort::High => &[ReasoningEffort::High],
286                ReasoningEffort::XHigh => &[ReasoningEffort::XHigh],
287                ReasoningEffort::Max => &[ReasoningEffort::Max],
288            },
289            Self::O3 => &[
290                ReasoningEffort::Low,
291                ReasoningEffort::Medium,
292                ReasoningEffort::High,
293            ],
294            Self::FivePointOne => &[
295                ReasoningEffort::None,
296                ReasoningEffort::Low,
297                ReasoningEffort::Medium,
298                ReasoningEffort::High,
299            ],
300            Self::Five | Self::FiveMini | Self::FiveNano => &[
301                ReasoningEffort::Minimal,
302                ReasoningEffort::Low,
303                ReasoningEffort::Medium,
304                ReasoningEffort::High,
305            ],
306            Self::FivePointFourPro | Self::FivePointFivePro => &[
307                ReasoningEffort::Medium,
308                ReasoningEffort::High,
309                ReasoningEffort::XHigh,
310            ],
311            Self::FivePointThreeCodex => &[
312                ReasoningEffort::Low,
313                ReasoningEffort::Medium,
314                ReasoningEffort::High,
315                ReasoningEffort::XHigh,
316            ],
317            Self::FivePointSixSol | Self::FivePointSixTerra | Self::FivePointSixLuna => &[
318                ReasoningEffort::None,
319                ReasoningEffort::Low,
320                ReasoningEffort::Medium,
321                ReasoningEffort::High,
322                ReasoningEffort::XHigh,
323                ReasoningEffort::Max,
324            ],
325            Self::FivePointTwo
326            | Self::FivePointFour
327            | Self::FivePointFive
328            | Self::FivePointFourMini
329            | Self::FivePointFourNano => &[
330                ReasoningEffort::None,
331                ReasoningEffort::Low,
332                ReasoningEffort::Medium,
333                ReasoningEffort::High,
334                ReasoningEffort::XHigh,
335            ],
336            _ => &[],
337        }
338    }
339
340    pub fn uses_responses_api(&self) -> bool {
341        match self {
342            Self::Custom {
343                supports_chat_completions,
344                ..
345            } => !*supports_chat_completions,
346            _ => true,
347        }
348    }
349
350    /// Returns whether the given model supports the `parallel_tool_calls` parameter.
351    ///
352    /// If the model does not support the parameter, do not pass it up, or the API will return an error.
353    pub fn supports_parallel_tool_calls(&self) -> bool {
354        match self {
355            Self::Four
356            | Self::FourOmniMini
357            | Self::Five
358            | Self::FiveMini
359            | Self::FivePointOne
360            | Self::FivePointTwo
361            | Self::FivePointThreeCodex
362            | Self::FivePointFour
363            | Self::FivePointFourMini
364            | Self::FivePointFourNano
365            | Self::FivePointFourPro
366            | Self::FivePointFive
367            | Self::FivePointFivePro
368            | Self::FivePointSixSol
369            | Self::FivePointSixTerra
370            | Self::FivePointSixLuna
371            | Self::FiveNano => true,
372            Self::O3 | Model::Custom { .. } => false,
373        }
374    }
375
376    /// Returns whether the given model supports the `prompt_cache_key` parameter.
377    ///
378    /// If the model does not support the parameter, do not pass it up.
379    pub fn supports_prompt_cache_key(&self) -> bool {
380        true
381    }
382
383    /// Whether this model supports server-side compaction via the
384    /// `context_management` request parameter. OpenAI doesn't publish a
385    /// support matrix, but the GPT-5.5 guide notes compaction is a feature
386    /// shared with GPT-5.4, and the compaction docs exercise the GPT-5.3
387    /// Codex line, so we treat everything from GPT-5.3 onward as supported.
388    ///
389    /// <https://developers.openai.com/api/docs/guides/compaction>
390    pub fn supports_compaction(&self) -> bool {
391        match self {
392            Self::FivePointThreeCodex
393            | Self::FivePointFourNano
394            | Self::FivePointFourMini
395            | Self::FivePointFour
396            | Self::FivePointFourPro
397            | Self::FivePointFive
398            | Self::FivePointFivePro
399            | Self::FivePointSixSol
400            | Self::FivePointSixTerra
401            | Self::FivePointSixLuna => true,
402            Self::Four
403            | Self::FourOmniMini
404            | Self::O3
405            | Self::Five
406            | Self::FiveMini
407            | Self::FiveNano
408            | Self::FivePointOne
409            | Self::FivePointTwo
410            | Self::Custom { .. } => false,
411        }
412    }
413
414    /// Whether OpenAI's Priority processing tier is available for this model.
415    /// Sourced from <https://openai.com/api-priority-processing/>. The `*-pro`,
416    /// `*-nano`, and legacy `gpt-4` variants are not eligible.
417    pub fn supports_priority(&self) -> bool {
418        match self {
419            Self::FourOmniMini
420            | Self::O3
421            | Self::Five
422            | Self::FiveMini
423            | Self::FivePointOne
424            | Self::FivePointTwo
425            | Self::FivePointThreeCodex
426            | Self::FivePointFourMini
427            | Self::FivePointFour
428            | Self::FivePointFive
429            | Self::FivePointSixSol
430            | Self::FivePointSixTerra
431            | Self::FivePointSixLuna => true,
432            Self::Four
433            | Self::FiveNano
434            | Self::FivePointFourNano
435            | Self::FivePointFourPro
436            | Self::FivePointFivePro
437            | Self::Custom { .. } => false,
438        }
439    }
440}
441
442#[cfg(test)]
443mod tests {
444    use super::{Model, ReasoningEffort};
445
446    #[test]
447    fn gpt_5_1_uses_none_reasoning_by_default() {
448        let expected_efforts = [
449            ReasoningEffort::None,
450            ReasoningEffort::Low,
451            ReasoningEffort::Medium,
452            ReasoningEffort::High,
453        ];
454
455        assert_eq!(
456            Model::FivePointOne.reasoning_effort(),
457            Some(ReasoningEffort::None)
458        );
459        assert_eq!(
460            Model::FivePointOne.supported_reasoning_efforts(),
461            expected_efforts.as_slice()
462        );
463    }
464
465    #[test]
466    fn newer_frontier_models_support_none_reasoning() {
467        let expected_efforts = [
468            ReasoningEffort::None,
469            ReasoningEffort::Low,
470            ReasoningEffort::Medium,
471            ReasoningEffort::High,
472            ReasoningEffort::XHigh,
473        ];
474
475        assert_eq!(
476            Model::FivePointTwo.reasoning_effort(),
477            Some(ReasoningEffort::None)
478        );
479        assert_eq!(
480            Model::FivePointTwo.supported_reasoning_efforts(),
481            expected_efforts.as_slice()
482        );
483        assert_eq!(
484            Model::FivePointFour.reasoning_effort(),
485            Some(ReasoningEffort::None)
486        );
487        assert_eq!(
488            Model::FivePointFour.supported_reasoning_efforts(),
489            expected_efforts.as_slice()
490        );
491        assert_eq!(
492            Model::FivePointFive.reasoning_effort(),
493            Some(ReasoningEffort::Medium)
494        );
495        assert_eq!(
496            Model::FivePointFive.supported_reasoning_efforts(),
497            expected_efforts.as_slice()
498        );
499    }
500
501    #[test]
502    fn newer_codex_models_support_low_reasoning_effort() {
503        let expected_efforts = [
504            ReasoningEffort::Low,
505            ReasoningEffort::Medium,
506            ReasoningEffort::High,
507            ReasoningEffort::XHigh,
508        ];
509
510        assert_eq!(
511            Model::FivePointThreeCodex.supported_reasoning_efforts(),
512            expected_efforts.as_slice()
513        );
514    }
515}
516
517#[derive(Debug, Serialize, Deserialize)]
518pub struct StreamOptions {
519    pub include_usage: bool,
520}
521
522impl Default for StreamOptions {
523    fn default() -> Self {
524        Self {
525            include_usage: true,
526        }
527    }
528}
529
530#[derive(Debug, Serialize, Deserialize)]
531pub struct Request {
532    pub model: String,
533    pub messages: Vec<RequestMessage>,
534    pub stream: bool,
535    #[serde(default, skip_serializing_if = "Option::is_none")]
536    pub stream_options: Option<StreamOptions>,
537    #[serde(default, skip_serializing_if = "Option::is_none")]
538    pub max_completion_tokens: Option<u64>,
539    #[serde(default, skip_serializing_if = "Option::is_none")]
540    pub max_tokens: Option<u64>,
541    #[serde(default, skip_serializing_if = "Vec::is_empty")]
542    pub stop: Vec<String>,
543    #[serde(default, skip_serializing_if = "Option::is_none")]
544    pub temperature: Option<f32>,
545    #[serde(default, skip_serializing_if = "Option::is_none")]
546    pub tool_choice: Option<ToolChoice>,
547    /// Whether to enable parallel function calling during tool use.
548    #[serde(default, skip_serializing_if = "Option::is_none")]
549    pub parallel_tool_calls: Option<bool>,
550    #[serde(default, skip_serializing_if = "Vec::is_empty")]
551    pub tools: Vec<ToolDefinition>,
552    #[serde(default, skip_serializing_if = "Option::is_none")]
553    pub prompt_cache_key: Option<String>,
554    #[serde(default, skip_serializing_if = "Option::is_none")]
555    pub reasoning_effort: Option<ReasoningEffort>,
556    #[serde(default, skip_serializing_if = "Option::is_none")]
557    pub service_tier: Option<ServiceTier>,
558}
559
560/// Service tier for OpenAI requests. Maps to the top-level `service_tier`
561/// field on Responses and Chat Completions. We only ever send `Priority`
562/// today (in response to Fast Mode being enabled); the other variants are
563/// included for symmetry with the API and so deserialization of echoed
564/// values does not fail.
565#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
566#[serde(rename_all = "snake_case")]
567pub enum ServiceTier {
568    Auto,
569    Default,
570    Flex,
571    Scale,
572    Priority,
573}
574
575#[derive(Debug, Serialize, Deserialize)]
576#[serde(rename_all = "lowercase")]
577pub enum ToolChoice {
578    Auto,
579    Required,
580    None,
581    #[serde(untagged)]
582    Other(ToolDefinition),
583}
584
585#[derive(Clone, Deserialize, Serialize, Debug)]
586#[serde(tag = "type", rename_all = "snake_case")]
587pub enum ToolDefinition {
588    #[allow(dead_code)]
589    Function { function: FunctionDefinition },
590}
591
592#[derive(Clone, Debug, Serialize, Deserialize)]
593pub struct FunctionDefinition {
594    pub name: String,
595    pub description: Option<String>,
596    pub parameters: Option<Value>,
597}
598
599#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
600#[serde(tag = "role", rename_all = "lowercase")]
601pub enum RequestMessage {
602    Assistant {
603        content: Option<MessageContent>,
604        #[serde(default, skip_serializing_if = "Vec::is_empty")]
605        tool_calls: Vec<ToolCall>,
606        #[serde(default, skip_serializing_if = "Option::is_none")]
607        reasoning_content: Option<String>,
608    },
609    User {
610        content: MessageContent,
611    },
612    System {
613        content: MessageContent,
614    },
615    Tool {
616        content: MessageContent,
617        tool_call_id: String,
618    },
619}
620
621#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)]
622#[serde(untagged)]
623pub enum MessageContent {
624    Plain(String),
625    Multipart(Vec<MessagePart>),
626}
627
628impl MessageContent {
629    pub fn empty() -> Self {
630        MessageContent::Multipart(vec![])
631    }
632
633    pub fn push_part(&mut self, part: MessagePart) {
634        match self {
635            MessageContent::Plain(text) => {
636                *self =
637                    MessageContent::Multipart(vec![MessagePart::Text { text: text.clone() }, part]);
638            }
639            MessageContent::Multipart(parts) if parts.is_empty() => match part {
640                MessagePart::Text { text } => *self = MessageContent::Plain(text),
641                MessagePart::Image { .. } => *self = MessageContent::Multipart(vec![part]),
642            },
643            MessageContent::Multipart(parts) => parts.push(part),
644        }
645    }
646}
647
648impl From<Vec<MessagePart>> for MessageContent {
649    fn from(mut parts: Vec<MessagePart>) -> Self {
650        if let [MessagePart::Text { text }] = parts.as_mut_slice() {
651            MessageContent::Plain(std::mem::take(text))
652        } else {
653            MessageContent::Multipart(parts)
654        }
655    }
656}
657
658#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)]
659#[serde(tag = "type")]
660pub enum MessagePart {
661    #[serde(rename = "text")]
662    Text { text: String },
663    #[serde(rename = "image_url")]
664    Image { image_url: ImageUrl },
665}
666
667#[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)]
668pub struct ImageUrl {
669    pub url: String,
670    #[serde(skip_serializing_if = "Option::is_none")]
671    pub detail: Option<String>,
672}
673
674#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
675pub struct ToolCall {
676    pub id: String,
677    #[serde(flatten)]
678    pub content: ToolCallContent,
679}
680
681#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
682#[serde(tag = "type", rename_all = "lowercase")]
683pub enum ToolCallContent {
684    Function { function: FunctionContent },
685}
686
687#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq)]
688pub struct FunctionContent {
689    pub name: String,
690    pub arguments: String,
691}
692
693#[derive(Clone, Serialize, Deserialize, Debug)]
694pub struct Response {
695    pub id: String,
696    pub object: String,
697    pub created: u64,
698    pub model: String,
699    pub choices: Vec<Choice>,
700    pub usage: Usage,
701}
702
703#[derive(Clone, Serialize, Deserialize, Debug)]
704pub struct Choice {
705    pub index: u32,
706    pub message: RequestMessage,
707    pub finish_reason: Option<String>,
708}
709
710#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
711pub struct ResponseMessageDelta {
712    pub role: Option<Role>,
713    pub content: Option<String>,
714    pub reasoning: Option<String>,
715    #[serde(default, skip_serializing_if = "is_none_or_empty")]
716    pub tool_calls: Option<Vec<ToolCallChunk>>,
717    #[serde(default, skip_serializing_if = "is_none_or_empty")]
718    pub reasoning_content: Option<String>,
719}
720
721#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
722pub struct ToolCallChunk {
723    pub index: usize,
724    pub id: Option<String>,
725
726    // There is also an optional `type` field that would determine if a
727    // function is there. Sometimes this streams in with the `function` before
728    // it streams in the `type`
729    pub function: Option<FunctionChunk>,
730}
731
732#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
733pub struct FunctionChunk {
734    pub name: Option<String>,
735    pub arguments: Option<String>,
736}
737
738#[derive(Clone, Serialize, Deserialize, Debug)]
739pub struct Usage {
740    pub prompt_tokens: Option<u64>,
741    pub completion_tokens: Option<u64>,
742    pub total_tokens: Option<u64>,
743}
744
745#[derive(Serialize, Deserialize, Debug)]
746pub struct ChoiceDelta {
747    pub index: u32,
748    pub delta: Option<ResponseMessageDelta>,
749    pub finish_reason: Option<String>,
750}
751
752#[derive(Error, Debug)]
753pub enum RequestError {
754    #[error("HTTP response error from {provider}'s API: status {status_code} - {body:?}")]
755    HttpResponseError {
756        provider: String,
757        status_code: StatusCode,
758        body: String,
759        headers: HeaderMap<HeaderValue>,
760    },
761    #[error(transparent)]
762    Other(#[from] anyhow::Error),
763}
764
765#[derive(Serialize, Deserialize, Debug)]
766pub struct ResponseStreamError {
767    message: String,
768}
769
770#[derive(Serialize, Deserialize, Debug)]
771#[serde(untagged)]
772pub enum ResponseStreamResult {
773    Ok(ResponseStreamEvent),
774    Err { error: ResponseStreamError },
775}
776
777#[derive(Serialize, Deserialize, Debug)]
778pub struct ResponseStreamEvent {
779    pub choices: Vec<ChoiceDelta>,
780    pub usage: Option<Usage>,
781}
782
783pub async fn non_streaming_completion(
784    client: &dyn HttpClient,
785    api_url: &str,
786    api_key: &str,
787    request: Request,
788) -> Result<Response, RequestError> {
789    let uri = format!("{api_url}/chat/completions");
790    let request_builder = HttpRequest::builder()
791        .method(Method::POST)
792        .uri(uri)
793        .header("Content-Type", "application/json")
794        .header("Authorization", format!("Bearer {}", api_key.trim()));
795
796    let request = request_builder
797        .body(AsyncBody::from(
798            serde_json::to_string(&request).map_err(|e| RequestError::Other(e.into()))?,
799        ))
800        .map_err(|e| RequestError::Other(e.into()))?;
801
802    let mut response = client.send(request).await?;
803    if response.status().is_success() {
804        let mut body = String::new();
805        response
806            .body_mut()
807            .read_to_string(&mut body)
808            .await
809            .map_err(|e| RequestError::Other(e.into()))?;
810
811        serde_json::from_str(&body).map_err(|e| RequestError::Other(e.into()))
812    } else {
813        let mut body = String::new();
814        response
815            .body_mut()
816            .read_to_string(&mut body)
817            .await
818            .map_err(|e| RequestError::Other(e.into()))?;
819
820        Err(RequestError::HttpResponseError {
821            provider: "openai".to_owned(),
822            status_code: response.status(),
823            body,
824            headers: response.headers().clone(),
825        })
826    }
827}
828
829pub async fn stream_completion(
830    client: &dyn HttpClient,
831    provider_name: &str,
832    api_url: &str,
833    api_key: &str,
834    request: Request,
835    extra_headers: &CustomHeaders,
836) -> Result<BoxStream<'static, Result<ResponseStreamEvent>>, RequestError> {
837    let uri = format!("{api_url}/chat/completions");
838    let request = HttpRequest::builder()
839        .method(Method::POST)
840        .uri(uri)
841        .header("Content-Type", "application/json")
842        .header("Authorization", format!("Bearer {}", api_key.trim()))
843        .extra_headers(extra_headers)
844        .body(AsyncBody::from(
845            serde_json::to_string(&request).map_err(|e| RequestError::Other(e.into()))?,
846        ))
847        .map_err(|e| RequestError::Other(e.into()))?;
848
849    let mut response = client.send(request).await?;
850    if response.status().is_success() {
851        let reader = BufReader::new(response.into_body());
852        Ok(reader
853            .lines()
854            .filter_map(|line| async move {
855                match line {
856                    Ok(line) => {
857                        let line = line.strip_prefix("data: ").or_else(|| line.strip_prefix("data:"))?;
858                        if line == "[DONE]" {
859                            None
860                        } else {
861                            match serde_json::from_str(line) {
862                                Ok(ResponseStreamResult::Ok(response)) => Some(Ok(response)),
863                                Ok(ResponseStreamResult::Err { error }) => {
864                                    Some(Err(anyhow!(error.message)))
865                                }
866                                Err(error) => {
867                                    log::error!(
868                                        "Failed to parse OpenAI response into ResponseStreamResult: `{}`\n\
869                                        Response: `{}`",
870                                        error,
871                                        line,
872                                    );
873                                    Some(Err(anyhow!(error)))
874                                }
875                            }
876                        }
877                    }
878                    Err(error) => Some(Err(anyhow!(error))),
879                }
880            })
881            .boxed())
882    } else {
883        let mut body = String::new();
884        response
885            .body_mut()
886            .read_to_string(&mut body)
887            .await
888            .map_err(|e| RequestError::Other(e.into()))?;
889
890        Err(RequestError::HttpResponseError {
891            provider: provider_name.to_owned(),
892            status_code: response.status(),
893            body,
894            headers: response.headers().clone(),
895        })
896    }
897}
898
899#[derive(Copy, Clone, Serialize, Deserialize)]
900pub enum OpenAiEmbeddingModel {
901    #[serde(rename = "text-embedding-3-small")]
902    TextEmbedding3Small,
903    #[serde(rename = "text-embedding-3-large")]
904    TextEmbedding3Large,
905}
906
907#[derive(Serialize)]
908struct OpenAiEmbeddingRequest<'a> {
909    model: OpenAiEmbeddingModel,
910    input: Vec<&'a str>,
911}
912
913#[derive(Deserialize)]
914pub struct OpenAiEmbeddingResponse {
915    pub data: Vec<OpenAiEmbedding>,
916}
917
918#[derive(Deserialize)]
919pub struct OpenAiEmbedding {
920    pub embedding: Vec<f32>,
921}
922
923pub fn embed<'a>(
924    client: &dyn HttpClient,
925    api_url: &str,
926    api_key: &str,
927    model: OpenAiEmbeddingModel,
928    texts: impl IntoIterator<Item = &'a str>,
929) -> impl 'static + Future<Output = Result<OpenAiEmbeddingResponse>> {
930    let uri = format!("{api_url}/embeddings");
931
932    let request = OpenAiEmbeddingRequest {
933        model,
934        input: texts.into_iter().collect(),
935    };
936    let body = AsyncBody::from(serde_json::to_string(&request).unwrap());
937    let request = HttpRequest::builder()
938        .method(Method::POST)
939        .uri(uri)
940        .header("Content-Type", "application/json")
941        .header("Authorization", format!("Bearer {}", api_key.trim()))
942        .body(body)
943        .map(|request| client.send(request));
944
945    async move {
946        let mut response = request?.await?;
947        let mut body = String::new();
948        response.body_mut().read_to_string(&mut body).await?;
949
950        anyhow::ensure!(
951            response.status().is_success(),
952            "error during embedding, status: {:?}, body: {:?}",
953            response.status(),
954            body
955        );
956        let response: OpenAiEmbeddingResponse =
957            serde_json::from_str(&body).context("failed to parse OpenAI embedding response")?;
958        Ok(response)
959    }
960}
961
962// -- Conversions to `language_model_core` types --
963
964impl From<RequestError> for language_model_core::LanguageModelCompletionError {
965    fn from(error: RequestError) -> Self {
966        match error {
967            RequestError::HttpResponseError {
968                provider,
969                status_code,
970                body,
971                headers,
972            } => {
973                let retry_after = headers
974                    .get(http_client::http::header::RETRY_AFTER)
975                    .and_then(|val| val.to_str().ok()?.parse::<u64>().ok())
976                    .map(std::time::Duration::from_secs);
977
978                Self::from_http_status(provider.into(), status_code, body, retry_after)
979            }
980            RequestError::Other(e) => Self::Other(e),
981        }
982    }
983}
984