Skip to main content

ai_agent/utils/
image_validation.rs

1//! Image validation utilities
2//!
3//! Validates that all images in messages are within the API size limit.
4
5use serde::{Deserialize, Serialize};
6
7/// Information about an oversized image
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct OversizedImage {
10    /// Index of the image (1-based)
11    pub index: usize,
12    /// Size in bytes
13    pub size: usize,
14}
15
16/// Error thrown when one or more images exceed the API size limit
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ImageSizeError {
19    /// List of oversized images
20    pub oversized_images: Vec<OversizedImage>,
21    /// Maximum allowed size
22    pub max_size: usize,
23}
24
25impl std::fmt::Display for ImageSizeError {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        let first_image = self.oversized_images.first();
28        if self.oversized_images.len() == 1 && first_image.is_some() {
29            let img = first_image.unwrap();
30            write!(
31                f,
32                "Image base64 size ({} bytes) exceeds API limit ({} bytes). Please resize the image before sending.",
33                Self::format_size(img.size),
34                Self::format_size(self.max_size)
35            )
36        } else {
37            let img_list = self
38                .oversized_images
39                .iter()
40                .map(|img| format!("Image {}: {} bytes", img.index, Self::format_size(img.size)))
41                .collect::<Vec<_>>()
42                .join(", ");
43            write!(
44                f,
45                "{} images exceed the API limit ({} bytes): {}. Please resize these images before sending.",
46                self.oversized_images.len(),
47                Self::format_size(self.max_size),
48                img_list
49            )
50        }
51    }
52}
53
54impl std::error::Error for ImageSizeError {}
55
56impl ImageSizeError {
57    fn format_size(bytes: usize) -> String {
58        if bytes >= 1_000_000 {
59            format!("{:.1}MB", bytes as f64 / 1_000_000.0)
60        } else if bytes >= 1_000 {
61            format!("{:.1}KB", bytes as f64 / 1_000.0)
62        } else {
63            format!("{}B", bytes)
64        }
65    }
66}
67
68/// API image maximum base64 size (5MB)
69pub const API_IMAGE_MAX_BASE64_SIZE: usize = 5 * 1024 * 1024;
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_error_message_single() {
77        let error = ImageSizeError {
78            oversized_images: vec![OversizedImage {
79                index: 1,
80                size: 6_000_000,
81            }],
82            max_size: API_IMAGE_MAX_BASE64_SIZE,
83        };
84        let msg = error.to_string();
85        assert!(msg.contains("Image base64 size"));
86    }
87
88    #[test]
89    fn test_error_message_multiple() {
90        let error = ImageSizeError {
91            oversized_images: vec![
92                OversizedImage {
93                    index: 1,
94                    size: 6_000_000,
95                },
96                OversizedImage {
97                    index: 2,
98                    size: 7_000_000,
99                },
100            ],
101            max_size: API_IMAGE_MAX_BASE64_SIZE,
102        };
103        let msg = error.to_string();
104        assert!(msg.contains("2 images"));
105    }
106}