daimon_provider_local/
generic.rs

1//! Generic OpenAI-compatible model provider.
2//!
3//! For any locally-hosted server that speaks the OpenAI chat/embeddings API
4//! but isn't one of [`crate::llamacpp`], [`crate::llamars`], or
5//! [`crate::ollama`] specifically — vLLM, LM Studio, llamafile, LocalAI, and
6//! similar. Unlike the other local providers, [`OpenAiCompatible::new`]
7//! requires an explicit base URL: there is no sensible default across such a
8//! wide range of servers and ports.
9//!
10//! ```ignore
11//! use daimon_core::{ChatRequest, Message, Model};
12//! use daimon_provider_local::generic::OpenAiCompatible;
13//!
14//! let model = OpenAiCompatible::new("http://localhost:8000")
15//!     .with_model("my-model")
16//!     .with_extra_field("repetition_penalty", serde_json::json!(1.1));
17//! let response = model
18//!     .generate(&ChatRequest::new(vec![Message::user("hello")]))
19//!     .await?;
20//! ```
21
22use std::time::Duration;
23
24use daimon_core::{
25    ChatRequest, ChatResponse, DaimonError, EmbeddingModel, Model, ResponseStream, Result,
26};
27
28use crate::openai_compat::{
29    EmbedRequest, Http, api_error, build_chat_request, parse_chat_response, parse_embed_response,
30    stream_chat_response,
31};
32
33/// Generic OpenAI-compatible model provider. No default base URL — always
34/// construct with [`OpenAiCompatible::new`].
35#[derive(Debug)]
36pub struct OpenAiCompatible {
37    http: Http,
38    model: Option<String>,
39    extra: serde_json::Map<String, serde_json::Value>,
40}
41
42impl OpenAiCompatible {
43    /// Create a client targeting the given base URL (e.g. `http://localhost:8000`).
44    pub fn new(base_url: impl Into<String>) -> Self {
45        let base_url = base_url.into();
46        Self {
47            http: Http::new(&base_url),
48            model: None,
49            extra: serde_json::Map::new(),
50        }
51    }
52
53    /// Override the base URL set at construction.
54    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
55        self.http.set_base_url(url);
56        self
57    }
58
59    /// Set the model name sent in the request body.
60    pub fn with_model(mut self, name: impl Into<String>) -> Self {
61        self.model = Some(name.into());
62        self
63    }
64
65    /// Set the API key / bearer token, if the server requires one.
66    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
67        self.http.set_api_key(key);
68        self
69    }
70
71    /// Set a custom timeout for HTTP requests.
72    pub fn with_timeout(mut self, timeout: Duration) -> Self {
73        self.http.set_timeout(timeout);
74        self
75    }
76
77    /// Set the maximum number of retries for transient (429 / 5xx) errors
78    /// on the initial request (default: 3).
79    pub fn with_max_retries(mut self, retries: u32) -> Self {
80        self.http.set_max_retries(retries);
81        self
82    }
83
84    /// Opts back into warn-and-send for an API key sent over a plaintext
85    /// `http://` base URL (default: hard error). Only use this for a
86    /// genuinely local, unauthenticated-but-keyed server.
87    pub fn allow_plaintext_api_key(mut self) -> Self {
88        self.http.set_allow_plaintext_api_key(true);
89        self
90    }
91
92    /// Add an arbitrary top-level field to every chat request body, for
93    /// server-specific sampling parameters this generic client doesn't know
94    /// about by name (e.g. vLLM's `repetition_penalty`).
95    pub fn with_extra_field(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
96        self.extra.insert(key.into(), value);
97        self
98    }
99}
100
101impl Model for OpenAiCompatible {
102    #[tracing::instrument(skip_all, fields(model = self.model.as_deref().unwrap_or("openai-compatible")))]
103    async fn generate(&self, request: &ChatRequest) -> Result<ChatResponse> {
104        let body = build_chat_request(
105            &request.messages,
106            &request.tools,
107            self.model.as_deref(),
108            request.temperature,
109            request.max_tokens,
110            false,
111            self.extra.clone(),
112        );
113
114        let response = self.http.post("/v1/chat/completions", &body).await?;
115        let status = response.status();
116        if !status.is_success() {
117            let text = response.text().await.unwrap_or_default();
118            return Err(api_error(status, &text, "OpenAI-compatible"));
119        }
120
121        let bytes = response.bytes().await.map_err(|e| {
122            DaimonError::Model(format!("OpenAI-compatible response read error: {e}"))
123        })?;
124        parse_chat_response(&bytes, "OpenAI-compatible")
125    }
126
127    #[tracing::instrument(skip_all, fields(model = self.model.as_deref().unwrap_or("openai-compatible")))]
128    async fn generate_stream(&self, request: &ChatRequest) -> Result<ResponseStream> {
129        let body = build_chat_request(
130            &request.messages,
131            &request.tools,
132            self.model.as_deref(),
133            request.temperature,
134            request.max_tokens,
135            true,
136            self.extra.clone(),
137        );
138
139        let response = self
140            .http
141            .post_streaming("/v1/chat/completions", &body)
142            .await?;
143        let status = response.status();
144        if !status.is_success() {
145            let text = response.text().await.unwrap_or_default();
146            return Err(api_error(status, &text, "OpenAI-compatible"));
147        }
148
149        Ok(stream_chat_response(response, "OpenAI-compatible"))
150    }
151}
152
153/// Generic OpenAI-compatible embedding model. No default base URL — always
154/// construct with [`OpenAiCompatibleEmbedding::new`].
155#[derive(Debug)]
156pub struct OpenAiCompatibleEmbedding {
157    http: Http,
158    model: Option<String>,
159    dimensions: usize,
160}
161
162impl OpenAiCompatibleEmbedding {
163    /// Create a client targeting the given base URL.
164    pub fn new(base_url: impl Into<String>) -> Self {
165        let base_url = base_url.into();
166        Self {
167            http: Http::new(&base_url),
168            model: None,
169            dimensions: 768,
170        }
171    }
172
173    /// Override the base URL set at construction.
174    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
175        self.http.set_base_url(url);
176        self
177    }
178
179    /// Set the model name sent in the request body.
180    pub fn with_model(mut self, name: impl Into<String>) -> Self {
181        self.model = Some(name.into());
182        self
183    }
184
185    /// Set the API key / bearer token, if the server requires one.
186    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
187        self.http.set_api_key(key);
188        self
189    }
190
191    /// Set a custom timeout for HTTP requests.
192    pub fn with_timeout(mut self, timeout: Duration) -> Self {
193        self.http.set_timeout(timeout);
194        self
195    }
196
197    /// Set the maximum number of retries for transient (429 / 5xx) errors
198    /// on the initial request (default: 3).
199    pub fn with_max_retries(mut self, retries: u32) -> Self {
200        self.http.set_max_retries(retries);
201        self
202    }
203
204    /// Opts back into warn-and-send for an API key sent over a plaintext
205    /// `http://` base URL (default: hard error). Only use this for a
206    /// genuinely local, unauthenticated-but-keyed server.
207    pub fn allow_plaintext_api_key(mut self) -> Self {
208        self.http.set_allow_plaintext_api_key(true);
209        self
210    }
211
212    /// Declare the dimensionality of the embeddings this server produces.
213    pub fn with_dimensions(mut self, dims: usize) -> Self {
214        self.dimensions = dims;
215        self
216    }
217}
218
219impl EmbeddingModel for OpenAiCompatibleEmbedding {
220    async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
221        let body = EmbedRequest {
222            model: self.model.as_deref(),
223            input: texts,
224        };
225        let resp = self.http.post("/v1/embeddings", &body).await?;
226        let status = resp.status();
227        if !status.is_success() {
228            let text = resp.text().await.unwrap_or_default();
229            return Err(api_error(status, &text, "OpenAI-compatible"));
230        }
231        let bytes = resp.bytes().await.map_err(|e| {
232            DaimonError::Model(format!("OpenAI-compatible embedding read error: {e}"))
233        })?;
234        parse_embed_response(&bytes, "OpenAI-compatible")
235    }
236
237    fn dimensions(&self) -> usize {
238        self.dimensions
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    #[test]
247    fn test_builder_defaults() {
248        let model = OpenAiCompatible::new("http://localhost:8000");
249        assert_eq!(model.http.base_url(), "http://localhost:8000");
250        assert!(model.model.is_none());
251        assert!(model.extra.is_empty());
252    }
253
254    #[test]
255    fn test_extra_field() {
256        let model = OpenAiCompatible::new("http://localhost:8000")
257            .with_extra_field("repetition_penalty", serde_json::json!(1.1));
258        assert_eq!(model.extra["repetition_penalty"], 1.1);
259    }
260
261    #[test]
262    fn test_builder_chain() {
263        let model = OpenAiCompatible::new("http://localhost:8000")
264            .with_model("my-model")
265            .with_api_key("secret")
266            .with_timeout(Duration::from_secs(15))
267            .with_max_retries(5);
268        assert_eq!(model.model.as_deref(), Some("my-model"));
269        assert_eq!(model.http.api_key(), Some("secret"));
270        assert_eq!(model.http.max_retries(), 5);
271    }
272
273    #[test]
274    fn test_embedding_builder_defaults() {
275        let embed = OpenAiCompatibleEmbedding::new("http://localhost:8000");
276        assert_eq!(embed.dimensions, 768);
277    }
278
279    #[test]
280    fn test_debug_redacts_api_key() {
281        let model =
282            OpenAiCompatible::new("http://localhost:8000").with_api_key("sk-supersecret-key-value");
283        let dbg = format!("{model:?}");
284        assert!(
285            !dbg.contains("sk-supersecret-key-value"),
286            "Debug output must not contain the plaintext API key: {dbg}"
287        );
288        assert!(dbg.contains("[redacted]"));
289    }
290
291    #[test]
292    fn test_embedding_debug_redacts_api_key() {
293        let embed = OpenAiCompatibleEmbedding::new("http://localhost:8000")
294            .with_api_key("sk-supersecret-embed-key");
295        let dbg = format!("{embed:?}");
296        assert!(
297            !dbg.contains("sk-supersecret-embed-key"),
298            "Debug output must not contain the plaintext API key: {dbg}"
299        );
300        assert!(dbg.contains("[redacted]"));
301    }
302
303    #[tokio::test]
304    async fn test_plaintext_api_key_over_http_is_blocked_by_default() {
305        let model = OpenAiCompatible::new("http://localhost:8000").with_api_key("secret");
306        let request = ChatRequest::new(vec![daimon_core::Message::user("hi")]);
307        let err = model.generate(&request).await.unwrap_err();
308        assert!(matches!(err, DaimonError::Builder(_)));
309    }
310
311    #[tokio::test]
312    async fn test_plaintext_api_key_allowed_when_opted_in_does_not_hard_error() {
313        // No server is listening, so this still errors — but it must be a
314        // transport/model error, not the plaintext-key Builder error, proving
315        // the opt-in bypassed the hard block.
316        let model = OpenAiCompatible::new("http://localhost:1")
317            .with_api_key("secret")
318            .with_max_retries(0)
319            .allow_plaintext_api_key();
320        let request = ChatRequest::new(vec![daimon_core::Message::user("hi")]);
321        let err = model.generate(&request).await.unwrap_err();
322        assert!(!matches!(err, DaimonError::Builder(_)));
323    }
324}
daimon_provider_local/generic.rs

daimon_provider_local/
generic.rs
