tenant.openagents/omega

1use std::io::{Cursor, Write};
2use std::sync::Arc;
3
4use anyhow::{Result, anyhow};
5use base64::{Engine as _, write::EncoderWriter};
6use gpui::{
7    App, AppContext as _, DevicePixels, Image, ImageFormat, ObjectFit, Size, Task, point, px, size,
8};
9use gpui_util::ResultExt;
10use image::GenericImageView as _;
11use image::codecs::png::PngEncoder;
12
13use language_model_core::{ImageSize, LanguageModelImage};
14
15/// Anthropic wants uploaded images to be smaller than this in both dimensions.
16const ANTHROPIC_SIZE_LIMIT: f32 = 1568.;
17
18/// Default per-image hard limit (in bytes) for the encoded image payload we send upstream.
19///
20/// NOTE: `LanguageModelImage.source` is base64-encoded PNG bytes (without the `data:` prefix).
21/// This limit is enforced on the encoded PNG bytes *before* base64 encoding.
22const DEFAULT_IMAGE_MAX_BYTES: usize = 5 * 1024 * 1024;
23
24/// Conservative cap on how many times we'll attempt to shrink/re-encode an image to fit
25/// `DEFAULT_IMAGE_MAX_BYTES`.
26const MAX_IMAGE_DOWNSCALE_PASSES: usize = 8;
27
28/// Extension trait for `LanguageModelImage` that provides GPUI-dependent functionality.
29pub trait LanguageModelImageExt {
30    const FORMAT: ImageFormat;
31    fn from_image(data: Arc<Image>, cx: &mut App) -> Task<Option<LanguageModelImage>>;
32    fn from_base64_image(data: &str, mime_type: &str) -> Result<Option<LanguageModelImage>>;
33}
34
35impl LanguageModelImageExt for LanguageModelImage {
36    const FORMAT: ImageFormat = ImageFormat::Png;
37
38    fn from_image(data: Arc<Image>, cx: &mut App) -> Task<Option<LanguageModelImage>> {
39        cx.background_spawn(async move {
40            let format = match data.format() {
41                ImageFormat::Png => image::ImageFormat::Png,
42                ImageFormat::Jpeg => image::ImageFormat::Jpeg,
43                ImageFormat::Webp => image::ImageFormat::WebP,
44                ImageFormat::Gif => image::ImageFormat::Gif,
45                ImageFormat::Bmp => image::ImageFormat::Bmp,
46                ImageFormat::Tiff => image::ImageFormat::Tiff,
47                ImageFormat::Ico => image::ImageFormat::Ico,
48                ImageFormat::Pnm => image::ImageFormat::Pnm,
49                ImageFormat::Svg => return None,
50            };
51            let dynamic_image =
52                image::load_from_memory_with_format(data.bytes(), format).log_err()?;
53            language_model_image_from_dynamic_image(dynamic_image)
54                .log_err()
55                .flatten()
56        })
57    }
58
59    fn from_base64_image(data: &str, mime_type: &str) -> Result<Option<LanguageModelImage>> {
60        let format = image::ImageFormat::from_mime_type(mime_type)
61            .ok_or_else(|| anyhow!("unsupported image MIME type `{}`", mime_type))?;
62        let bytes = base64::engine::general_purpose::STANDARD.decode(data.as_bytes())?;
63        let dynamic_image = image::load_from_memory_with_format(&bytes, format)?;
64        language_model_image_from_dynamic_image(dynamic_image)
65    }
66}
67
68fn language_model_image_from_dynamic_image(
69    dynamic_image: image::DynamicImage,
70) -> Result<Option<LanguageModelImage>> {
71    let width = dynamic_image.width();
72    let height = dynamic_image.height();
73    let image_size = size(DevicePixels(width as i32), DevicePixels(height as i32));
74
75    // First apply any provider-specific dimension constraints we know about (Anthropic).
76    let mut processed_image = if image_size.width.0 > ANTHROPIC_SIZE_LIMIT as i32
77        || image_size.height.0 > ANTHROPIC_SIZE_LIMIT as i32
78    {
79        let new_bounds = ObjectFit::ScaleDown.get_bounds(
80            gpui::Bounds {
81                origin: point(px(0.0), px(0.0)),
82                size: size(px(ANTHROPIC_SIZE_LIMIT), px(ANTHROPIC_SIZE_LIMIT)),
83            },
84            image_size,
85        );
86        dynamic_image.resize(
87            new_bounds.size.width.into(),
88            new_bounds.size.height.into(),
89            image::imageops::FilterType::Triangle,
90        )
91    } else {
92        dynamic_image
93    };
94
95    // Then enforce a default per-image size cap on the encoded PNG bytes.
96    //
97    // We always send PNG bytes (either original PNG bytes, or re-encoded PNG) base64'd.
98    // The upstream provider limit we want to respect is effectively on the binary image
99    // payload size, so we enforce against the encoded PNG bytes before base64 encoding.
100    let mut encoded_png = encode_png_bytes(&processed_image)?;
101    for _pass in 0..MAX_IMAGE_DOWNSCALE_PASSES {
102        if encoded_png.len() <= DEFAULT_IMAGE_MAX_BYTES {
103            break;
104        }
105
106        // Scale down geometrically to converge quickly. We don't know the final PNG size
107        // as a function of pixels, so we iteratively shrink.
108        let (width, height) = processed_image.dimensions();
109        if width <= 1 || height <= 1 {
110            break;
111        }
112
113        // Shrink by ~15% each pass (0.85). This is a compromise between speed and
114        // preserving image detail.
115        let new_width = ((width as f32) * 0.85).round().max(1.0) as u32;
116        let new_height = ((height as f32) * 0.85).round().max(1.0) as u32;
117
118        processed_image =
119            processed_image.resize(new_width, new_height, image::imageops::FilterType::Triangle);
120        encoded_png = encode_png_bytes(&processed_image)?;
121    }
122
123    if encoded_png.len() > DEFAULT_IMAGE_MAX_BYTES {
124        // Still too large after multiple passes; treat as non-convertible for now.
125        // (Provider-specific handling can be introduced later.)
126        return Ok(None);
127    }
128
129    // Now base64 encode the PNG bytes.
130    let base64_image = encode_bytes_as_base64(encoded_png.as_slice())?;
131
132    // SAFETY: The base64 encoder should not produce non-UTF8.
133    let source = unsafe { String::from_utf8_unchecked(base64_image) };
134
135    Ok(Some(LanguageModelImage {
136        source: source.into(),
137    }))
138}
139
140fn encode_png_bytes(image: &image::DynamicImage) -> Result<Vec<u8>> {
141    let mut png = Vec::new();
142    image.write_with_encoder(PngEncoder::new(&mut png))?;
143    Ok(png)
144}
145
146fn encode_bytes_as_base64(bytes: &[u8]) -> Result<Vec<u8>> {
147    let mut base64_image = Vec::new();
148    {
149        let mut base64_encoder = EncoderWriter::new(
150            Cursor::new(&mut base64_image),
151            &base64::engine::general_purpose::STANDARD,
152        );
153        base64_encoder.write_all(bytes)?;
154    }
155    Ok(base64_image)
156}
157
158/// Convert a core `ImageSize` to a gpui `Size<DevicePixels>`.
159pub fn image_size_to_gpui(size: ImageSize) -> Size<DevicePixels> {
160    Size {
161        width: DevicePixels(size.width),
162        height: DevicePixels(size.height),
163    }
164}
165
166/// Convert a gpui `Size<DevicePixels>` to a core `ImageSize`.
167pub fn gpui_size_to_image_size(size: Size<DevicePixels>) -> ImageSize {
168    ImageSize {
169        width: size.width.0,
170        height: size.height.0,
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use gpui::TestAppContext;
178
179    fn base64_to_png_bytes(base64: &str) -> Vec<u8> {
180        base64::engine::general_purpose::STANDARD
181            .decode(base64)
182            .expect("valid base64")
183    }
184
185    fn png_dimensions(png_bytes: &[u8]) -> (u32, u32) {
186        let img = image::load_from_memory(png_bytes).expect("valid png");
187        (img.width(), img.height())
188    }
189
190    fn make_noisy_png_bytes(width: u32, height: u32) -> Vec<u8> {
191        use image::{ImageBuffer, Rgba};
192        use std::hash::{Hash, Hasher};
193
194        let img = ImageBuffer::from_fn(width, height, |x, y| {
195            let mut hasher = std::hash::DefaultHasher::new();
196            (x, y, width, height).hash(&mut hasher);
197            let h = hasher.finish();
198            Rgba([h as u8, (h >> 8) as u8, (h >> 16) as u8, 255])
199        });
200
201        let mut buf = Cursor::new(Vec::new());
202        img.write_with_encoder(PngEncoder::new(&mut buf))
203            .expect("encode");
204        buf.into_inner()
205    }
206
207    #[gpui::test]
208    async fn test_from_image_downscales_to_default_5mb_limit(cx: &mut TestAppContext) {
209        let raw_png = make_noisy_png_bytes(4096, 4096);
210        assert!(
211            raw_png.len() > DEFAULT_IMAGE_MAX_BYTES,
212            "Test image should exceed the 5 MB limit (actual: {} bytes)",
213            raw_png.len()
214        );
215
216        let image = Arc::new(gpui::Image::from_bytes(ImageFormat::Png, raw_png.clone()));
217        let lm_image = cx
218            .update(|cx| LanguageModelImage::from_image(Arc::clone(&image), cx))
219            .await
220            .expect("from_image should succeed");
221
222        assert_downscaled_from_original(lm_image.source.as_ref(), 4096, 4096);
223
224        let base64_png = base64::engine::general_purpose::STANDARD.encode(raw_png);
225        let lm_image = LanguageModelImage::from_base64_image(&base64_png, "image/png")
226            .expect("from_base64_image should not error")
227            .expect("from_base64_image should succeed");
228
229        assert_downscaled_from_original(lm_image.source.as_ref(), 4096, 4096);
230    }
231
232    #[test]
233    fn test_from_base64_image_converts_jpeg_to_png() {
234        use image::ImageEncoder as _;
235
236        let mut jpeg_bytes = Vec::new();
237        image::codecs::jpeg::JpegEncoder::new(&mut jpeg_bytes)
238            .write_image(&[255, 0, 0], 1, 1, image::ExtendedColorType::Rgb8)
239            .expect("encode jpeg");
240        let jpeg_data = base64::engine::general_purpose::STANDARD.encode(jpeg_bytes);
241
242        let image = LanguageModelImage::from_base64_image(&jpeg_data, "image/jpeg")
243            .expect("from_base64_image should not error")
244            .expect("from_base64_image should succeed");
245        let png_bytes = base64_to_png_bytes(image.source.as_ref());
246
247        assert_eq!(
248            image::guess_format(&png_bytes).expect("guess image format"),
249            image::ImageFormat::Png
250        );
251        assert_eq!(png_dimensions(&png_bytes), (1, 1));
252    }
253
254    fn assert_downscaled_from_original(base64_png: &str, width: u32, height: u32) {
255        let decoded_png = base64_to_png_bytes(base64_png);
256        assert!(
257            decoded_png.len() <= DEFAULT_IMAGE_MAX_BYTES,
258            "Encoded PNG should be ≤ {} bytes after downscale, but was {} bytes",
259            DEFAULT_IMAGE_MAX_BYTES,
260            decoded_png.len()
261        );
262
263        let (downsized_width, downsized_height) = png_dimensions(&decoded_png);
264        assert!(
265            downsized_width < width && downsized_height < height,
266            "Dimensions should have shrunk: got {}×{}",
267            downsized_width,
268            downsized_height
269        );
270    }
271}
272