use std::borrow::Cow;
use std::io::Cursor;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use image::imageops::{self, FilterType};
use image::{DynamicImage, ImageOutputFormat};
use yapi_types::message::{ContentBlock, ImageContent};
use yapi_types::models::ImageResize;
const CONVERT_FAILED: &str =
"[Image omitted: could not be converted to a supported inline image format.]";
const RESIZE_FAILED: &str =
"[Image omitted: could not be resized below the inline image size limit.]";
const DEFAULT_MAX_BYTES: u64 = 4_718_592;
#[derive(Debug, PartialEq)]
pub struct Processed {
pub image: ImageContent,
pub hints: Vec<String>,
}
pub fn process(
bytes: &[u8],
mime_type: &str,
auto_resize: bool,
limits: Option<&ImageResize>,
) -> Result<Processed, &'static str> {
let base = mime_type
.split(';')
.next()
.unwrap_or_default()
.trim()
.to_lowercase();
let (bytes, mime_type, converted_from) = match base.as_str() {
"image/png" => (Cow::Borrowed(bytes), "image/png", None),
"image/jpeg" | "image/jpg" => (Cow::Borrowed(bytes), "image/jpeg", None),
"image/gif" => (Cow::Borrowed(bytes), "image/gif", None),
"image/webp" => (Cow::Borrowed(bytes), "image/webp", None),
_ => {
let png = decode(bytes)
.and_then(|image| encode(&image, ImageOutputFormat::Png))
.ok_or(CONVERT_FAILED)?;
(Cow::Owned(png), "image/png", Some(base))
}
};
let (image, note) = if auto_resize {
resize(&bytes, mime_type, limits).ok_or(RESIZE_FAILED)?
} else {
(inline(&bytes, mime_type), None)
};
let mut hints = Vec::new();
if let Some(from) = converted_from.filter(|from| !from.is_empty() && *from != image.mime_type) {
hints.push(format!(
"[Image converted from {from} to {}.]",
image.mime_type
));
}
hints.extend(note);
Ok(Processed { image, hints })
}
pub fn normalize_prompt(
images: Vec<ImageContent>,
auto_resize: bool,
limits: Option<&ImageResize>,
) -> (Vec<ImageContent>, Vec<String>) {
let mut normalized = Vec::with_capacity(images.len());
let mut hints = Vec::new();
for image in images {
let bytes = STANDARD.decode(&image.data).unwrap_or_default();
match process(&bytes, &image.mime_type, auto_resize, limits) {
Ok(processed) => {
normalized.push(processed.image);
hints.extend(processed.hints);
}
Err(message) => hints.push(message.to_owned()),
}
}
(normalized, hints)
}
pub fn normalize_tool_result(
content: &[ContentBlock],
auto_resize: bool,
limits: Option<&ImageResize>,
) -> Option<Vec<ContentBlock>> {
let mut changed = false;
let mut normalized = Vec::with_capacity(content.len());
for block in content {
let ContentBlock::Image(image) = block else {
normalized.push(block.clone());
continue;
};
let processed = STANDARD
.decode(&image.data)
.ok()
.and_then(|bytes| process(&bytes, &image.mime_type, auto_resize, limits).ok());
match processed {
Some(processed) if processed.image != *image || !processed.hints.is_empty() => {
normalized.push(ContentBlock::Image(processed.image));
if !processed.hints.is_empty() {
normalized.push(ContentBlock::text(processed.hints.join("\n")));
}
changed = true;
}
_ => normalized.push(block.clone()),
}
}
changed.then_some(normalized)
}
fn resize(
bytes: &[u8],
mime_type: &str,
limits: Option<&ImageResize>,
) -> Option<(ImageContent, Option<String>)> {
let resized = resize_image(bytes, mime_type, limits)?;
let note = resized.was_resized.then(|| {
format!(
"[Image: original {}x{}, displayed at {}x{}. Multiply coordinates by {} to map to original image.]",
resized.original_width,
resized.original_height,
resized.width,
resized.height,
yapi_types::js::to_fixed(
f64::from(resized.original_width) / f64::from(resized.width),
2
)
)
});
Some((resized.image, note))
}
#[derive(Debug, PartialEq, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct Resized {
#[serde(flatten)]
pub image: ImageContent,
pub original_width: u32,
pub original_height: u32,
pub width: u32,
pub height: u32,
pub was_resized: bool,
}
pub fn resize_image(
bytes: &[u8],
mime_type: &str,
limits: Option<&ImageResize>,
) -> Option<Resized> {
let limits = limits.cloned().unwrap_or_default();
let max_width = limits.max_width.unwrap_or(2000);
let max_height = limits.max_height.unwrap_or(2000);
let max_bytes = limits.max_bytes.unwrap_or(DEFAULT_MAX_BYTES);
let image = decode(bytes)?;
let (original_width, original_height) = (image.width(), image.height());
let result = |image, width, height, was_resized| Resized {
image,
original_width,
original_height,
width,
height,
was_resized,
};
if original_width <= max_width
&& original_height <= max_height
&& base64_len(bytes.len()) < max_bytes
{
let image = inline(bytes, mime_type);
return Some(result(image, original_width, original_height, false));
}
let scale = |value: u32, max: u32, other: u32| {
yapi_types::js::round(f64::from(value) * f64::from(max) / f64::from(other)) as u32
};
let (mut width, mut height) = (original_width, original_height);
if width > max_width {
height = scale(height, max_width, width);
width = max_width;
}
if height > max_height {
width = scale(width, max_height, height);
height = max_height;
}
let quality = limits.jpeg_quality.unwrap_or(80);
let fallbacks = [85, 70, 55, 40]
.into_iter()
.filter(|other| *other != quality);
let formats: Vec<_> = [ImageOutputFormat::Png, ImageOutputFormat::Jpeg(quality)]
.into_iter()
.chain(fallbacks.map(ImageOutputFormat::Jpeg))
.collect();
loop {
if width == 0 || height == 0 {
return None;
}
let resized = DynamicImage::ImageRgba8(imageops::resize(
&image,
width,
height,
FilterType::Lanczos3,
));
for format in &formats {
let encoded = encode(&resized, format.clone())?;
if base64_len(encoded.len()) < max_bytes {
let mime_type = match format {
ImageOutputFormat::Png => "image/png",
_ => "image/jpeg",
};
return Some(result(inline(&encoded, mime_type), width, height, true));
}
}
let shrink = |value: u32| ((f64::from(value) * 0.75).floor() as u32).max(1);
let next = (shrink(width), shrink(height));
if next == (width, height) {
return None;
}
(width, height) = next;
}
}
fn inline(bytes: &[u8], mime_type: &str) -> ImageContent {
ImageContent {
data: STANDARD.encode(bytes),
mime_type: mime_type.to_owned(),
}
}
fn base64_len(len: usize) -> u64 {
len.div_ceil(3) as u64 * 4
}
fn decode(bytes: &[u8]) -> Option<DynamicImage> {
let image = image::load_from_memory(bytes).ok()?.to_rgba8();
Some(DynamicImage::ImageRgba8(match exif_orientation(bytes) {
2 => imageops::flip_horizontal(&image),
3 => imageops::rotate180(&image),
4 => imageops::flip_vertical(&image),
5 => imageops::flip_horizontal(&imageops::rotate90(&image)),
6 => imageops::rotate90(&image),
7 => imageops::flip_horizontal(&imageops::rotate270(&image)),
8 => imageops::rotate270(&image),
_ => image,
}))
}
pub fn convert_to_png(bytes: &[u8]) -> Option<Vec<u8>> {
encode(&decode(bytes)?, ImageOutputFormat::Png)
}
pub fn to_png(bytes: &[u8]) -> Option<Vec<u8>> {
let image = DynamicImage::ImageRgba8(image::load_from_memory(bytes).ok()?.to_rgba8());
encode(&image, ImageOutputFormat::Png)
}
fn encode(image: &DynamicImage, format: ImageOutputFormat) -> Option<Vec<u8>> {
let mut bytes = Cursor::new(Vec::new());
image.write_to(&mut bytes, format).ok()?;
Some(bytes.into_inner())
}
fn exif_orientation(bytes: &[u8]) -> u16 {
let tiff = if bytes.starts_with(&[0xff, 0xd8]) {
jpeg_tiff_offset(bytes)
} else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
webp_tiff_offset(bytes)
} else {
None
};
tiff.and_then(|start| tiff_orientation(bytes, start))
.unwrap_or(1)
}
fn has_exif_header(bytes: &[u8], offset: usize) -> bool {
bytes.get(offset..offset + 6) == Some(b"Exif\0\0")
}
fn jpeg_tiff_offset(bytes: &[u8]) -> Option<usize> {
let mut offset = 2;
while offset + 1 < bytes.len() {
if bytes[offset] != 0xff {
return None;
}
let marker = bytes[offset + 1];
if marker == 0xff {
offset += 1;
continue;
}
if marker == 0xe1 {
if offset + 10 > bytes.len() {
return None;
}
if has_exif_header(bytes, offset + 4) {
return Some(offset + 10);
}
}
let length = usize::from(u16::from_be_bytes(array(bytes, offset + 2)?));
offset += 2 + length;
}
None
}
fn webp_tiff_offset(bytes: &[u8]) -> Option<usize> {
let mut offset = 12;
while offset + 8 <= bytes.len() {
let size = u32::from_le_bytes(array(bytes, offset + 4)?) as usize;
let data = offset + 8;
if &bytes[offset..offset + 4] == b"EXIF" {
if data.saturating_add(size) > bytes.len() {
return None;
}
return Some(if size >= 6 && has_exif_header(bytes, data) {
data + 6
} else {
data
});
}
offset = data.checked_add(size)?.checked_add(size % 2)?;
}
None
}
fn tiff_orientation(bytes: &[u8], start: usize) -> Option<u16> {
let little_endian = array::<8>(bytes, start)?[..2] == *b"II";
let read16 = |at: usize| {
let pair = array(bytes, at)?;
Some(if little_endian {
u16::from_le_bytes(pair)
} else {
u16::from_be_bytes(pair)
})
};
let read32 = |at: usize| {
let quad = array(bytes, at)?;
Some(if little_endian {
u32::from_le_bytes(quad)
} else {
u32::from_be_bytes(quad)
})
};
let ifd = start.checked_add(read32(start + 4)? as usize)?;
for entry in 0..usize::from(read16(ifd)?) {
let at = ifd + 2 + entry * 12;
array::<12>(bytes, at)?;
if read16(at)? == 0x0112 {
return read16(at + 8).filter(|value| (1..=8).contains(value));
}
}
None
}
pub(crate) fn array<const N: usize>(bytes: &[u8], at: usize) -> Option<[u8; N]> {
bytes.get(at..at.checked_add(N)?)?.try_into().ok()
}
#[cfg(test)]
mod tests {
#![allow(
clippy::unwrap_used,
reason = "test helpers; a panic is a test failure"
)]
use super::*;
use image::{Rgba, RgbaImage};
fn pixels(width: u32, height: u32, noisy: bool) -> DynamicImage {
let mut seed = 7u32;
DynamicImage::ImageRgba8(RgbaImage::from_fn(width, height, |x, y| {
seed = seed.wrapping_mul(1_103_515_245).wrapping_add(12_345);
let noise = if noisy { (seed >> 16) as u8 } else { 0 };
Rgba([(x % 256) as u8 ^ noise, (y % 256) as u8, noise, 255])
}))
}
fn file(image: &DynamicImage, format: ImageOutputFormat) -> Vec<u8> {
encode(image, format).unwrap()
}
fn limits(width: u32, height: u32, bytes: Option<u64>) -> ImageResize {
ImageResize {
max_width: Some(width),
max_height: Some(height),
max_bytes: bytes,
jpeg_quality: None,
}
}
fn dimensions(image: &ImageContent) -> (u32, u32) {
let decoded = image::load_from_memory(&STANDARD.decode(&image.data).unwrap()).unwrap();
(decoded.width(), decoded.height())
}
#[test]
fn small_images_unchanged() {
let png = file(&pixels(30, 20, false), ImageOutputFormat::Png);
let processed = process(&png, "image/png", true, None).unwrap();
assert_eq!(processed.image.data, STANDARD.encode(&png));
assert_eq!(processed.image.mime_type, "image/png");
assert!(processed.hints.is_empty());
}
#[test]
fn large_images_scaled() {
let png = file(&pixels(2100, 10, false), ImageOutputFormat::Png);
let processed = process(&png, "image/png", true, None).unwrap();
assert_eq!(processed.image.mime_type, "image/png");
assert_eq!(dimensions(&processed.image), (2000, 10));
assert_eq!(
processed.hints,
[
"[Image: original 2100x10, displayed at 2000x10. Multiply coordinates by 1.05 to map to original image.]"
]
);
let tall = file(&pixels(30, 90, false), ImageOutputFormat::Png);
let limits = limits(40, 40, None);
let processed = process(&tall, "image/png", true, Some(&limits)).unwrap();
assert_eq!(dimensions(&processed.image), (13, 40));
assert_eq!(
processed.hints,
[
"[Image: original 30x90, displayed at 13x40. Multiply coordinates by 2.31 to map to original image.]"
]
);
let processed = process(&tall, "image/png", false, Some(&limits)).unwrap();
assert_eq!(processed.image.data, STANDARD.encode(&tall));
assert!(processed.hints.is_empty());
}
#[test]
fn jpeg_replaces_large_png() {
let image = pixels(64, 64, true);
let png = file(&image, ImageOutputFormat::Png);
let limits = limits(2000, 2000, Some(base64_len(png.len())));
let processed = process(&png, "image/png", true, Some(&limits)).unwrap();
assert_eq!(processed.image.mime_type, "image/jpeg");
assert_eq!(
STANDARD.decode(&processed.image.data).unwrap(),
file(&image, ImageOutputFormat::Jpeg(80))
);
assert_eq!(
processed.hints,
[
"[Image: original 64x64, displayed at 64x64. Multiply coordinates by 1.00 to map to original image.]"
]
);
}
#[test]
fn shrinks_then_omits() {
let png = file(&pixels(40, 40, true), ImageOutputFormat::Png);
let limits = limits(2000, 2000, Some(1200));
let processed = process(&png, "image/png", true, Some(&limits)).unwrap();
assert!(processed.image.data.len() < 1200);
let (width, height) = dimensions(&processed.image);
assert!(width < 40 && width == height);
let none_fit = ImageResize {
max_bytes: Some(1),
..ImageResize::default()
};
assert_eq!(
process(&png, "image/png", true, Some(&none_fit)),
Err(RESIZE_FAILED)
);
}
#[test]
fn exif_upright_before_resize() {
let jpeg = file(&pixels(4, 2, false), ImageOutputFormat::Jpeg(90));
let mut exif = b"\xff\xe1\x00\x22Exif\0\0MM\0\x2a\0\0\0\x08\0\x01".to_vec();
exif.extend_from_slice(b"\x01\x12\0\x03\0\0\0\x01\0\x06\0\0\0\0\0\0");
let rotated = [&jpeg[..2], &exif, &jpeg[2..]].concat();
assert_eq!(exif_orientation(&rotated), 6);
let limits = limits(3, 3, None);
let processed = process(&rotated, "image/jpeg", true, Some(&limits)).unwrap();
assert_eq!(dimensions(&processed.image), (2, 3));
assert_eq!(
processed.hints,
[
"[Image: original 2x4, displayed at 2x3. Multiply coordinates by 1.00 to map to original image.]"
]
);
}
#[test]
fn to_png_ignores_exif() {
let jpeg = file(&pixels(4, 2, false), ImageOutputFormat::Jpeg(90));
let mut exif = b"\xff\xe1\x00\x22Exif\0\0MM\0\x2a\0\0\0\x08\0\x01".to_vec();
exif.extend_from_slice(b"\x01\x12\0\x03\0\0\0\x01\0\x06\0\0\0\0\0\0");
let rotated = [&jpeg[..2], &exif, &jpeg[2..]].concat();
let png = to_png(&rotated).unwrap();
assert!(png.starts_with(b"\x89PNG"));
let decoded = image::load_from_memory(&png).unwrap();
assert_eq!((decoded.width(), decoded.height()), (4, 2));
}
fn converted(bytes: &[u8], mime_type: &str) -> Vec<u8> {
let mut png = None;
for auto_resize in [true, false] {
let processed = process(bytes, mime_type, auto_resize, None)
.unwrap_or_else(|note| panic!("{mime_type}: {note}"));
assert_eq!(processed.image.mime_type, "image/png");
assert_eq!(
processed.hints,
[format!("[Image converted from {mime_type} to image/png.]")]
);
let data = STANDARD.decode(&processed.image.data).unwrap();
assert!(png.replace(data.clone()).is_none_or(|first| first == data));
}
png.unwrap()
}
const TIFF: &str = "SUkqAAgAAAAJAAABAwABAAAAAgAAAAEBAwABAAAAAgAAAAIBAwADAAAAegAAAAMBAwABAAAAAQAAAAYBAwABAAAAAgAAABEBBAABAAAAgAAAABUBAwABAAAAAwAAABYBAwABAAAAAgAAABcBBAABAAAADAAAAAAAAAAIAAgACAD/AAAA/wAAAP////8=";
fn dds() -> Vec<u8> {
let mut bytes = b"DDS ".to_vec();
let mut header = [0u32; 31];
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");
header[26] = 0x1000; bytes.extend(header.iter().flat_map(|word| word.to_le_bytes()));
bytes.extend([0x00, 0xf8, 0x00, 0x00, 0, 0, 0, 0]);
bytes
}
#[test]
fn other_formats_become_png() {
let image = pixels(3, 2, false);
let png = file(&image, ImageOutputFormat::Png);
let rgba16 = DynamicImage::ImageRgba16(image.to_rgba16());
let rgb = DynamicImage::ImageRgb8(image.to_rgb8());
let pnm = ImageOutputFormat::Pnm(image::codecs::pnm::PnmSubtype::Pixmap(
image::codecs::pnm::SampleEncoding::Binary,
));
for (mime_type, bytes) in [
("image/bmp", file(&image, ImageOutputFormat::Bmp)),
("image/x-icon", file(&image, ImageOutputFormat::Ico)),
("image/x-portable-pixmap", file(&rgb, pnm)),
(
"image/x-farbfeld",
file(&rgba16, ImageOutputFormat::Farbfeld),
),
] {
assert_eq!(converted(&bytes, mime_type), png, "{mime_type}");
}
let tiff = RgbaImage::from_fn(2, 2, |x, y| match (x, y) {
(0, 0) => Rgba([255, 0, 0, 255]),
(1, 0) => Rgba([0, 255, 0, 255]),
(0, 1) => Rgba([0, 0, 255, 255]),
_ => Rgba([255; 4]),
});
assert_eq!(
converted(&STANDARD.decode(TIFF).unwrap(), "image/tiff"),
file(&DynamicImage::ImageRgba8(tiff), ImageOutputFormat::Png)
);
let red = DynamicImage::ImageRgba8(RgbaImage::from_pixel(4, 4, Rgba([255, 0, 0, 255])));
assert_eq!(
converted(&dds(), "image/vnd.ms-dds"),
file(&red, ImageOutputFormat::Png)
);
let mut hdr = Vec::new();
image::codecs::hdr::HdrEncoder::new(&mut hdr)
.encode(&[image::Rgb([0.5f32, 0.25, 1.0]); 6], 3, 2)
.unwrap();
let decoded = image::load_from_memory(&converted(&hdr, "image/vnd.radiance")).unwrap();
assert_eq!((decoded.width(), decoded.height()), (3, 2));
assert_eq!(
process(b"not an image", "image/tiff", true, None),
Err(CONVERT_FAILED)
);
assert_eq!(
process(b"not an image", "image/png", true, None),
Err(RESIZE_FAILED)
);
}
#[test]
fn normalizes_prompts_and_results() {
let png = ImageContent {
data: STANDARD.encode(file(&pixels(3, 2, false), ImageOutputFormat::Png)),
mime_type: "image/png".into(),
};
let bmp = ImageContent {
data: STANDARD.encode(file(&pixels(3, 2, false), ImageOutputFormat::Bmp)),
mime_type: "image/bmp".into(),
};
let broken = ImageContent {
data: "AAAA".into(),
mime_type: "image/png".into(),
};
let (images, hints) =
normalize_prompt(vec![png.clone(), bmp.clone(), broken.clone()], true, None);
assert_eq!(images[0], png);
assert_eq!(images[1].mime_type, "image/png");
assert_eq!(images.len(), 2);
assert_eq!(
hints,
[
"[Image converted from image/bmp to image/png.]",
RESIZE_FAILED
]
);
let unchanged = [ContentBlock::text("a"), ContentBlock::Image(png.clone())];
assert_eq!(normalize_tool_result(&unchanged, true, None), None);
let content = [
ContentBlock::Image(bmp),
ContentBlock::Image(broken.clone()),
ContentBlock::text("a"),
];
let normalized = normalize_tool_result(&content, true, None).unwrap();
assert_eq!(
normalized,
[
ContentBlock::Image(images[1].clone()),
ContentBlock::text("[Image converted from image/bmp to image/png.]"),
ContentBlock::Image(broken),
ContentBlock::text("a"),
]
);
}
}