mistralrs_server_core/
util.rs

1//! ## General utilities.
2
3use image::DynamicImage;
4use mistralrs_core::AudioInput;
5use mistralrs_core::MistralRs;
6use std::error::Error;
7use std::sync::Arc;
8use tokio::{
9    fs::{self, File},
10    io::AsyncReadExt,
11};
12
13/// Parses and loads an image from a URL, file path, or data URL.
14///
15/// This function accepts various input formats and attempts to parse them in order:
16/// 1. First tries to parse as a complete URL (http/https/file/data schemes)
17/// 2. If that fails, checks if it's a local file path and converts to file URL
18/// 3. Finally falls back to treating it as a malformed URL and returns an error
19///
20/// ### Arguments
21///
22/// * `url_unparsed` - A string that can be:
23///   - An HTTP/HTTPS URL (e.g., "<https://example.com/image.png>")
24///   - A file path (e.g., "/path/to/image.jpg" or "image.png")
25///   - A data URL with base64 encoded image (e.g., "data:image/png;base64,...")
26///   - A file URL (e.g., "file:///path/to/image.jpg")
27///
28/// ### Examples
29///
30/// ```ignore
31/// use mistralrs_server_core::util::parse_image_url;
32///
33/// // Load from HTTP URL
34/// let image = parse_image_url("https://example.com/photo.jpg").await?;
35///
36/// // Load from local file path
37/// let image = parse_image_url("./assets/logo.png").await?;
38///
39/// // Load from data URL
40/// let image = parse_image_url("data:image/png;base64,iVBORw0KGgoAAAANS...").await?;
41///
42/// // Load from file URL
43/// let image = parse_image_url("file:///home/user/picture.jpg").await?;
44/// ```
45pub async fn parse_image_url(url_unparsed: &str) -> Result<DynamicImage, anyhow::Error> {
46    let url = if let Ok(url) = url::Url::parse(url_unparsed) {
47        url
48    } else if File::open(url_unparsed).await.is_ok() {
49        url::Url::from_file_path(std::path::absolute(url_unparsed)?)
50            .map_err(|_| anyhow::anyhow!("Could not parse file path: {}", url_unparsed))?
51    } else {
52        url::Url::parse(url_unparsed)
53            .map_err(|_| anyhow::anyhow!("Could not parse as base64 data: {}", url_unparsed))?
54    };
55
56    let bytes = if url.scheme() == "http" || url.scheme() == "https" {
57        // Read from http
58        match reqwest::get(url.clone()).await {
59            Ok(http_resp) => http_resp.bytes().await?.to_vec(),
60            Err(e) => anyhow::bail!(e),
61        }
62    } else if url.scheme() == "file" {
63        let path = url
64            .to_file_path()
65            .map_err(|_| anyhow::anyhow!("Could not parse file path: {}", url))?;
66
67        if let Ok(mut f) = File::open(&path).await {
68            // Read from local file
69            let metadata = fs::metadata(&path).await?;
70            let mut buffer = vec![0; metadata.len() as usize];
71            f.read_exact(&mut buffer).await?;
72            buffer
73        } else {
74            anyhow::bail!("Could not open file at path: {}", url);
75        }
76    } else if url.scheme() == "data" {
77        // Decode with base64
78        let data_url = data_url::DataUrl::process(url.as_str())?;
79        data_url.decode_to_vec()?.0
80    } else {
81        anyhow::bail!("Unsupported URL scheme: {}", url.scheme());
82    };
83
84    Ok(image::load_from_memory(&bytes)?)
85}
86
87/// Parses and loads an audio file from a URL, file path, or data URL.
88pub async fn parse_audio_url(url_unparsed: &str) -> Result<AudioInput, anyhow::Error> {
89    let url = if let Ok(url) = url::Url::parse(url_unparsed) {
90        url
91    } else if File::open(url_unparsed).await.is_ok() {
92        url::Url::from_file_path(std::path::absolute(url_unparsed)?)
93            .map_err(|_| anyhow::anyhow!("Could not parse file path: {}", url_unparsed))?
94    } else {
95        url::Url::parse(url_unparsed)
96            .map_err(|_| anyhow::anyhow!("Could not parse as base64 data: {}", url_unparsed))?
97    };
98
99    let bytes = if url.scheme() == "http" || url.scheme() == "https" {
100        match reqwest::get(url.clone()).await {
101            Ok(http_resp) => http_resp.bytes().await?.to_vec(),
102            Err(e) => anyhow::bail!(e),
103        }
104    } else if url.scheme() == "file" {
105        let path = url
106            .to_file_path()
107            .map_err(|_| anyhow::anyhow!("Could not parse file path: {}", url))?;
108
109        if let Ok(mut f) = File::open(&path).await {
110            let metadata = fs::metadata(&path).await?;
111            let mut buffer = vec![0; metadata.len() as usize];
112            f.read_exact(&mut buffer).await?;
113            buffer
114        } else {
115            anyhow::bail!("Could not open file at path: {}", url);
116        }
117    } else if url.scheme() == "data" {
118        let data_url = data_url::DataUrl::process(url.as_str())?;
119        data_url.decode_to_vec()?.0
120    } else {
121        anyhow::bail!("Unsupported URL scheme: {}", url.scheme());
122    };
123
124    AudioInput::from_bytes(&bytes)
125}
126
127/// Validates that the requested model matches one of the loaded models.
128///
129/// This function checks if the model parameter from an OpenAI API request
130/// matches one of the models that are currently loaded by the server.
131///
132/// The special model name "default" can be used to bypass this validation,
133/// which is useful for clients that require a model parameter but want
134/// to use the default model.
135///
136/// ### Arguments
137///
138/// * `requested_model` - The model name from the API request
139/// * `state` - The MistralRs state containing the loaded models info
140///
141/// ### Returns
142///
143/// Returns `Ok(())` if the model is available or if "default" is specified, otherwise returns an error.
144pub fn validate_model_name(
145    requested_model: &str,
146    state: Arc<MistralRs>,
147) -> Result<(), anyhow::Error> {
148    // Allow "default" as a special case to bypass validation
149    if requested_model == "default" {
150        return Ok(());
151    }
152
153    let available_models = state
154        .list_models()
155        .map_err(|e| anyhow::anyhow!("Failed to get available models: {}", e))?;
156
157    if available_models.is_empty() {
158        anyhow::bail!("No models are currently loaded.");
159    }
160
161    if !available_models.contains(&requested_model.to_string()) {
162        anyhow::bail!(
163            "Requested model '{}' is not available. Available models: {}. Use 'default' to use the default model.",
164            requested_model,
165            available_models.join(", ")
166        );
167    }
168    Ok(())
169}
170
171/// Sanitize error messages to remove internal implementation details like stack traces.
172/// This ensures that sensitive internal information is not exposed to API clients.
173///
174/// The function traverses the error chain to find the deepest (root) error and returns its message.
175/// This is useful for API responses where we want to provide meaningful error information
176/// without exposing internal stack traces or implementation details.
177///
178/// ### Arguments
179///
180/// * `error` - The error to sanitize
181///
182/// ### Returns
183///
184/// The message from the root cause error in the error chain
185///
186/// ### Examples
187///
188/// ```ignore
189/// use mistralrs_server_core::util::sanitize_error_message;
190///
191/// // For a simple error without chain
192/// let error = std::io::Error::new(std::io::ErrorKind::NotFound, "File not found");
193/// assert_eq!(sanitize_error_message(&error), "File not found");
194///
195/// // For chained errors, returns the root cause
196/// let root = std::io::Error::new(std::io::ErrorKind::PermissionDenied, "Access denied");
197/// let wrapped = anyhow::Error::new(root).context("Failed to read file");
198/// // This would return "Access denied" instead of "Failed to read file"
199/// ```
200pub fn sanitize_error_message(error: &(dyn Error + 'static)) -> String {
201    // Traverse the error chain to find the deepest (root) error and return its message.
202    let mut current: &dyn Error = error;
203
204    // Keep traversing until we find an error with no source
205    while let Some(source) = current.source() {
206        current = source;
207    }
208
209    // Return the message of the root cause error
210    current.to_string()
211}
212
213#[cfg(test)]
214mod tests {
215    use image::GenericImageView;
216
217    use super::*;
218
219    #[tokio::test]
220    async fn test_parse_image_url() {
221        // from URL
222        let url = "https://www.rust-lang.org/logos/rust-logo-32x32.png";
223        let image = parse_image_url(url).await.unwrap();
224        assert_eq!(image.dimensions(), (32, 32));
225
226        let url = "http://www.rust-lang.org/logos/rust-logo-32x32.png";
227        let image = parse_image_url(url).await.unwrap();
228        assert_eq!(image.dimensions(), (32, 32));
229
230        // from file path
231        let url = "resources/rust-logo-32x32.png";
232        let image = parse_image_url(url).await.unwrap();
233        assert_eq!(image.dimensions(), (32, 32));
234
235        // URL must be an absolute path
236        let absolute_path = std::path::absolute(url).unwrap();
237        let url = format!("file://{}", absolute_path.as_os_str().to_str().unwrap());
238        let image = parse_image_url(&url).await.unwrap();
239        assert_eq!(image.dimensions(), (32, 32));
240
241        // from base64 encoded image (rust-logo-32x32.png)
242        let url = "
243        iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAHhElEQVR4AZXVA5Aky9bA8f/JzKrq
244        npleX9u2bdu2bdu2bdv29z1e2zZm7k5PT3dXZeZ56I6J3o03sbG/iDLOSQuT6fptF11OREYDj4uR
245        i9UmAAdHa9ZH6QP+n8kg1+26HJMU44rG+4OjL3YqCv+693HOwcHiTJeYY2NUch/PLI3sOdZY82lU
246        Xbynp3yzEXMH8CCTINfuujzDEXQlVN9sju8/uFHPTy2KWLVpWsl9ZGQCvY2AF0ulu0RTBRHIi1AV
247        iZU0sSd0dWWXZKVsUeAVhiFX7roCwzGDA9rXV6uaqH/YcmnmPEQGg4IYIoLAYRHRABcaQIGuNMVa
248        IS98tZnnjOxJK4AwDDlzs0XoNGUmlWDsPr/98ucLIerrPVlCI8KAWMAQAYWXo8rKipyuMDewuaAv
249        g6wMgEa6M0dX6ugdqOPQxSs96WqlcukqoEoHuWiHZelki3yF/vHVV0OhdCUJfzZyQlYiiPlR4RxV
250        bgKqAbNthDto2Q64U6ACbAicKzCtAON6Uqr1HAk5XYlZEXiNDnLaBgvQxqiSzPdLX70PNT9U/pN9
251        0xNdSjT2UoXjJ84+x6ygwMQ/bSdyOnCgamSqSpmBepOY53OliXHAh7TJsesuCMBMU/XM/+dvve/9
252        PhgYl2X8Xi8IWZkobAg8xuQjx24L3KEamaY7oX/8IDZ6ukZkCwDvA8gpGy1EG9Vq44fRpXTa3oZv
253        BVeIQERQQBFUQQGE4frWj+3hdyxQtei2oHe4UDB1KvxWL34EpqPNLjzdWKYZXVqpr3fgdDV2QSJZ
254        A4M3loC0gqu0ggsgrXMQhlEBlgR2Au6OyF+AWby4hbvU4xVtRF2x7OQ7a+QbOWKN+Rjp4lF/NOLZ
255        o0sZvw96MIJPM6IYVEFFAFrnTEhF6CSqdHgaWEeEzQXuc9EzlYv8VPdkwtHAOS4P8Fsw52A40Mc4
256        rRp5ICKzR2WhCC8hsrgqFaWlXRPfCfJtRIxCVGQWYFoAERCU9rY2AKqXAO/7qHFA7YIi4ccczgFw
257        U490G/7WV7/KZdm0/YVHxBwcka2jyEKI7K3KdQorarvqI4aIXAWcRQcv5ixBjgZFVBEUg2KJiyFM
258        i+p2EpWB3L+UJHbamPsfMo37uFoj/8RjKqkRmiqAfKcioPKjwoCCjWKIGALtBERmi5glFHGAgswC
259        ur4AoEO1YDReAvITaNVIfInEVWOzoJwaBqGSwCeuVucTMebUIsTzBCHCD9G63dH4tCh4m5qIELAE
260        ERRDRHbT/24AAhMch5rgqSDm4AIARpS1uZ3k4SqLEsUCnFoX84nLupPLaoN+/8xZ6kUBYr82wT9N
261        W7AlghgChnbwdlMICCgCgMLQmaiAcDAdqtJ1R0+ie4VQrNCFosh5TpjJSe6fVGWnqFrx/2Nwe7G0
262        233oqNIxL9D5iSIIiCLwenu8VwEy9ad4cTNL9AQMXqlaa/7uxqt7SiQcVCvijQGDUR1Nh4jRdvt3
263        BJdjgLPbISsKVwHbgSBA66gVgY2A252GSlQ90eRNECEP4MUd5AO3u5Els2Gtrjd2p5a81iSYZB7v
264        ssUKl6UBm0dkRECI0k4Ag8LsiiyhkAHvANsDGwIVBQRoH/cW+CiKORBtt70oYi1i/I2p8LsL8BLW
265        3FHLwwZZYkcMBlDM68rQsEMZCk4EFNkN2A0EhTOB44D3gWWYgC4HvB4RsI5gLJlEGj72q5pHrY1f
266        mmpuqiD7RB+lK7UYUWIMRCxDHU4mCA5IR/tT0PIcbQoTvBMR1BfEEEjTlIZXKaVm32DcByYYR0Y4
267        0ahWvIIRxWiEULSDT9zZBGUChpZrAIZLQsWAS4gKZXzFKCcaBWMUSNypwNq1PL7dlbZe0qgovK7I
268        i4q8oPCywl//szG08TrwZscquF773kvACwovgbwOhugjXaljoFl8Z4ysHdFTI4qLKOO9q46o8Jei
269        VswWFMoOGj6iToyKfKKwIMgTAmcJyvB48j9bRM4DlqaVwPr4BiUTEAzdJowyxvwlQhXARQSAclc2
270        a+H101rD10ZWulbsq4GE5qKOdOoc+5kaORM4i0lQpAIcDnyIcjC+USEGxpSgVq9+4JKskZg4a3v0
271        IPussxSdTGNw2jrJD7Zcpmg2iaKr9LrRqMpLWLsE8DrDQ1UXA15HZDppDq5p0Zum6TbUBmq4LJsO
272        +JEOro6j0xQjyrMzWNCs1/4Y210a21+El0blvdU+NwZCeFGNuQGRe4APgCotFWA+YtwK1d3QiAb/
273        j9EujBmXKL9U69UscRUncfaJE5Dd112G4RT1hmZpQuYM3+UJRYBoE8QYMA40gmrrKIKGCNaSaMHU
274        iQfvaeYBQBiG7LTKEgxndI/przbii001lR7HqnVXphmU3EdCFHLjEI1ojKQWyonQI4pGT72Rv2OE
275        r7qtrAaMYBiy1xpLMClSTk/Nm/6EZtAHUms3y1BKzvCHZESEMVqnXkRyHwjIAxbdLEnsqcBJTILz
276        xjApU46pnBe8fwF4pb+/8Ywv/DK9zbCKsfWXUBhf+A1dOX00S+xfgc3L3dmKWSn7iklDjthxbSaH
277        c7YCVIAfi6JYn5bHjTHTGmurQJXJ8C/um928G9zK4gAAAABJRU5ErkJggg==
278        ";
279
280        let url = format!("data:image/png;base64,{url}");
281        let image = parse_image_url(&url).await.unwrap();
282        assert_eq!(image.dimensions(), (32, 32));
283
284        // audio from base64
285        let audio_b64 = "UklGRiYAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQIAAAAAAA==";
286        let url = format!("data:audio/wav;base64,{audio_b64}");
287        let audio = parse_audio_url(&url).await.unwrap();
288        assert_eq!(audio.sample_rate, 8000);
289        assert_eq!(audio.samples.len(), 1);
290    }
291
292    #[test]
293    fn test_sanitize_error_message_with_backtrace() {
294        // Test error with backtrace
295        let error_with_backtrace = "Failed to parse Forge Provider response: A weight is negative, too large or not a valid number
296  0: candle_core::error::Error::bt
297  1: mistralrs_core::sampler::Sampler::sample_multinomial
298  2: mistralrs_core::sampler::Sampler::sample_top_kp_min_p
299  3: mistralrs_core::sampler::Sampler::sample
300  4: mistralrs_core::pipeline::sampling::sample_sequence::{{closure}}";
301
302        struct TestError(String);
303        impl std::fmt::Display for TestError {
304            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305                write!(f, "{}", self.0)
306            }
307        }
308        impl std::fmt::Debug for TestError {
309            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
310                write!(f, "{}", self.0)
311            }
312        }
313        impl std::error::Error for TestError {}
314
315        let error = TestError(error_with_backtrace.to_string());
316        let sanitized = sanitize_error_message(&error);
317
318        // Since TestError has no source(), it should return the full message including backtrace
319        assert_eq!(sanitized, error_with_backtrace);
320        // The improved solution returns the root error as-is when there's no error chain
321    }
322
323    #[test]
324    fn test_sanitize_error_message_without_backtrace() {
325        // Test error without backtrace
326        let simple_error = "Simple error message without backtrace";
327
328        struct TestError(String);
329        impl std::fmt::Display for TestError {
330            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
331                write!(f, "{}", self.0)
332            }
333        }
334        impl std::fmt::Debug for TestError {
335            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
336                write!(f, "{}", self.0)
337            }
338        }
339        impl std::error::Error for TestError {}
340
341        let error = TestError(simple_error.to_string());
342        let sanitized = sanitize_error_message(&error);
343
344        assert_eq!(sanitized, simple_error);
345    }
346
347    #[test]
348    fn test_sanitize_error_message_with_chain() {
349        // Test error chain - the root cause should be extracted
350        use std::fmt;
351
352        #[derive(Debug)]
353        struct RootError;
354        impl fmt::Display for RootError {
355            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
356                write!(f, "Root cause: Database connection failed")
357            }
358        }
359        impl std::error::Error for RootError {}
360
361        #[derive(Debug)]
362        struct MiddleError(Box<dyn std::error::Error>);
363        impl fmt::Display for MiddleError {
364            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
365                write!(f, "Middle error: Service unavailable")
366            }
367        }
368        impl std::error::Error for MiddleError {
369            fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
370                Some(&*self.0)
371            }
372        }
373
374        #[derive(Debug)]
375        struct TopError(Box<dyn std::error::Error>);
376        impl fmt::Display for TopError {
377            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
378                write!(
379                    f,
380                    "Top error: Request failed with backtrace\n  0: some::module::function"
381                )
382            }
383        }
384        impl std::error::Error for TopError {
385            fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
386                Some(&*self.0)
387            }
388        }
389
390        let root = RootError;
391        let middle = MiddleError(Box::new(root));
392        let top = TopError(Box::new(middle));
393
394        let sanitized = sanitize_error_message(&top);
395
396        // Should return the root cause, not the top-level error with backtrace
397        assert_eq!(sanitized, "Root cause: Database connection failed");
398        assert!(!sanitized.contains("backtrace"));
399        assert!(!sanitized.contains("Request failed"));
400    }
401}