1use 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
15const CONVERT_FAILED: &str =
17 "[Image omitted: could not be converted to a supported inline image format.]";
18const RESIZE_FAILED: &str =
20 "[Image omitted: could not be resized below the inline image size limit.]";
21
22const DEFAULT_MAX_BYTES: u64 = 4_718_592;
24
25#[derive(Debug, PartialEq)]
27pub struct Processed {
28 pub image: ImageContent,
30 pub hints: Vec<String>,
32}
33
34pub 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
79pub 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 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
103pub 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
136fn 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#[derive(Debug, PartialEq, serde::Serialize)]
162#[serde(rename_all = "camelCase")]
163pub struct Resized {
164 #[serde(flatten)]
166 pub image: ImageContent,
167 pub original_width: u32,
169 pub original_height: u32,
171 pub width: u32,
173 pub height: u32,
175 pub was_resized: bool,
177}
178
179pub 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 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
260fn inline(bytes: &[u8], mime_type: &str) -> ImageContent {
262 ImageContent {
263 data: STANDARD.encode(bytes),
264 mime_type: mime_type.to_owned(),
265 }
266}
267
268fn base64_len(len: usize) -> u64 {
270 len.div_ceil(3) as u64 * 4
271}
272
273fn 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
289pub fn convert_to_png(bytes: &[u8]) -> Option<Vec<u8>> {
292 encode(&decode(bytes)?, ImageOutputFormat::Png)
293}
294
295pub 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
303fn 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
310fn 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 return Some(if size >= 6 && has_exif_header(bytes, data) {
364 data + 6
365 } else {
366 data
367 });
368 }
369 offset = data.checked_add(size)?.checked_add(size % 2)?;
371 }
372 None
373}
374
375fn 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 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
406pub(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 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 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 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 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 const TIFF: &str = "SUkqAAgAAAAJAAABAwABAAAAAgAAAAEBAwABAAAAAgAAAAIBAwADAAAAegAAAAMBAwABAAAAAQAAAAYBAwABAAAAAgAAABEBBAABAAAAgAAAABUBAwABAAAAAwAAABYBAwABAAAAAgAAABcBBAABAAAADAAAAAAAAAAIAAgACAD/AAAA/wAAAP////8=";
576
577 fn dds() -> Vec<u8> {
579 let mut bytes = b"DDS ".to_vec();
580 let mut header = [0u32; 31];
581 header[0] = 124; header[1] = 0x1 | 0x2 | 0x4 | 0x1000; header[2] = 4; header[3] = 4; header[4] = 8; header[18] = 32; header[19] = 0x4; header[20] = u32::from_le_bytes(*b"DXT1");
589 header[26] = 0x1000; 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 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 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}