agent_io/llm/openai_compatible/
mod.rs

1//! OpenAI-compatible API base implementation
2//!
3//! This module provides a base implementation for any LLM provider that uses
4//! the OpenAI-compatible API format (Ollama, OpenRouter, DeepSeek, Groq, etc.)
5
6mod request;
7mod response;
8mod types;
9
10use async_trait::async_trait;
11use derive_builder::Builder;
12use futures::StreamExt;
13use reqwest::{Client, StatusCode};
14use std::time::Duration;
15
16use crate::llm::{
17    BaseChatModel, ChatCompletion, ChatStream, LlmError, Message, ToolChoice, ToolDefinition,
18};
19
20use types::*;
21
22/// OpenAI-compatible Chat Model base implementation
23#[derive(Builder, Clone)]
24#[builder(pattern = "owned", build_fn(skip))]
25pub struct ChatOpenAICompatible {
26    /// Model identifier
27    #[builder(setter(into))]
28    pub(super) model: String,
29    /// API key (optional for some providers like Ollama)
30    #[builder(setter(into), default = "None")]
31    pub(super) api_key: Option<String>,
32    /// Base URL for the API
33    #[builder(setter(into))]
34    pub(super) base_url: String,
35    /// Provider name for identification
36    #[builder(setter(into))]
37    pub(super) provider: String,
38    /// Temperature for sampling
39    #[builder(default = "0.2")]
40    pub(super) temperature: f32,
41    /// Maximum completion tokens
42    #[builder(default = "Some(4096)")]
43    pub(super) max_completion_tokens: Option<u64>,
44    /// HTTP client
45    #[builder(setter(skip))]
46    pub(super) client: Client,
47    /// Context window size
48    #[builder(setter(skip))]
49    pub(super) context_window: u64,
50    /// Whether to include Bearer prefix in auth header
51    #[builder(default = "true")]
52    pub(super) use_bearer_auth: bool,
53}
54
55pub(crate) struct OpenAICompatibleProviderConfig<'a> {
56    pub provider: &'a str,
57    pub default_base_url: &'a str,
58    pub api_key_env: Option<&'a str>,
59    pub base_url_env: Option<&'a str>,
60    pub use_bearer_auth: bool,
61    pub default_temperature: f32,
62}
63
64impl ChatOpenAICompatible {
65    /// Create a builder for configuration
66    pub fn builder() -> ChatOpenAICompatibleBuilder {
67        ChatOpenAICompatibleBuilder::default()
68    }
69
70    pub(crate) fn build_provider(
71        config: OpenAICompatibleProviderConfig<'_>,
72        model: Option<String>,
73        api_key: Option<String>,
74        base_url: Option<String>,
75        temperature: Option<f32>,
76        max_tokens: Option<u64>,
77    ) -> Result<Self, LlmError> {
78        let model = model.ok_or_else(|| LlmError::Config("model is required".into()))?;
79
80        let api_key = match (api_key, config.api_key_env) {
81            (Some(key), _) => Some(key),
82            (None, Some(env_var)) => std::env::var(env_var).ok(),
83            (None, None) => None,
84        };
85
86        if api_key.is_none() && config.api_key_env.is_some() {
87            return Err(LlmError::Config(format!(
88                "{} not set",
89                config.api_key_env.unwrap_or_default()
90            )));
91        }
92
93        let base_url = base_url
94            .or_else(|| {
95                config
96                    .base_url_env
97                    .and_then(|env_var| std::env::var(env_var).ok())
98            })
99            .unwrap_or_else(|| config.default_base_url.to_string());
100
101        ChatOpenAICompatible::builder()
102            .model(model)
103            .base_url(base_url)
104            .provider(config.provider)
105            .api_key(api_key)
106            .use_bearer_auth(config.use_bearer_auth)
107            .temperature(temperature.unwrap_or(config.default_temperature))
108            .max_completion_tokens(max_tokens)
109            .build()
110    }
111
112    fn map_error_status(status: StatusCode, body: String) -> LlmError {
113        match status {
114            StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => LlmError::Auth(body),
115            StatusCode::NOT_FOUND => LlmError::ModelNotFound(body),
116            StatusCode::TOO_MANY_REQUESTS => LlmError::RateLimit,
117            _ => LlmError::Api(format!("API error ({}): {}", status, body)),
118        }
119    }
120
121    /// Build the HTTP client
122    fn build_client() -> Client {
123        Client::builder()
124            .timeout(Duration::from_secs(120))
125            .build()
126            .expect("Failed to create HTTP client")
127    }
128
129    /// Default context window
130    fn default_context_window() -> u64 {
131        128_000
132    }
133
134    /// Get the API URL
135    fn api_url(&self) -> String {
136        format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
137    }
138}
139
140impl ChatOpenAICompatibleBuilder {
141    pub fn build(&self) -> Result<ChatOpenAICompatible, LlmError> {
142        let model = self
143            .model
144            .clone()
145            .ok_or_else(|| LlmError::Config("model is required".into()))?;
146        let base_url = self
147            .base_url
148            .clone()
149            .ok_or_else(|| LlmError::Config("base_url is required".into()))?;
150        let provider = self
151            .provider
152            .clone()
153            .ok_or_else(|| LlmError::Config("provider is required".into()))?;
154
155        Ok(ChatOpenAICompatible {
156            client: ChatOpenAICompatible::build_client(),
157            context_window: ChatOpenAICompatible::default_context_window(),
158            model,
159            api_key: self.api_key.clone().flatten(),
160            base_url,
161            provider,
162            temperature: self.temperature.unwrap_or(0.2),
163            max_completion_tokens: self.max_completion_tokens.flatten(),
164            use_bearer_auth: self.use_bearer_auth.unwrap_or(true),
165        })
166    }
167}
168
169#[async_trait]
170impl BaseChatModel for ChatOpenAICompatible {
171    fn model(&self) -> &str {
172        &self.model
173    }
174
175    fn provider(&self) -> &str {
176        &self.provider
177    }
178
179    fn context_window(&self) -> Option<u64> {
180        Some(self.context_window)
181    }
182
183    async fn invoke(
184        &self,
185        messages: Vec<Message>,
186        tools: Option<Vec<ToolDefinition>>,
187        tool_choice: Option<ToolChoice>,
188    ) -> Result<ChatCompletion, LlmError> {
189        let request = self.build_request(messages, tools, tool_choice, false)?;
190
191        let mut req = self
192            .client
193            .post(self.api_url())
194            .header("Content-Type", "application/json");
195
196        if let Some(ref api_key) = self.api_key {
197            if self.use_bearer_auth {
198                req = req.header("Authorization", format!("Bearer {}", api_key));
199            } else {
200                req = req.header("Authorization", api_key.clone());
201            }
202        }
203
204        let response = req.json(&request).send().await?;
205
206        if !response.status().is_success() {
207            let status = response.status();
208            let body = response.text().await.unwrap_or_default();
209            return Err(Self::map_error_status(
210                status,
211                format!("{}: {}", self.provider, body),
212            ));
213        }
214
215        let completion: OpenAICompatibleResponse = response.json().await?;
216        Ok(Self::parse_response(completion))
217    }
218
219    async fn invoke_stream(
220        &self,
221        messages: Vec<Message>,
222        tools: Option<Vec<ToolDefinition>>,
223        tool_choice: Option<ToolChoice>,
224    ) -> Result<ChatStream, LlmError> {
225        let request = self.build_request(messages, tools, tool_choice, true)?;
226
227        let mut req = self
228            .client
229            .post(self.api_url())
230            .header("Content-Type", "application/json");
231
232        if let Some(ref api_key) = self.api_key {
233            if self.use_bearer_auth {
234                req = req.header("Authorization", format!("Bearer {}", api_key));
235            } else {
236                req = req.header("Authorization", api_key.clone());
237            }
238        }
239
240        let response = req.json(&request).send().await?;
241
242        if !response.status().is_success() {
243            let status = response.status();
244            let body = response.text().await.unwrap_or_default();
245            return Err(Self::map_error_status(
246                status,
247                format!("{}: {}", self.provider, body),
248            ));
249        }
250
251        let stream = response.bytes_stream().filter_map(|result| async move {
252            match result {
253                Ok(bytes) => {
254                    let text = String::from_utf8_lossy(&bytes);
255                    Self::parse_stream_chunk(&text)
256                }
257                Err(e) => Some(Err(LlmError::Stream(e.to_string()))),
258            }
259        });
260
261        Ok(Box::pin(stream))
262    }
263
264    fn supports_vision(&self) -> bool {
265        // Most OpenAI-compatible providers support vision
266        true
267    }
268}
agent_io/llm/openai_compatible/mod.rs

agent_io/llm/openai_compatible/
mod.rs
