agent_io/llm/openai/
mod.rs

1//! OpenAI Chat Model implementation
2
3mod request;
4mod response;
5mod types;
6
7use async_trait::async_trait;
8use derive_builder::Builder;
9use futures::StreamExt;
10use reqwest::{Client, StatusCode};
11use std::time::Duration;
12
13use crate::llm::{
14    BaseChatModel, ChatCompletion, ChatStream, LlmError, Message, ToolChoice, ToolDefinition,
15};
16
17use types::*;
18
19const OPENAI_BASE_URL: &str = "https://api.openai.com/v1";
20const CHAT_COMPLETIONS_PATH: &str = "/chat/completions";
21
22/// OpenAI Chat Model
23#[derive(Builder, Clone)]
24#[builder(pattern = "owned", build_fn(skip))]
25pub struct ChatOpenAI {
26    /// Model identifier
27    #[builder(setter(into))]
28    pub(super) model: String,
29    /// API key
30    pub(super) api_key: String,
31    /// Base URL (for proxies)
32    #[builder(setter(into, strip_option), default = "None")]
33    pub(super) base_url: Option<String>,
34    /// Temperature for sampling
35    #[builder(default = "0.2")]
36    pub(super) temperature: f32,
37    /// Maximum completion tokens
38    #[builder(default = "Some(4096)")]
39    pub(super) max_completion_tokens: Option<u64>,
40    /// HTTP client
41    #[builder(setter(skip))]
42    pub(super) client: Client,
43    /// Context window size
44    #[builder(setter(skip))]
45    pub(super) context_window: u64,
46}
47
48impl ChatOpenAI {
49    /// Create a new OpenAI chat model
50    pub fn new(model: impl Into<String>) -> Result<Self, LlmError> {
51        let api_key = std::env::var("OPENAI_API_KEY")
52            .map_err(|_| LlmError::Config("OPENAI_API_KEY not set".into()))?;
53        let base_url = std::env::var("OPENAI_BASE_URL").ok();
54
55        let mut builder = Self::builder().model(model).api_key(api_key);
56        if let Some(url) = base_url {
57            builder = builder.base_url(url);
58        }
59        builder.build()
60    }
61
62    /// Create a builder for configuration
63    pub fn builder() -> ChatOpenAIBuilder {
64        ChatOpenAIBuilder::default()
65    }
66
67    /// Check if this is a reasoning model (o1, o3, o4, gpt-5)
68    fn is_reasoning_model(&self) -> bool {
69        let model_lower = self.model.to_lowercase();
70        model_lower.starts_with("o1")
71            || model_lower.starts_with("o3")
72            || model_lower.starts_with("o4")
73            || model_lower.starts_with("gpt-5")
74    }
75
76    /// Get the API URL
77    fn api_url(&self) -> String {
78        let base = self.base_url.as_deref().unwrap_or(OPENAI_BASE_URL);
79        format!("{}{}", base.trim_end_matches('/'), CHAT_COMPLETIONS_PATH)
80    }
81
82    fn map_error_status(status: StatusCode, body: String) -> LlmError {
83        match status {
84            StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => LlmError::Auth(body),
85            StatusCode::NOT_FOUND => LlmError::ModelNotFound(body),
86            StatusCode::TOO_MANY_REQUESTS => LlmError::RateLimit,
87            _ => LlmError::Api(format!("OpenAI API error ({}): {}", status, body)),
88        }
89    }
90
91    /// Build the HTTP client
92    fn build_client() -> Client {
93        Client::builder()
94            .timeout(Duration::from_secs(120))
95            .build()
96            .expect("Failed to create HTTP client")
97    }
98
99    /// Get context window for model
100    fn get_context_window(model: &str) -> u64 {
101        let model_lower = model.to_lowercase();
102
103        // GPT-4o family
104        if model_lower.contains("gpt-4o") || model_lower.contains("gpt-4-turbo") {
105            128_000
106        }
107        // GPT-4
108        else if model_lower.starts_with("gpt-4") {
109            8_192
110        }
111        // GPT-3.5
112        else if model_lower.starts_with("gpt-3.5") {
113            16_385
114        }
115        // O1/O3/O4 reasoning models
116        else if model_lower.starts_with("o1")
117            || model_lower.starts_with("o3")
118            || model_lower.starts_with("o4")
119        {
120            200_000
121        }
122        // Default
123        else {
124            128_000
125        }
126    }
127}
128
129impl ChatOpenAIBuilder {
130    pub fn build(&self) -> Result<ChatOpenAI, LlmError> {
131        let model = self
132            .model
133            .clone()
134            .ok_or_else(|| LlmError::Config("model is required".into()))?;
135        let api_key = self
136            .api_key
137            .clone()
138            .ok_or_else(|| LlmError::Config("api_key is required".into()))?;
139
140        Ok(ChatOpenAI {
141            context_window: ChatOpenAI::get_context_window(&model),
142            client: ChatOpenAI::build_client(),
143            model,
144            api_key,
145            base_url: self.base_url.clone().flatten(),
146            temperature: self.temperature.unwrap_or(0.2),
147            max_completion_tokens: self.max_completion_tokens.flatten(),
148        })
149    }
150}
151
152#[async_trait]
153impl BaseChatModel for ChatOpenAI {
154    fn model(&self) -> &str {
155        &self.model
156    }
157
158    fn provider(&self) -> &str {
159        "openai"
160    }
161
162    fn context_window(&self) -> Option<u64> {
163        Some(self.context_window)
164    }
165
166    async fn invoke(
167        &self,
168        messages: Vec<Message>,
169        tools: Option<Vec<ToolDefinition>>,
170        tool_choice: Option<ToolChoice>,
171    ) -> Result<ChatCompletion, LlmError> {
172        let request = self.build_request(messages, tools, tool_choice, false)?;
173
174        let response = self
175            .client
176            .post(self.api_url())
177            .header("Authorization", format!("Bearer {}", self.api_key))
178            .header("Content-Type", "application/json")
179            .json(&request)
180            .send()
181            .await?;
182
183        if !response.status().is_success() {
184            let status = response.status();
185            let body = response.text().await.unwrap_or_default();
186            return Err(Self::map_error_status(status, body));
187        }
188        let body = response.text().await?;
189        tracing::debug!("OpenAI raw response: {}", body);
190
191        // Some proxies always return SSE format regardless of stream=false.
192        // Detect by checking if the body starts with "data:"
193        if body.trim_start().starts_with("data:") {
194            return self.parse_sse_as_completion(&body);
195        }
196
197        let completion: OpenAIResponse = serde_json::from_str(&body).map_err(|e| {
198            LlmError::Api(format!(
199                "Failed to parse response: {}\nBody: {}",
200                e,
201                &body[..body.len().min(500)]
202            ))
203        })?;
204        Ok(self.parse_response(completion))
205    }
206
207    async fn invoke_stream(
208        &self,
209        messages: Vec<Message>,
210        tools: Option<Vec<ToolDefinition>>,
211        tool_choice: Option<ToolChoice>,
212    ) -> Result<ChatStream, LlmError> {
213        let request = self.build_request(messages, tools, tool_choice, true)?;
214
215        let response = self
216            .client
217            .post(self.api_url())
218            .header("Authorization", format!("Bearer {}", self.api_key))
219            .header("Content-Type", "application/json")
220            .json(&request)
221            .send()
222            .await?;
223
224        if !response.status().is_success() {
225            let status = response.status();
226            let body = response.text().await.unwrap_or_default();
227            return Err(Self::map_error_status(status, body));
228        }
229
230        let stream = response.bytes_stream().filter_map(|result| async move {
231            match result {
232                Ok(bytes) => {
233                    let text = String::from_utf8_lossy(&bytes);
234                    Self::parse_stream_chunk(&text)
235                }
236                Err(e) => Some(Err(LlmError::Stream(e.to_string()))),
237            }
238        });
239
240        Ok(Box::pin(stream))
241    }
242
243    fn supports_vision(&self) -> bool {
244        let model_lower = self.model.to_lowercase();
245        model_lower.contains("gpt-4o")
246            || model_lower.contains("gpt-4-turbo")
247            || model_lower.contains("gpt-4-vision")
248            || model_lower.contains("gpt-4.1")
249    }
250}
agent_io/llm/openai/mod.rs

agent_io/llm/openai/
mod.rs
