cortexkit_model_batch/
image.rs

1//! The item image shape every model-batch role shares, its validation, and
2//! its canonical form.
3//!
4//! An image is a `data:` URL string or `{content_type, base64}`. Remote URLs
5//! are never accepted. Validation decodes the base64, so an image a module
6//! admits is one the provider can read, and sizes are the decoded bytes.
7//! The canonical form, `{content_type, base64}` with the decoded bytes
8//! re-encoded as standard padded base64, is what batch identity hashes: one
9//! image sent in either spelling is the same batch (`RULES.md` §2, §13).
10
11use base64::{
12    alphabet,
13    engine::{general_purpose::STANDARD, DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
14    Engine,
15};
16use serde::{Deserialize, Serialize};
17
18use crate::errors::Refusal;
19
20/// Image content types a provider accepts.
21pub mod content_types {
22    pub const PNG: &str = "image/png";
23    pub const JPEG: &str = "image/jpeg";
24    pub const WEBP: &str = "image/webp";
25    pub const ALL: &[&str] = &[PNG, JPEG, WEBP];
26}
27
28/// Largest image, in decoded bytes (4 MiB), as Cloudflare Workers AI states
29/// it for its image-taking batch models.
30pub const IMAGE_MAX_BYTES: usize = 4 * 1024 * 1024;
31/// Largest total of an item's images, in decoded bytes (8 MiB).
32pub const IMAGES_MAX_TOTAL_BYTES: usize = 8 * 1024 * 1024;
33/// Largest image, in pixels (16 megapixels). The role crates do not decode
34/// image formats, so this is a runner's check, not [`check_image`]'s.
35pub const IMAGE_MAX_PIXELS: u64 = 16_000_000;
36
37/// Standard base64, padded or not. Trailing bits that a canonical encoding
38/// would leave zero are refused, so one byte string has one accepted
39/// spelling per padding choice.
40const LENIENT: GeneralPurpose = GeneralPurpose::new(
41    &alphabet::STANDARD,
42    GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
43);
44
45/// An image in its object form: `{content_type, base64}`.
46#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
47pub struct InlineImage {
48    /// One of the model's accepted formats ([`content_types::ALL`] unless
49    /// the model states fewer).
50    pub content_type: String,
51    pub base64: String,
52}
53
54impl InlineImage {
55    pub fn new(content_type: impl Into<String>, base64: impl Into<String>) -> Self {
56        Self {
57            content_type: content_type.into(),
58            base64: base64.into(),
59        }
60    }
61}
62
63/// An image: a `data:` URL string, or `{content_type, base64}`. Remote URLs
64/// are never accepted.
65#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
66#[serde(untagged)]
67pub enum Image {
68    DataUrl(String),
69    Inline(InlineImage),
70}
71
72/// An image's content type and decoded bytes.
73#[derive(Clone, Debug, PartialEq, Eq)]
74pub struct DecodedImage {
75    pub content_type: String,
76    pub bytes: Vec<u8>,
77}
78
79impl Image {
80    /// The image's size in decoded bytes, estimated from its base64 length.
81    /// Exact for well-formed base64; [`Image::decode`] is the check.
82    pub fn approx_bytes(&self) -> usize {
83        let payload = match self {
84            Self::DataUrl(url) => url.split_once(',').map_or("", |(_, data)| data),
85            Self::Inline(inline) => &inline.base64,
86        };
87        payload.trim_end_matches('=').len() * 3 / 4
88    }
89
90    /// The content type and decoded bytes, or why the image is malformed.
91    /// A `data:` URL must be `data:<content type>;base64,<payload>`. The
92    /// content type is not checked against a format list here.
93    pub fn decode(&self) -> Result<DecodedImage, String> {
94        let (content_type, payload) = match self {
95            Self::DataUrl(url) => {
96                let Some(rest) = url.strip_prefix("data:") else {
97                    return Err(
98                        "an image string must be a data: URL; remote URLs are refused".into(),
99                    );
100                };
101                let Some((media, payload)) = rest.split_once(',') else {
102                    return Err("a data: URL has no payload".into());
103                };
104                let Some(content_type) = media.strip_suffix(";base64") else {
105                    return Err("a data: URL is data:<content type>;base64,<payload>".into());
106                };
107                (content_type, payload)
108            }
109            Self::Inline(inline) => (inline.content_type.as_str(), inline.base64.as_str()),
110        };
111        let bytes = LENIENT
112            .decode(payload)
113            .map_err(|e| format!("the image's base64 does not decode: {e}"))?;
114        if bytes.is_empty() {
115            return Err("an image has no bytes".into());
116        }
117        Ok(DecodedImage {
118            content_type: content_type.to_owned(),
119            bytes,
120        })
121    }
122
123    /// The canonical form batch identity hashes: the content type and the
124    /// decoded bytes as standard padded base64.
125    pub fn canonical(&self) -> Result<InlineImage, String> {
126        let decoded = self.decode()?;
127        Ok(InlineImage::new(
128            decoded.content_type,
129            STANDARD.encode(decoded.bytes),
130        ))
131    }
132}
133
134/// A model's image limits per item, as an extractor `role.describe` entry
135/// states them: `{max, max_bytes_each, max_bytes_total, formats}`.
136#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
137#[non_exhaustive]
138pub struct ImageLimits {
139    /// Most images in one item.
140    pub max: u64,
141    /// Largest image, in decoded bytes.
142    pub max_bytes_each: u64,
143    /// Largest total of one item's images, in decoded bytes.
144    pub max_bytes_total: u64,
145    /// The content types accepted, a subset of [`content_types::ALL`].
146    pub formats: Vec<String>,
147}
148
149impl ImageLimits {
150    pub fn new(max: u64, max_bytes_each: u64, max_bytes_total: u64, formats: Vec<String>) -> Self {
151        Self {
152            max,
153            max_bytes_each,
154            max_bytes_total,
155            formats,
156        }
157    }
158
159    /// Workers AI's limits ([`IMAGE_MAX_BYTES`] each,
160    /// [`IMAGES_MAX_TOTAL_BYTES`] in total, every format) with at most
161    /// `max` images.
162    pub fn workers_ai(max: u64) -> Self {
163        Self::new(
164            max,
165            IMAGE_MAX_BYTES as u64,
166            IMAGES_MAX_TOTAL_BYTES as u64,
167            content_types::ALL.iter().map(|t| (*t).to_owned()).collect(),
168        )
169    }
170}
171
172/// Check one image's form, content type, base64 and decoded size against
173/// `limits`, returning its decoded size. The error says why it is refused.
174pub fn check_image_with(image: &Image, limits: &ImageLimits) -> Result<u64, String> {
175    if let Image::DataUrl(url) = image {
176        if !url.starts_with("data:") {
177            return Err("an image string must be a data: URL; remote URLs are refused".into());
178        }
179    }
180    let content_type = match image {
181        Image::DataUrl(url) => url
182            .split_once(',')
183            .map(|(media, _)| media["data:".len()..].split(';').next().unwrap_or(""))
184            .ok_or_else(|| "a data: URL has no payload".to_owned())?,
185        Image::Inline(inline) => inline.content_type.as_str(),
186    };
187    if !limits.formats.iter().any(|format| format == content_type) {
188        return Err(format!("unsupported image content type {content_type}"));
189    }
190    let size = image.decode()?.bytes.len() as u64;
191    if size > limits.max_bytes_each {
192        return Err(format!(
193            "an image is at most {} bytes",
194            limits.max_bytes_each
195        ));
196    }
197    Ok(size)
198}
199
200/// Check one image against Workers AI's limits ([`ImageLimits::workers_ai`]).
201pub fn check_image(image: &Image) -> Result<(), String> {
202    check_image_with(image, &ImageLimits::workers_ai(u64::MAX)).map(|_| ())
203}
204
205/// Check an item's images: their count against `limits.max` (`None` when
206/// the model takes no images), each image ([`check_image_with`]), and their
207/// decoded total. `field` is the item's `images` path (`items[i].images`);
208/// a refusal names it, or `field[j]` for one image.
209pub fn check_images(
210    model: &str,
211    field: &str,
212    images: &[Image],
213    limits: Option<&ImageLimits>,
214) -> Result<(), Refusal> {
215    let max = limits.map_or(0, |limits| limits.max);
216    if images.len() as u64 > max {
217        return Err(Refusal::over_limit(
218            field,
219            max,
220            if limits.is_none() {
221                format!("{model} takes no images")
222            } else {
223                format!("{} images; at most {max}", images.len())
224            },
225        ));
226    }
227    let Some(limits) = limits else {
228        return Ok(());
229    };
230    let mut total: u64 = 0;
231    for (j, image) in images.iter().enumerate() {
232        let size = check_image_with(image, limits)
233            .map_err(|why| Refusal::invalid_params(format!("{field}[{j}]"), why))?;
234        total = total.saturating_add(size);
235    }
236    if total > limits.max_bytes_total {
237        return Err(Refusal::over_limit(
238            field,
239            limits.max_bytes_total,
240            format!("{total} image bytes; at most {}", limits.max_bytes_total),
241        ));
242    }
243    Ok(())
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    #[test]
251    fn remote_urls_unknown_types_and_bad_base64_are_refused() {
252        assert!(check_image(&Image::DataUrl("https://example.com/a.png".into())).is_err());
253        assert!(check_image(&Image::DataUrl("data:image/gif;base64,AAAA".into())).is_err());
254        assert!(check_image(&Image::DataUrl("data:image/png;base64,AAAA".into())).is_ok());
255        assert!(check_image(&Image::Inline(InlineImage::new("image/webp", "AAAA"))).is_ok());
256        assert!(check_image(&Image::Inline(InlineImage::new("image/bmp", "AAAA"))).is_err());
257        let malformed =
258            check_image(&Image::Inline(InlineImage::new("image/png", "A!A="))).unwrap_err();
259        assert!(malformed.contains("does not decode"), "{malformed}");
260        let plain = check_image(&Image::DataUrl("data:image/png,AAAA".into())).unwrap_err();
261        assert!(plain.contains(";base64"), "{plain}");
262        assert!(check_image(&Image::Inline(InlineImage::new("image/png", ""))).is_err());
263    }
264
265    #[test]
266    fn both_spellings_have_one_canonical_form() {
267        let url = Image::DataUrl("data:image/png;base64,AAECAw".into());
268        let inline = Image::Inline(InlineImage::new("image/png", "AAECAw=="));
269        assert_eq!(url.canonical().unwrap(), inline.canonical().unwrap());
270        assert_eq!(url.canonical().unwrap().base64, "AAECAw==");
271    }
272
273    #[test]
274    fn decoded_sizes_are_checked_against_the_limits() {
275        let limits = ImageLimits::new(2, 3, 5, vec![content_types::PNG.into()]);
276        let three = Image::Inline(InlineImage::new("image/png", "AAAA"));
277        let four = Image::Inline(InlineImage::new("image/png", "AAAAAA=="));
278        assert_eq!(check_image_with(&three, &limits), Ok(3));
279        assert!(check_image_with(&four, &limits).is_err());
280        let refusal = check_images(
281            "m",
282            "items[0].images",
283            &[three.clone(), three.clone()],
284            Some(&limits),
285        )
286        .unwrap_err();
287        assert_eq!(refusal.field(), Some("items[0].images"));
288        assert_eq!(refusal.detail.limit, Some(5));
289        let refusal = check_images("m", "items[0].images", &[three], None).unwrap_err();
290        assert_eq!(refusal.detail.limit, Some(0));
291        assert_eq!(refusal.message, "m takes no images");
292    }
293}
cortexkit_model_batch/image.rs

cortexkit_model_batch/
image.rs
