Skip to main content

yapi_core/
images.rs

1//! pi's image pipeline: what an image becomes before it enters the
2//! conversation. Formats models do not accept become PNG, and with
3//! `images.autoResize` images are fitted to the model's limits.
4
5use std::borrow::Cow;
6use std::io::Cursor;
7
8use base64::Engine as _;
9use base64::engine::general_purpose::STANDARD;
10use image::imageops::{self, FilterType};
11use image::{DynamicImage, ImageOutputFormat};
12use yapi_types::message::{ContentBlock, ImageContent};
13use yapi_types::models::ImageResize;
14
15/// pi's note for an image that cannot be decoded for conversion.
16const CONVERT_FAILED: &str =
17    "[Image omitted: could not be converted to a supported inline image format.]";
18/// pi's note for an image no encoding fits within the size limit.
19const RESIZE_FAILED: &str =
20    "[Image omitted: could not be resized below the inline image size limit.]";
21
22/// pi's default limit: 4.5 MB of base64, below Anthropic's 5 MB.
23const DEFAULT_MAX_BYTES: u64 = 4_718_592;
24
25/// An image ready for the model, with pi's notes on how it changed.
26#[derive(Debug, PartialEq)]
27pub struct Processed {
28    /// The image to send.
29    pub image: ImageContent,
30    /// pi's conversion and dimension notes, in order.
31    pub hints: Vec<String>,
32}
33
34/// pi's `processImage`: an image that is not PNG, JPEG, GIF or WebP becomes
35/// PNG when the `image` crate, with Photon's formats, recognizes it from its
36/// bytes (BMP, TIFF, ICO, PNM, HDR, DDS or farbfeld), and with `auto_resize`
37/// the image is fitted to `limits`, pi's defaults where unset: 2000×2000 and
38/// 4.5 MB of base64. The error is pi's omission note.
39pub fn process(
40    bytes: &[u8],
41    mime_type: &str,
42    auto_resize: bool,
43    limits: Option<&ImageResize>,
44) -> Result<Processed, &'static str> {
45    let base = mime_type
46        .split(';')
47        .next()
48        .unwrap_or_default()
49        .trim()
50        .to_lowercase();
51    let (bytes, mime_type, converted_from) = match base.as_str() {
52        "image/png" => (Cow::Borrowed(bytes), "image/png", None),
53        "image/jpeg" | "image/jpg" => (Cow::Borrowed(bytes), "image/jpeg", None),
54        "image/gif" => (Cow::Borrowed(bytes), "image/gif", None),
55        "image/webp" => (Cow::Borrowed(bytes), "image/webp", None),
56        _ => {
57            let png = decode(bytes)
58                .and_then(|image| encode(&image, ImageOutputFormat::Png))
59                .ok_or(CONVERT_FAILED)?;
60            (Cow::Owned(png), "image/png", Some(base))
61        }
62    };
63    let (image, note) = if auto_resize {
64        resize(&bytes, mime_type, limits).ok_or(RESIZE_FAILED)?
65    } else {
66        (inline(&bytes, mime_type), None)
67    };
68    let mut hints = Vec::new();
69    if let Some(from) = converted_from.filter(|from| !from.is_empty() && *from != image.mime_type) {
70        hints.push(format!(
71            "[Image converted from {from} to {}.]",
72            image.mime_type
73        ));
74    }
75    hints.extend(note);
76    Ok(Processed { image, hints })
77}
78
79/// pi's `_normalizePromptImages`: a prompt's images processed for the
80/// model, with pi's notes, which include the omission note of each image
81/// that fails.
82pub fn normalize_prompt(
83    images: Vec<ImageContent>,
84    auto_resize: bool,
85    limits: Option<&ImageResize>,
86) -> (Vec<ImageContent>, Vec<String>) {
87    let mut normalized = Vec::with_capacity(images.len());
88    let mut hints = Vec::new();
89    for image in images {
90        // Undecodable base64 fails as the empty image it decodes to in pi.
91        let bytes = STANDARD.decode(&image.data).unwrap_or_default();
92        match process(&bytes, &image.mime_type, auto_resize, limits) {
93            Ok(processed) => {
94                normalized.push(processed.image);
95                hints.extend(processed.hints);
96            }
97            Err(message) => hints.push(message.to_owned()),
98        }
99    }
100    (normalized, hints)
101}
102
103/// pi's `normalizeToolResultImages`: each image of a tool result processed
104/// as for the model, followed by a text block of pi's notes when there are
105/// any. An image that fails keeps its block. `None` when nothing changed.
106pub fn normalize_tool_result(
107    content: &[ContentBlock],
108    auto_resize: bool,
109    limits: Option<&ImageResize>,
110) -> Option<Vec<ContentBlock>> {
111    let mut changed = false;
112    let mut normalized = Vec::with_capacity(content.len());
113    for block in content {
114        let ContentBlock::Image(image) = block else {
115            normalized.push(block.clone());
116            continue;
117        };
118        let processed = STANDARD
119            .decode(&image.data)
120            .ok()
121            .and_then(|bytes| process(&bytes, &image.mime_type, auto_resize, limits).ok());
122        match processed {
123            Some(processed) if processed.image != *image || !processed.hints.is_empty() => {
124                normalized.push(ContentBlock::Image(processed.image));
125                if !processed.hints.is_empty() {
126                    normalized.push(ContentBlock::text(processed.hints.join("\n")));
127                }
128                changed = true;
129            }
130            _ => normalized.push(block.clone()),
131        }
132    }
133    changed.then_some(normalized)
134}
135
136/// pi's `resizeImage` with its `formatDimensionNote`.
137fn resize(
138    bytes: &[u8],
139    mime_type: &str,
140    limits: Option<&ImageResize>,
141) -> Option<(ImageContent, Option<String>)> {
142    let resized = resize_image(bytes, mime_type, limits)?;
143    let note = resized.was_resized.then(|| {
144        format!(
145            "[Image: original {}x{}, displayed at {}x{}. Multiply coordinates by {} to map to original image.]",
146            resized.original_width,
147            resized.original_height,
148            resized.width,
149            resized.height,
150            yapi_types::js::to_fixed(
151                f64::from(resized.original_width) / f64::from(resized.width),
152                2
153            )
154        )
155    });
156    Some((resized.image, note))
157}
158
159/// pi's `ResizedImage`: an image fitted to the limits, with its dimensions
160/// before and after.
161#[derive(Debug, PartialEq, serde::Serialize)]
162#[serde(rename_all = "camelCase")]
163pub struct Resized {
164    /// The image to send.
165    #[serde(flatten)]
166    pub image: ImageContent,
167    /// Width as decoded, turned upright.
168    pub original_width: u32,
169    /// Height as decoded, turned upright.
170    pub original_height: u32,
171    /// Width as sent.
172    pub width: u32,
173    /// Height as sent.
174    pub height: u32,
175    /// Whether the image was scaled and encoded again.
176    pub was_resized: bool,
177}
178
179/// pi's `resizeImage`: the image unchanged when it is within `limits`, pi's
180/// defaults where unset (2000×2000 and 4.5 MB of base64). Otherwise the image
181/// scaled to fit, as the first of PNG and JPEG at falling qualities that fits
182/// the byte limit, shrinking by a quarter until one does. `None` when the
183/// image does not decode or no encoding fits.
184pub fn resize_image(
185    bytes: &[u8],
186    mime_type: &str,
187    limits: Option<&ImageResize>,
188) -> Option<Resized> {
189    let limits = limits.cloned().unwrap_or_default();
190    let max_width = limits.max_width.unwrap_or(2000);
191    let max_height = limits.max_height.unwrap_or(2000);
192    let max_bytes = limits.max_bytes.unwrap_or(DEFAULT_MAX_BYTES);
193    let image = decode(bytes)?;
194    let (original_width, original_height) = (image.width(), image.height());
195    let result = |image, width, height, was_resized| Resized {
196        image,
197        original_width,
198        original_height,
199        width,
200        height,
201        was_resized,
202    };
203    if original_width <= max_width
204        && original_height <= max_height
205        && base64_len(bytes.len()) < max_bytes
206    {
207        let image = inline(bytes, mime_type);
208        return Some(result(image, original_width, original_height, false));
209    }
210    let scale = |value: u32, max: u32, other: u32| {
211        yapi_types::js::round(f64::from(value) * f64::from(max) / f64::from(other)) as u32
212    };
213    let (mut width, mut height) = (original_width, original_height);
214    if width > max_width {
215        height = scale(height, max_width, width);
216        width = max_width;
217    }
218    if height > max_height {
219        width = scale(width, max_height, height);
220        height = max_height;
221    }
222    let quality = limits.jpeg_quality.unwrap_or(80);
223    let fallbacks = [85, 70, 55, 40]
224        .into_iter()
225        .filter(|other| *other != quality);
226    let formats: Vec<_> = [ImageOutputFormat::Png, ImageOutputFormat::Jpeg(quality)]
227        .into_iter()
228        .chain(fallbacks.map(ImageOutputFormat::Jpeg))
229        .collect();
230    loop {
231        // pi's Photon fails to encode an empty image.
232        if width == 0 || height == 0 {
233            return None;
234        }
235        let resized = DynamicImage::ImageRgba8(imageops::resize(
236            &image,
237            width,
238            height,
239            FilterType::Lanczos3,
240        ));
241        for format in &formats {
242            let encoded = encode(&resized, format.clone())?;
243            if base64_len(encoded.len()) < max_bytes {
244                let mime_type = match format {
245                    ImageOutputFormat::Png => "image/png",
246                    _ => "image/jpeg",
247                };
248                return Some(result(inline(&encoded, mime_type), width, height, true));
249            }
250        }
251        let shrink = |value: u32| ((f64::from(value) * 0.75).floor() as u32).max(1);
252        let next = (shrink(width), shrink(height));
253        if next == (width, height) {
254            return None;
255        }
256        (width, height) = next;
257    }
258}
259
260/// `bytes` as an image block.
261fn inline(bytes: &[u8], mime_type: &str) -> ImageContent {
262    ImageContent {
263        data: STANDARD.encode(bytes),
264        mime_type: mime_type.to_owned(),
265    }
266}
267
268/// The length of `len` bytes in base64.
269fn base64_len(len: usize) -> u64 {
270    len.div_ceil(3) as u64 * 4
271}
272
273/// The image as RGBA, as Photon holds it, turned upright by its EXIF
274/// orientation as pi does.
275fn decode(bytes: &[u8]) -> Option<DynamicImage> {
276    let image = image::load_from_memory(bytes).ok()?.to_rgba8();
277    Some(DynamicImage::ImageRgba8(match exif_orientation(bytes) {
278        2 => imageops::flip_horizontal(&image),
279        3 => imageops::rotate180(&image),
280        4 => imageops::flip_vertical(&image),
281        5 => imageops::flip_horizontal(&imageops::rotate90(&image)),
282        6 => imageops::rotate90(&image),
283        7 => imageops::flip_horizontal(&imageops::rotate270(&image)),
284        8 => imageops::rotate270(&image),
285        _ => image,
286    }))
287}
288
289/// pi's `convertImageBytesToPng`: `bytes` decoded, turned upright by
290/// their EXIF orientation and encoded as PNG. `None` when they do not decode.
291pub fn convert_to_png(bytes: &[u8]) -> Option<Vec<u8>> {
292    encode(&decode(bytes)?, ImageOutputFormat::Png)
293}
294
295/// The clipboard's `convertToPng`: `bytes` decoded to RGBA and encoded as
296/// PNG, with no EXIF orientation and no resizing. `None` when they do not
297/// decode.
298pub fn to_png(bytes: &[u8]) -> Option<Vec<u8>> {
299    let image = DynamicImage::ImageRgba8(image::load_from_memory(bytes).ok()?.to_rgba8());
300    encode(&image, ImageOutputFormat::Png)
301}
302
303/// `image` in `format`. JPEG drops the alpha channel, as with Photon.
304fn encode(image: &DynamicImage, format: ImageOutputFormat) -> Option<Vec<u8>> {
305    let mut bytes = Cursor::new(Vec::new());
306    image.write_to(&mut bytes, format).ok()?;
307    Some(bytes.into_inner())
308}
309
310/// pi's `getExifOrientation`: the EXIF orientation of a JPEG or WebP image,
311/// or 1 when it has none.
312fn exif_orientation(bytes: &[u8]) -> u16 {
313    let tiff = if bytes.starts_with(&[0xff, 0xd8]) {
314        jpeg_tiff_offset(bytes)
315    } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
316        webp_tiff_offset(bytes)
317    } else {
318        None
319    };
320    tiff.and_then(|start| tiff_orientation(bytes, start))
321        .unwrap_or(1)
322}
323
324fn has_exif_header(bytes: &[u8], offset: usize) -> bool {
325    bytes.get(offset..offset + 6) == Some(b"Exif\0\0")
326}
327
328fn jpeg_tiff_offset(bytes: &[u8]) -> Option<usize> {
329    let mut offset = 2;
330    while offset + 1 < bytes.len() {
331        if bytes[offset] != 0xff {
332            return None;
333        }
334        let marker = bytes[offset + 1];
335        if marker == 0xff {
336            offset += 1;
337            continue;
338        }
339        if marker == 0xe1 {
340            if offset + 10 > bytes.len() {
341                return None;
342            }
343            if has_exif_header(bytes, offset + 4) {
344                return Some(offset + 10);
345            }
346        }
347        let length = usize::from(u16::from_be_bytes(array(bytes, offset + 2)?));
348        offset += 2 + length;
349    }
350    None
351}
352
353fn webp_tiff_offset(bytes: &[u8]) -> Option<usize> {
354    let mut offset = 12;
355    while offset + 8 <= bytes.len() {
356        let size = u32::from_le_bytes(array(bytes, offset + 4)?) as usize;
357        let data = offset + 8;
358        if &bytes[offset..offset + 4] == b"EXIF" {
359            if data.saturating_add(size) > bytes.len() {
360                return None;
361            }
362            // Some files put "Exif\0\0" before the TIFF header.
363            return Some(if size >= 6 && has_exif_header(bytes, data) {
364                data + 6
365            } else {
366                data
367            });
368        }
369        // Chunks are padded to an even size.
370        offset = data.checked_add(size)?.checked_add(size % 2)?;
371    }
372    None
373}
374
375/// The orientation tag of the TIFF header at `start`, if it has a valid one.
376fn tiff_orientation(bytes: &[u8], start: usize) -> Option<u16> {
377    let little_endian = array::<8>(bytes, start)?[..2] == *b"II";
378    let read16 = |at: usize| {
379        let pair = array(bytes, at)?;
380        Some(if little_endian {
381            u16::from_le_bytes(pair)
382        } else {
383            u16::from_be_bytes(pair)
384        })
385    };
386    let read32 = |at: usize| {
387        let quad = array(bytes, at)?;
388        Some(if little_endian {
389            u32::from_le_bytes(quad)
390        } else {
391            u32::from_be_bytes(quad)
392        })
393    };
394    let ifd = start.checked_add(read32(start + 4)? as usize)?;
395    for entry in 0..usize::from(read16(ifd)?) {
396        let at = ifd + 2 + entry * 12;
397        // pi reads an entry only when all 12 bytes are there.
398        array::<12>(bytes, at)?;
399        if read16(at)? == 0x0112 {
400            return read16(at + 8).filter(|value| (1..=8).contains(value));
401        }
402    }
403    None
404}
405
406/// The `N` bytes at `at`, if there are that many.
407pub(crate) fn array<const N: usize>(bytes: &[u8], at: usize) -> Option<[u8; N]> {
408    bytes.get(at..at.checked_add(N)?)?.try_into().ok()
409}
410
411#[cfg(test)]
412mod tests {
413    #![allow(
414        clippy::unwrap_used,
415        reason = "test helpers; a panic is a test failure"
416    )]
417
418    use super::*;
419    use image::{Rgba, RgbaImage};
420
421    /// A `width`×`height` image: a gradient, or noise that compresses badly.
422    fn pixels(width: u32, height: u32, noisy: bool) -> DynamicImage {
423        let mut seed = 7u32;
424        DynamicImage::ImageRgba8(RgbaImage::from_fn(width, height, |x, y| {
425            seed = seed.wrapping_mul(1_103_515_245).wrapping_add(12_345);
426            let noise = if noisy { (seed >> 16) as u8 } else { 0 };
427            Rgba([(x % 256) as u8 ^ noise, (y % 256) as u8, noise, 255])
428        }))
429    }
430
431    fn file(image: &DynamicImage, format: ImageOutputFormat) -> Vec<u8> {
432        encode(image, format).unwrap()
433    }
434
435    fn limits(width: u32, height: u32, bytes: Option<u64>) -> ImageResize {
436        ImageResize {
437            max_width: Some(width),
438            max_height: Some(height),
439            max_bytes: bytes,
440            jpeg_quality: None,
441        }
442    }
443
444    fn dimensions(image: &ImageContent) -> (u32, u32) {
445        let decoded = image::load_from_memory(&STANDARD.decode(&image.data).unwrap()).unwrap();
446        (decoded.width(), decoded.height())
447    }
448
449    #[test]
450    fn small_images_unchanged() {
451        let png = file(&pixels(30, 20, false), ImageOutputFormat::Png);
452        let processed = process(&png, "image/png", true, None).unwrap();
453        assert_eq!(processed.image.data, STANDARD.encode(&png));
454        assert_eq!(processed.image.mime_type, "image/png");
455        assert!(processed.hints.is_empty());
456    }
457
458    #[test]
459    fn large_images_scaled() {
460        let png = file(&pixels(2100, 10, false), ImageOutputFormat::Png);
461        let processed = process(&png, "image/png", true, None).unwrap();
462        assert_eq!(processed.image.mime_type, "image/png");
463        assert_eq!(dimensions(&processed.image), (2000, 10));
464        assert_eq!(
465            processed.hints,
466            [
467                "[Image: original 2100x10, displayed at 2000x10. Multiply coordinates by 1.05 to map to original image.]"
468            ]
469        );
470
471        let tall = file(&pixels(30, 90, false), ImageOutputFormat::Png);
472        let limits = limits(40, 40, None);
473        let processed = process(&tall, "image/png", true, Some(&limits)).unwrap();
474        assert_eq!(dimensions(&processed.image), (13, 40));
475        assert_eq!(
476            processed.hints,
477            [
478                "[Image: original 30x90, displayed at 13x40. Multiply coordinates by 2.31 to map to original image.]"
479            ]
480        );
481        // Without auto-resize the image passes as it is.
482        let processed = process(&tall, "image/png", false, Some(&limits)).unwrap();
483        assert_eq!(processed.image.data, STANDARD.encode(&tall));
484        assert!(processed.hints.is_empty());
485    }
486
487    #[test]
488    fn jpeg_replaces_large_png() {
489        let image = pixels(64, 64, true);
490        let png = file(&image, ImageOutputFormat::Png);
491        let limits = limits(2000, 2000, Some(base64_len(png.len())));
492        let processed = process(&png, "image/png", true, Some(&limits)).unwrap();
493        assert_eq!(processed.image.mime_type, "image/jpeg");
494        assert_eq!(
495            STANDARD.decode(&processed.image.data).unwrap(),
496            file(&image, ImageOutputFormat::Jpeg(80))
497        );
498        assert_eq!(
499            processed.hints,
500            [
501                "[Image: original 64x64, displayed at 64x64. Multiply coordinates by 1.00 to map to original image.]"
502            ]
503        );
504    }
505
506    #[test]
507    fn shrinks_then_omits() {
508        let png = file(&pixels(40, 40, true), ImageOutputFormat::Png);
509        let limits = limits(2000, 2000, Some(1200));
510        let processed = process(&png, "image/png", true, Some(&limits)).unwrap();
511        assert!(processed.image.data.len() < 1200);
512        let (width, height) = dimensions(&processed.image);
513        assert!(width < 40 && width == height);
514        let none_fit = ImageResize {
515            max_bytes: Some(1),
516            ..ImageResize::default()
517        };
518        assert_eq!(
519            process(&png, "image/png", true, Some(&none_fit)),
520            Err(RESIZE_FAILED)
521        );
522    }
523
524    #[test]
525    fn exif_upright_before_resize() {
526        let jpeg = file(&pixels(4, 2, false), ImageOutputFormat::Jpeg(90));
527        // An APP1 segment with one IFD (image file directory) entry:
528        // orientation 6, rotate clockwise.
529        let mut exif = b"\xff\xe1\x00\x22Exif\0\0MM\0\x2a\0\0\0\x08\0\x01".to_vec();
530        exif.extend_from_slice(b"\x01\x12\0\x03\0\0\0\x01\0\x06\0\0\0\0\0\0");
531        let rotated = [&jpeg[..2], &exif, &jpeg[2..]].concat();
532        assert_eq!(exif_orientation(&rotated), 6);
533        let limits = limits(3, 3, None);
534        let processed = process(&rotated, "image/jpeg", true, Some(&limits)).unwrap();
535        assert_eq!(dimensions(&processed.image), (2, 3));
536        assert_eq!(
537            processed.hints,
538            [
539                "[Image: original 2x4, displayed at 2x3. Multiply coordinates by 1.00 to map to original image.]"
540            ]
541        );
542    }
543
544    #[test]
545    fn to_png_ignores_exif() {
546        let jpeg = file(&pixels(4, 2, false), ImageOutputFormat::Jpeg(90));
547        let mut exif = b"\xff\xe1\x00\x22Exif\0\0MM\0\x2a\0\0\0\x08\0\x01".to_vec();
548        exif.extend_from_slice(b"\x01\x12\0\x03\0\0\0\x01\0\x06\0\0\0\0\0\0");
549        let rotated = [&jpeg[..2], &exif, &jpeg[2..]].concat();
550        let png = to_png(&rotated).unwrap();
551        assert!(png.starts_with(b"\x89PNG"));
552        let decoded = image::load_from_memory(&png).unwrap();
553        assert_eq!((decoded.width(), decoded.height()), (4, 2));
554    }
555
556    /// Runs `bytes` through `process` as `mime_type` with and without resizing
557    /// and returns the PNG it becomes, after checking pi's conversion note.
558    fn converted(bytes: &[u8], mime_type: &str) -> Vec<u8> {
559        let mut png = None;
560        for auto_resize in [true, false] {
561            let processed = process(bytes, mime_type, auto_resize, None)
562                .unwrap_or_else(|note| panic!("{mime_type}: {note}"));
563            assert_eq!(processed.image.mime_type, "image/png");
564            assert_eq!(
565                processed.hints,
566                [format!("[Image converted from {mime_type} to image/png.]")]
567            );
568            let data = STANDARD.decode(&processed.image.data).unwrap();
569            assert!(png.replace(data.clone()).is_none_or(|first| first == data));
570        }
571        png.unwrap()
572    }
573
574    /// A 2x2 RGB TIFF holding red, green, blue and white.
575    const TIFF: &str = "SUkqAAgAAAAJAAABAwABAAAAAgAAAAEBAwABAAAAAgAAAAIBAwADAAAAegAAAAMBAwABAAAAAQAAAAYBAwABAAAAAgAAABEBBAABAAAAgAAAABUBAwABAAAAAwAAABYBAwABAAAAAgAAABcBBAABAAAADAAAAAAAAAAIAAgACAD/AAAA/wAAAP////8=";
576
577    /// A 4x4 DDS of one DXT1 block that is solid red.
578    fn dds() -> Vec<u8> {
579        let mut bytes = b"DDS ".to_vec();
580        let mut header = [0u32; 31];
581        header[0] = 124; // size
582        header[1] = 0x1 | 0x2 | 0x4 | 0x1000; // caps, height, width, pixel format
583        header[2] = 4; // height
584        header[3] = 4; // width
585        header[4] = 8; // linear size
586        header[18] = 32; // pixel format size
587        header[19] = 0x4; // fourcc
588        header[20] = u32::from_le_bytes(*b"DXT1");
589        header[26] = 0x1000; // texture
590        bytes.extend(header.iter().flat_map(|word| word.to_le_bytes()));
591        bytes.extend([0x00, 0xf8, 0x00, 0x00, 0, 0, 0, 0]);
592        bytes
593    }
594
595    #[test]
596    fn other_formats_become_png() {
597        let image = pixels(3, 2, false);
598        let png = file(&image, ImageOutputFormat::Png);
599        let rgba16 = DynamicImage::ImageRgba16(image.to_rgba16());
600        // PNM has no alpha channel, so it takes the opaque pixels as RGB.
601        let rgb = DynamicImage::ImageRgb8(image.to_rgb8());
602        let pnm = ImageOutputFormat::Pnm(image::codecs::pnm::PnmSubtype::Pixmap(
603            image::codecs::pnm::SampleEncoding::Binary,
604        ));
605        for (mime_type, bytes) in [
606            ("image/bmp", file(&image, ImageOutputFormat::Bmp)),
607            ("image/x-icon", file(&image, ImageOutputFormat::Ico)),
608            ("image/x-portable-pixmap", file(&rgb, pnm)),
609            (
610                "image/x-farbfeld",
611                file(&rgba16, ImageOutputFormat::Farbfeld),
612            ),
613        ] {
614            assert_eq!(converted(&bytes, mime_type), png, "{mime_type}");
615        }
616
617        let tiff = RgbaImage::from_fn(2, 2, |x, y| match (x, y) {
618            (0, 0) => Rgba([255, 0, 0, 255]),
619            (1, 0) => Rgba([0, 255, 0, 255]),
620            (0, 1) => Rgba([0, 0, 255, 255]),
621            _ => Rgba([255; 4]),
622        });
623        assert_eq!(
624            converted(&STANDARD.decode(TIFF).unwrap(), "image/tiff"),
625            file(&DynamicImage::ImageRgba8(tiff), ImageOutputFormat::Png)
626        );
627
628        let red = DynamicImage::ImageRgba8(RgbaImage::from_pixel(4, 4, Rgba([255, 0, 0, 255])));
629        assert_eq!(
630            converted(&dds(), "image/vnd.ms-dds"),
631            file(&red, ImageOutputFormat::Png)
632        );
633
634        // Radiance HDR stores shared-exponent floats, so only the size is exact.
635        let mut hdr = Vec::new();
636        image::codecs::hdr::HdrEncoder::new(&mut hdr)
637            .encode(&[image::Rgb([0.5f32, 0.25, 1.0]); 6], 3, 2)
638            .unwrap();
639        let decoded = image::load_from_memory(&converted(&hdr, "image/vnd.radiance")).unwrap();
640        assert_eq!((decoded.width(), decoded.height()), (3, 2));
641
642        assert_eq!(
643            process(b"not an image", "image/tiff", true, None),
644            Err(CONVERT_FAILED)
645        );
646        assert_eq!(
647            process(b"not an image", "image/png", true, None),
648            Err(RESIZE_FAILED)
649        );
650    }
651
652    #[test]
653    fn normalizes_prompts_and_results() {
654        let png = ImageContent {
655            data: STANDARD.encode(file(&pixels(3, 2, false), ImageOutputFormat::Png)),
656            mime_type: "image/png".into(),
657        };
658        let bmp = ImageContent {
659            data: STANDARD.encode(file(&pixels(3, 2, false), ImageOutputFormat::Bmp)),
660            mime_type: "image/bmp".into(),
661        };
662        let broken = ImageContent {
663            data: "AAAA".into(),
664            mime_type: "image/png".into(),
665        };
666
667        let (images, hints) =
668            normalize_prompt(vec![png.clone(), bmp.clone(), broken.clone()], true, None);
669        assert_eq!(images[0], png);
670        assert_eq!(images[1].mime_type, "image/png");
671        assert_eq!(images.len(), 2);
672        assert_eq!(
673            hints,
674            [
675                "[Image converted from image/bmp to image/png.]",
676                RESIZE_FAILED
677            ]
678        );
679
680        let unchanged = [ContentBlock::text("a"), ContentBlock::Image(png.clone())];
681        assert_eq!(normalize_tool_result(&unchanged, true, None), None);
682        let content = [
683            ContentBlock::Image(bmp),
684            ContentBlock::Image(broken.clone()),
685            ContentBlock::text("a"),
686        ];
687        let normalized = normalize_tool_result(&content, true, None).unwrap();
688        assert_eq!(
689            normalized,
690            [
691                ContentBlock::Image(images[1].clone()),
692                ContentBlock::text("[Image converted from image/bmp to image/png.]"),
693                ContentBlock::Image(broken),
694                ContentBlock::text("a"),
695            ]
696        );
697    }
698}