llm_browser_testkit/
vision.rs

1//! Vision support — screenshot capture, tiling, downscaling, and JPEG
2//! encoding for LLM visual assertions.
3//!
4//! When an `assert` step sets `screenshot = true`, the runner captures the
5//! **full scrollable page** (via the CDP `captureBeyondViewport` flag —
6//! nothing below the fold is skipped) and splits it into viewport-tall
7//! bands ("tiles") from the top, covering at most
8//! `scenario::ScenarioConfig::screenshot_max_height` (default `"20x"` =
9//! four viewports). Each tile is downscaled independently so its longest
10//! edge is at most [`default_max_dimension`] (the page's `[config]
11//! screenshot_max_dimension` wins when set), JPEG-encoded, and sent as its
12//! own OpenAI-compatible `image_url` content part next to the text prompt —
13//! detail is preserved at every depth, and the tile count (hence token
14//! cost) is bounded by the height cap and [`max_tiles`].
15//!
16//! Downscaling and cropping happen in Rust (no page JS, no fragile
17//! `canvas` evaluation): the PNG bytes are decoded with the `image` crate,
18//! banded, resized with Lanczos filtering, and re-encoded as quality-85
19//! JPEG in a single capture call.
20
21use base64::engine::general_purpose::STANDARD;
22use base64::Engine as _;
23use headless_chrome::protocol::cdp::Page::CaptureScreenshot;
24use headless_chrome::protocol::cdp::Page::CaptureScreenshotFormatOption;
25use headless_chrome::Tab;
26use image::DynamicImage;
27
28/// Default longest edge (px) of screenshots sent to vision endpoints.
29pub const DEFAULT_MAX_DIMENSION: u32 = 1400;
30/// Hard ceiling on the number of tiles attached to a single vision assert,
31/// so a huge `screenshot_max_height` can never produce an unbounded
32/// request (or token bill).
33pub const MAX_TILES: u32 = 30;
34/// JPEG quality used for the encoded screenshot.
35const JPEG_QUALITY: u8 = 85;
36
37/// Captures the full page and returns one JPEG data URL per viewport-tall
38/// band, ordered from the top of the page down — suitable for the
39/// OpenAI-compatible `image_url` content-part array.
40///
41/// The capture is covered from the top to `height_cap` pixels (pages
42/// shorter than the cap or the viewport are returned as a single tile).
43/// Each band is downscaled independently so its longest edge is at most
44/// `max_dimension` (no upscaling; `0` disables resizing). The number of
45/// tiles is bounded by [`max_tiles`].
46///
47/// # Errors
48///
49/// Returns a description when the CDP screenshot capture, PNG decode,
50/// crop, resize, or JPEG encode fails.
51pub fn capture_screenshot_data_urls(
52    tab: &Tab,
53    max_dimension: u32,
54    height_cap: u32,
55    band_height: u32,
56) -> Result<Vec<String>, String> {
57    let png = tab
58        .call_method(CaptureScreenshot {
59            format: Some(CaptureScreenshotFormatOption::Png),
60            quality: None,
61            clip: None,
62            from_surface: Some(true),
63            capture_beyond_viewport: Some(true),
64            optimize_for_speed: None,
65        })
66        .map_err(|e| format!("screenshot capture failed: {e}"))?
67        .data;
68    let png_bytes = STANDARD
69        .decode(png)
70        .map_err(|e| format!("screenshot base64 decode failed: {e}"))?;
71
72    let img = image::load_from_memory(&png_bytes)
73        .map_err(|e| format!("screenshot decode failed: {e}"))?;
74
75    // Coverage from the top: the height cap bounds token cost, tile bands
76    // keep 1:1 detail for everything that is covered.
77    let coverage = img.height().min(height_cap);
78    let count = tile_count(coverage, band_height);
79    let mut tiles: Vec<String> = Vec::new();
80    for i in 0..count {
81        let top = i * band_height;
82        let band_h = (coverage - top).min(band_height);
83        let band = img.crop_imm(0, top, img.width(), band_h);
84        tiles.push(encode_jpeg(band, max_dimension)?);
85    }
86    Ok(tiles)
87}
88
89/// Number of viewport-tall bands covering `coverage` pixels, clamped to
90/// [`max_tiles`] so a single assert can never send an unbounded number of
91/// image parts.
92#[must_use]
93pub fn tile_count(coverage: u32, band_height: u32) -> u32 {
94    if coverage == 0 {
95        return 0;
96    }
97    if band_height == 0 {
98        return 1;
99    }
100    coverage.div_ceil(band_height).min(MAX_TILES)
101}
102
103/// Downscales (optional) and JPEG-encodes a single band into a data URL.
104fn encode_jpeg(img: DynamicImage, max_dimension: u32) -> Result<String, String> {
105    let (width, height) = (img.width(), img.height());
106    let longest = width.max(height);
107    let resized = if max_dimension > 0 && longest > max_dimension {
108        let scale = f64::from(max_dimension) / f64::from(longest);
109        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
110        let new_width = (f64::from(width) * scale).round().max(1.0) as u32;
111        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
112        let new_height = (f64::from(height) * scale).round().max(1.0) as u32;
113        img.resize(new_width, new_height, image::imageops::FilterType::Lanczos3)
114    } else {
115        img
116    };
117
118    let mut jpeg = Vec::new();
119    {
120        let mut encoder =
121            image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg, JPEG_QUALITY);
122        encoder
123            .encode_image(&resized)
124            .map_err(|e| format!("screenshot encode failed: {e}"))?;
125    }
126
127    Ok(format!("data:image/jpeg;base64,{}", STANDARD.encode(jpeg)))
128}
129
130#[cfg(test)]
131mod tests {
132    use super::{tile_count, DEFAULT_MAX_DIMENSION, MAX_TILES};
133
134    #[test]
135    fn default_max_dimension_sane() {
136        // 1280x720 viewports stay untouched; 1920-wide screens are scaled
137        // down to keep vision tokens reasonable.
138        const {
139            assert!(DEFAULT_MAX_DIMENSION >= 1280 && DEFAULT_MAX_DIMENSION <= 1600);
140        }
141    }
142
143    #[test]
144    fn tile_count_bands_and_clamps() {
145        assert_eq!(tile_count(0, 720), 0);
146        assert_eq!(tile_count(720, 720), 1);
147        assert_eq!(tile_count(1440, 720), 2);
148        assert_eq!(tile_count(1500, 720), 3);
149        assert_eq!(tile_count(2880, 720), 4);
150        assert_eq!(tile_count(720, 0), 1);
151        assert_eq!(tile_count(MAX_TILES * 720 + 1, 720), MAX_TILES);
152    }
153}
llm_browser_testkit/vision.rs

llm_browser_testkit/
vision.rs
