2ee77bc867
0.87.14 followup. The 0.87.14 test used MissedTickBehavior::Burst —
any post-prep yield replayed missed ticks and could fake the >= 2
threshold. Switch to Skip and assert on `ticks_post - ticks_pre >= 50`,
which sits well above the inlined-prep ceiling (~5) and far below
the spawn_blocking floor (~700+). Mutation-confirmed.
Also corrects: prepare_analysis doc count ("five fallback paths" → 4).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1302 lines
49 KiB
Rust
1302 lines
49 KiB
Rust
//! The OpenAI-compatible vision HTTP client + the long-page slicing wrapper.
|
||
//!
|
||
//! [`VisionClient::analyze`] decides per page:
|
||
//! * **Normal-aspect page** → one combined call (OCR + tags + scene + safety).
|
||
//! * **Tall page** → slice it into budget-sized bands at native resolution,
|
||
//! OCR each (Pass A), merge with seam-dedup, then one grounding call that
|
||
//! gets the whole (downscaled) image **plus** the merged OCR text and
|
||
//! returns tags + scene + safety (Pass B). Only OCR is merged.
|
||
//!
|
||
//! The geometry ([`plan_slices`]), merge ([`merge_ocr`]) and request builders
|
||
//! are pure so they're unit-tested without a live model; only the `post_chat`
|
||
//! path does I/O.
|
||
|
||
use std::collections::HashSet;
|
||
use std::io::Cursor;
|
||
|
||
use anyhow::{anyhow, Context};
|
||
use base64::Engine;
|
||
use image::imageops::FilterType;
|
||
use image::DynamicImage;
|
||
use serde_json::json;
|
||
|
||
use crate::analysis::prompt::{
|
||
self, MAX_OCR_PIECES, MAX_OCR_TEXT_CHARS, MAX_SCENE_CHARS, MAX_TAGS,
|
||
};
|
||
use crate::config::{AnalysisConfig, ResponseFormat};
|
||
use crate::domain::page_analysis::{ContentWarning, OcrResult, VisionAnalysis};
|
||
|
||
/// Vision client built from [`AnalysisConfig`]. Cheap to clone (holds a
|
||
/// `reqwest::Client`, which is internally `Arc`-backed).
|
||
#[derive(Clone)]
|
||
pub struct VisionClient {
|
||
http: reqwest::Client,
|
||
endpoint: String,
|
||
model: String,
|
||
api_key: Option<String>,
|
||
max_tokens: u32,
|
||
response_format: ResponseFormat,
|
||
frequency_penalty: f64,
|
||
temperature: f64,
|
||
/// System prompt for the single-call (normal-aspect) path.
|
||
system_prompt: String,
|
||
/// Pass-A OCR-only prompt (tall-page slices).
|
||
ocr_prompt: String,
|
||
/// Pass-B grounding prompt (tags/scene/safety).
|
||
grounding_prompt: String,
|
||
slice: SliceParams,
|
||
}
|
||
|
||
/// Pure geometry inputs, mirrored from [`AnalysisConfig`].
|
||
#[derive(Clone, Copy, Debug)]
|
||
struct SliceParams {
|
||
max_pixels: u32,
|
||
min_slice_height: u32,
|
||
overlap: f64,
|
||
tall_threshold: f64,
|
||
max_slices: usize,
|
||
}
|
||
|
||
/// Result of the (blocking-pool) prep pass: a self-contained set of byte
|
||
/// payloads the async HTTP loop can dispatch without further CPU work.
|
||
/// The enum mirrors [`Plan`] but carries actual JPEGs instead of geometry.
|
||
enum PreparedAnalysis {
|
||
/// `image::load_from_memory` failed — fall back to sending the raw bytes.
|
||
Undecodable,
|
||
/// One vision call against `jpeg` (the encoded whole page).
|
||
Single { jpeg: Vec<u8> },
|
||
/// Per-band OCR calls plus a final whole-image grounding call. Each
|
||
/// slice tuple is `(y0, y1, encoded_jpeg)`.
|
||
Sliced {
|
||
slices: Vec<(u32, u32, Vec<u8>)>,
|
||
whole: Vec<u8>,
|
||
},
|
||
}
|
||
|
||
/// CPU-only image work: decode → plan → optional width-reduce → slice →
|
||
/// JPEG encode. Returns `Ok(Undecodable)` (never `Err`) when the bytes
|
||
/// don't parse, so the caller can fall through to the raw-pass-through
|
||
/// behaviour.
|
||
///
|
||
/// **Diagnosability.** The four fallback paths below (decode failure,
|
||
/// `render_whole` on `Single`, `render_slice` mid-loop, `render_whole`
|
||
/// on `Sliced`) all return `Undecodable`, which the caller treats as
|
||
/// "send raw bytes". That's correct semantically — a server can
|
||
/// sometimes salvage a broken page — but it also makes a silent
|
||
/// JPEG-encoder regression indistinguishable from "the page was just
|
||
/// garbage." Emit a `warn` on each fallback so an operator can grep
|
||
/// for "vision prep fell through" and tell the two apart.
|
||
fn prepare_analysis(image: &[u8], params: SliceParams) -> PreparedAnalysis {
|
||
let Some(img) = image::load_from_memory(image).ok() else {
|
||
tracing::warn!(
|
||
bytes = image.len(),
|
||
"vision prep fell through to Undecodable: image::load_from_memory failed"
|
||
);
|
||
return PreparedAnalysis::Undecodable;
|
||
};
|
||
match plan_slices(img.width(), img.height(), ¶ms) {
|
||
Plan::Single { .. } => match render_whole(&img, params.max_pixels) {
|
||
Some(jpeg) => PreparedAnalysis::Single { jpeg },
|
||
None => {
|
||
tracing::warn!(
|
||
width = img.width(),
|
||
height = img.height(),
|
||
"vision prep fell through to Undecodable: render_whole on Single plan failed"
|
||
);
|
||
PreparedAnalysis::Undecodable
|
||
}
|
||
},
|
||
Plan::Sliced { width, height, bands } => {
|
||
let work = if width == img.width() && height == img.height() {
|
||
img.clone()
|
||
} else {
|
||
img.resize_exact(width, height, FilterType::Triangle)
|
||
};
|
||
let mut slices = Vec::with_capacity(bands.len());
|
||
for (y0, y1) in &bands {
|
||
let Some(jpeg) = render_slice(&work, *y0, *y1, params.max_pixels) else {
|
||
// A failed slice falls back to the undecodable path —
|
||
// the server then gets one combined call instead of N
|
||
// broken slice calls.
|
||
tracing::warn!(
|
||
y0 = *y0,
|
||
y1 = *y1,
|
||
width,
|
||
height,
|
||
"vision prep fell through to Undecodable: render_slice failed"
|
||
);
|
||
return PreparedAnalysis::Undecodable;
|
||
};
|
||
slices.push((*y0, *y1, jpeg));
|
||
}
|
||
let Some(whole) = render_whole(&img, params.max_pixels) else {
|
||
tracing::warn!(
|
||
width = img.width(),
|
||
height = img.height(),
|
||
"vision prep fell through to Undecodable: render_whole on Sliced plan failed"
|
||
);
|
||
return PreparedAnalysis::Undecodable;
|
||
};
|
||
PreparedAnalysis::Sliced { slices, whole }
|
||
}
|
||
}
|
||
}
|
||
|
||
impl VisionClient {
|
||
pub fn new(http: reqwest::Client, cfg: &AnalysisConfig) -> Self {
|
||
Self {
|
||
http,
|
||
endpoint: cfg.endpoint.clone(),
|
||
model: cfg.model.clone(),
|
||
api_key: cfg.api_key.clone(),
|
||
max_tokens: cfg.max_tokens,
|
||
response_format: cfg.response_format,
|
||
frequency_penalty: cfg.frequency_penalty,
|
||
temperature: cfg.temperature,
|
||
system_prompt: cfg.system_prompt.clone(),
|
||
ocr_prompt: cfg.ocr_prompt.clone(),
|
||
grounding_prompt: cfg.grounding_prompt.clone(),
|
||
slice: SliceParams {
|
||
max_pixels: cfg.max_pixels,
|
||
min_slice_height: cfg.min_slice_height,
|
||
overlap: cfg.slice_overlap,
|
||
tall_threshold: cfg.tall_aspect_threshold,
|
||
max_slices: cfg.max_slices,
|
||
},
|
||
}
|
||
}
|
||
|
||
/// Analyze one page image. `mime` is the stored content type (used only
|
||
/// for the fallback when the image can't be decoded locally).
|
||
pub async fn analyze(&self, image: &[u8], mime: &str) -> anyhow::Result<VisionAnalysis> {
|
||
// All the CPU-heavy work — JPEG/PNG decode, optional width reduce,
|
||
// per-band slice, JPEG re-encode — runs on the blocking pool so it
|
||
// doesn't starve the tokio runtime (axum handlers, SSE streams,
|
||
// other daemons share the same threads). A single hop pays the
|
||
// overhead once per page; the per-call work stays on the blocking
|
||
// worker until the HTTP loop below picks back up.
|
||
let prepared = {
|
||
let bytes = image.to_vec();
|
||
let params = self.slice;
|
||
tokio::task::spawn_blocking(move || prepare_analysis(&bytes, params))
|
||
.await
|
||
.map_err(|e| anyhow!("vision prep join: {e}"))?
|
||
};
|
||
|
||
match prepared {
|
||
PreparedAnalysis::Undecodable => {
|
||
// Send raw bytes in a single combined call and let the server
|
||
// cope (preserves prior behavior). Base64 encoding still runs
|
||
// on the runtime; for an undecodable page this is the
|
||
// happy path's exit and stays brief.
|
||
let url = format!("data:{mime};base64,{}", b64(image));
|
||
let body = build_request_body(
|
||
&self.model,
|
||
self.max_tokens,
|
||
&url,
|
||
self.response_format,
|
||
self.frequency_penalty,
|
||
self.temperature,
|
||
&self.system_prompt,
|
||
);
|
||
parse_chat_completion(&self.post_chat(body).await?)
|
||
}
|
||
PreparedAnalysis::Single { jpeg } => {
|
||
let body = build_request_body(
|
||
&self.model,
|
||
self.max_tokens,
|
||
&data_url(&jpeg),
|
||
self.response_format,
|
||
self.frequency_penalty,
|
||
self.temperature,
|
||
&self.system_prompt,
|
||
);
|
||
parse_chat_completion(&self.post_chat(body).await?)
|
||
}
|
||
PreparedAnalysis::Sliced { slices, whole } => {
|
||
tracing::debug!(
|
||
bands = slices.len(),
|
||
"analysis: slicing tall page"
|
||
);
|
||
// Pass A: OCR each band (keep its y-range for seam dedup).
|
||
let mut ocrs: Vec<SliceOcr> = Vec::with_capacity(slices.len());
|
||
for (y0, y1, jpeg) in &slices {
|
||
let body = build_ocr_body(
|
||
&self.model,
|
||
self.max_tokens,
|
||
self.response_format,
|
||
self.frequency_penalty,
|
||
self.temperature,
|
||
&self.ocr_prompt,
|
||
&data_url(jpeg),
|
||
);
|
||
let parsed = parse_chat_completion(&self.post_chat(body).await?)?;
|
||
ocrs.push(SliceOcr {
|
||
y0: *y0 as f64,
|
||
y1: *y1 as f64,
|
||
pieces: parsed.ocr_results,
|
||
});
|
||
}
|
||
let merged = merge_ocr(ocrs);
|
||
|
||
// Pass B: ground tags/scene/safety on the whole image + OCR.
|
||
let ocr_text = merged
|
||
.iter()
|
||
.map(|o| format!("[{}] {}", o.kind, o.text))
|
||
.collect::<Vec<_>>()
|
||
.join("\n");
|
||
let body = build_grounding_body(
|
||
&self.model,
|
||
self.max_tokens,
|
||
self.response_format,
|
||
self.frequency_penalty,
|
||
self.temperature,
|
||
&self.grounding_prompt,
|
||
&data_url(&whole),
|
||
&ocr_text,
|
||
);
|
||
let g = parse_chat_completion(&self.post_chat(body).await?)?;
|
||
|
||
Ok(sanitize(VisionAnalysis {
|
||
ocr_results: merged,
|
||
tagging_results: g.tagging_results,
|
||
scene_description: g.scene_description,
|
||
safety_flag: g.safety_flag,
|
||
}))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// POST a chat/completions body and return the parsed JSON value,
|
||
/// surfacing the server's error body on a non-2xx (different
|
||
/// OpenAI-compatible servers reject different fields).
|
||
async fn post_chat(&self, body: serde_json::Value) -> anyhow::Result<serde_json::Value> {
|
||
let mut req = self.http.post(&self.endpoint).json(&body);
|
||
if let Some(key) = &self.api_key {
|
||
req = req.bearer_auth(key);
|
||
}
|
||
let resp = req.send().await.context("vision request failed")?;
|
||
let status = resp.status();
|
||
if !status.is_success() {
|
||
let body = resp.text().await.unwrap_or_default();
|
||
anyhow::bail!(
|
||
"vision endpoint returned {status}: {}",
|
||
body.trim().chars().take(800).collect::<String>()
|
||
);
|
||
}
|
||
resp.json().await.context("vision response was not JSON")
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Geometry
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[derive(Debug, PartialEq, Eq)]
|
||
enum Plan {
|
||
/// Whole-page single call (dims are the working dims, informational).
|
||
Single { width: u32, height: u32 },
|
||
/// Slice the page into `bands` (y0,y1) over a `width × height` working
|
||
/// image; `width` is reduced from native only for pathologically wide
|
||
/// pages.
|
||
Sliced {
|
||
width: u32,
|
||
height: u32,
|
||
bands: Vec<(u32, u32)>,
|
||
},
|
||
}
|
||
|
||
/// Plan how to feed a `w×h` page to the model: one call, or N budget-sized
|
||
/// vertical slices at native resolution (portrait *or* landscape). See the
|
||
/// module/plan docs for the rules.
|
||
fn plan_slices(w: u32, h: u32, p: &SliceParams) -> Plan {
|
||
if w == 0 || h == 0 {
|
||
return Plan::Single { width: w, height: h };
|
||
}
|
||
let budget = p.max_pixels.max(1) as f64;
|
||
let overlap = p.overlap.clamp(0.0, 0.9);
|
||
|
||
// Aspect guard: a page wider than this would yield hair-thin strips, so
|
||
// (and only then) reduce the working width.
|
||
let max_slice_width = (budget / p.min_slice_height.max(1) as f64).floor().max(1.0);
|
||
let (work_w, work_h) = if (w as f64) > max_slice_width {
|
||
let s = max_slice_width / w as f64;
|
||
(
|
||
max_slice_width as u32,
|
||
((h as f64) * s).round().max(1.0) as u32,
|
||
)
|
||
} else {
|
||
(w, h)
|
||
};
|
||
|
||
// Budget-sized slice height for this width.
|
||
let slice_h_budget = (budget / work_w as f64).floor().clamp(1.0, work_h as f64);
|
||
if (work_h as f64) <= slice_h_budget * p.tall_threshold.max(1.0) {
|
||
return Plan::Single {
|
||
width: work_w,
|
||
height: work_h,
|
||
};
|
||
}
|
||
|
||
// Would the budget-sized slices exceed the cap? If so, grow the slice
|
||
// height so exactly `max_slices` bands cover the page (coarser; the
|
||
// oversized crop is downscaled to budget by `render_slice`).
|
||
let step_for = |sh: f64| (sh * (1.0 - overlap)).max(1.0);
|
||
let bands_for = |sh: f64| (((work_h as f64 - sh) / step_for(sh)).ceil() as i64 + 1).max(1);
|
||
let slice_h = if bands_for(slice_h_budget) > p.max_slices as i64 {
|
||
(work_h as f64 / (1.0 + (p.max_slices as f64 - 1.0) * (1.0 - overlap)))
|
||
.ceil()
|
||
.clamp(1.0, work_h as f64)
|
||
} else {
|
||
slice_h_budget
|
||
} as u32;
|
||
|
||
let step = ((slice_h as f64) * (1.0 - overlap)).round().max(1.0) as u32;
|
||
let mut bands: Vec<(u32, u32)> = Vec::new();
|
||
let mut y0 = 0u32;
|
||
loop {
|
||
let y1 = (y0 + slice_h).min(work_h);
|
||
bands.push((y0, y1));
|
||
if y1 >= work_h || bands.len() >= p.max_slices {
|
||
break;
|
||
}
|
||
y0 += step;
|
||
}
|
||
// Guarantee coverage to the bottom edge.
|
||
if let Some(last) = bands.last_mut() {
|
||
last.1 = work_h;
|
||
}
|
||
Plan::Sliced {
|
||
width: work_w,
|
||
height: work_h,
|
||
bands,
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Image rendering
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Resize `img` down to `max_pixels` (preserving aspect) only if it exceeds
|
||
/// the budget, then JPEG-encode. An image already within budget is encoded
|
||
/// at native resolution.
|
||
fn render_whole(img: &DynamicImage, max_pixels: u32) -> Option<Vec<u8>> {
|
||
let pixels = img.width() as u64 * img.height() as u64;
|
||
let out = if pixels > max_pixels as u64 && pixels > 0 {
|
||
let s = (max_pixels as f64 / pixels as f64).sqrt();
|
||
let nw = ((img.width() as f64 * s).round() as u32).max(1);
|
||
let nh = ((img.height() as f64 * s).round() as u32).max(1);
|
||
img.resize_exact(nw, nh, FilterType::Triangle)
|
||
} else {
|
||
img.clone()
|
||
};
|
||
encode_jpeg(&out)
|
||
}
|
||
|
||
/// Crop a vertical band `[y0, y1)` from the working image and encode it,
|
||
/// downscaling to budget only if the band is oversized (the coarse fallback).
|
||
fn render_slice(work: &DynamicImage, y0: u32, y1: u32, max_pixels: u32) -> Option<Vec<u8>> {
|
||
let h = y1.saturating_sub(y0).max(1).min(work.height().saturating_sub(y0).max(1));
|
||
let crop = work.crop_imm(0, y0, work.width(), h);
|
||
render_whole(&crop, max_pixels)
|
||
}
|
||
|
||
fn encode_jpeg(img: &DynamicImage) -> Option<Vec<u8>> {
|
||
let mut buf = Vec::new();
|
||
img.to_rgb8()
|
||
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Jpeg)
|
||
.ok()?;
|
||
Some(buf)
|
||
}
|
||
|
||
fn b64(bytes: &[u8]) -> String {
|
||
base64::engine::general_purpose::STANDARD.encode(bytes)
|
||
}
|
||
|
||
fn data_url(jpeg: &[u8]) -> String {
|
||
format!("data:image/jpeg;base64,{}", b64(jpeg))
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// OCR merge (seam dedup)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// How many pieces at each seam edge to consider for dedup.
|
||
const SEAM: usize = 8;
|
||
|
||
/// One slice's OCR plus the vertical band `[y0, y1)` it covered in the
|
||
/// working image — so per-piece `y` fractions map to page-global positions.
|
||
pub struct SliceOcr {
|
||
pub y0: f64,
|
||
pub y1: f64,
|
||
pub pieces: Vec<OcrResult>,
|
||
}
|
||
|
||
/// One merged-list entry with the geometry used for dedup decisions.
|
||
struct Entry {
|
||
piece: OcrResult,
|
||
/// Page-global y of the text (working-image space); `None` when the
|
||
/// slice's positions weren't usable.
|
||
page_y: Option<f64>,
|
||
/// Center of the slice this piece came from.
|
||
center: f64,
|
||
}
|
||
|
||
/// Merge per-slice OCR (top-to-bottom) into one ordered list, de-duplicating
|
||
/// text that straddles a slice seam.
|
||
///
|
||
/// Pairing is *seam-scoped* (first `SEAM` pieces of a slice vs. the previous
|
||
/// slice's tail). A pair is a duplicate when the two are **position-close**
|
||
/// (page-global y within a text-line tolerance, same kind) OR text-similar —
|
||
/// so even a mis-OCR'd boundary line ("slightly different text, same spot")
|
||
/// is caught. The kept copy is the one whose text was **more central in its
|
||
/// slice** (farther from the cut = less cropped); when positions are
|
||
/// unavailable it falls back to keeping the longer transcription. Repeats in
|
||
/// non-adjacent slices are preserved.
|
||
fn merge_ocr(slices: Vec<SliceOcr>) -> Vec<OcrResult> {
|
||
let mut out: Vec<Entry> = Vec::new();
|
||
let mut prev_start = 0usize;
|
||
for (si, slice) in slices.into_iter().enumerate() {
|
||
let center = (slice.y0 + slice.y1) / 2.0;
|
||
let slice_h = (slice.y1 - slice.y0).max(1.0);
|
||
let tol = 0.04 * slice_h; // ~one text line
|
||
let cur_start = out.len();
|
||
|
||
// Self-calibrate the model's y values: if they look like pixels
|
||
// (max > 1.5) normalize by the max; require some spread to trust
|
||
// them at all (a degenerate constant column is worthless).
|
||
let ys: Vec<f64> = slice.pieces.iter().filter_map(|p| p.y).collect();
|
||
let (ymin, ymax) = ys.iter().fold((f64::MAX, f64::MIN), |(lo, hi), &v| {
|
||
(lo.min(v), hi.max(v))
|
||
});
|
||
let calib = if ymax > 1.5 { ymax } else { 1.0 };
|
||
let usable = ys.len() >= 2 && (ymax - ymin) / calib > 0.05;
|
||
|
||
for (pi, mut piece) in slice.pieces.into_iter().enumerate() {
|
||
let page_y = if usable {
|
||
piece.y.map(|y| slice.y0 + (y / calib).clamp(0.0, 1.0) * slice_h)
|
||
} else {
|
||
None
|
||
};
|
||
piece.y = None; // internal-only, never persisted
|
||
|
||
if si > 0 && pi < SEAM {
|
||
let win = cur_start.saturating_sub(SEAM).max(prev_start);
|
||
let np = normalize_text(&piece.text);
|
||
let matched = (win..cur_start).find(|&j| {
|
||
let e = &out[j];
|
||
let pos = match (e.page_y, page_y) {
|
||
(Some(a), Some(b)) => {
|
||
(a - b).abs() <= tol && e.piece.kind == piece.kind
|
||
}
|
||
_ => false,
|
||
};
|
||
pos || (!np.is_empty() && similar(&normalize_text(&e.piece.text), &np))
|
||
});
|
||
if let Some(j) = matched {
|
||
let keep_new = match (out[j].page_y, page_y) {
|
||
(Some(ey), Some(ny)) => {
|
||
(ny - center).abs() < (ey - out[j].center).abs()
|
||
}
|
||
_ => piece.text.chars().count() > out[j].piece.text.chars().count(),
|
||
};
|
||
if keep_new {
|
||
out[j] = Entry { piece, page_y, center };
|
||
}
|
||
continue;
|
||
}
|
||
}
|
||
out.push(Entry { piece, page_y, center });
|
||
}
|
||
prev_start = cur_start;
|
||
}
|
||
out.into_iter().map(|e| e.piece).collect()
|
||
}
|
||
|
||
/// Lowercased, alphanumeric-only, whitespace-collapsed form for comparison.
|
||
fn normalize_text(s: &str) -> String {
|
||
let mut t = String::with_capacity(s.len());
|
||
for c in s.chars() {
|
||
if c.is_alphanumeric() {
|
||
t.extend(c.to_lowercase());
|
||
} else {
|
||
t.push(' ');
|
||
}
|
||
}
|
||
t.split_whitespace().collect::<Vec<_>>().join(" ")
|
||
}
|
||
|
||
/// Heuristic similarity over already-normalized strings: equal, containment
|
||
/// (one inside the other, ≥3 chars), or word-set Jaccard ≥ 0.6.
|
||
fn similar(a: &str, b: &str) -> bool {
|
||
if a == b {
|
||
return true;
|
||
}
|
||
if a.is_empty() || b.is_empty() {
|
||
return false;
|
||
}
|
||
let (short, long) = if a.len() <= b.len() { (a, b) } else { (b, a) };
|
||
if short.len() >= 3 && long.contains(short) {
|
||
return true;
|
||
}
|
||
let wa: HashSet<&str> = a.split_whitespace().collect();
|
||
let wb: HashSet<&str> = b.split_whitespace().collect();
|
||
if wa.is_empty() || wb.is_empty() {
|
||
return false;
|
||
}
|
||
let inter = wa.intersection(&wb).count() as f64;
|
||
let union = wa.union(&wb).count() as f64;
|
||
inter / union >= 0.6
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Request builders
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Combined single-call body (OCR + tags + scene + safety). Public + this
|
||
/// exact signature so the existing tests keep working.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn build_request_body(
|
||
model: &str,
|
||
max_tokens: u32,
|
||
data_url: &str,
|
||
response_format: ResponseFormat,
|
||
frequency_penalty: f64,
|
||
temperature: f64,
|
||
system_prompt: &str,
|
||
) -> serde_json::Value {
|
||
build_body(
|
||
model,
|
||
max_tokens,
|
||
response_format,
|
||
frequency_penalty,
|
||
temperature,
|
||
system_prompt,
|
||
"page_analysis",
|
||
prompt::output_json_schema(),
|
||
data_url,
|
||
None,
|
||
)
|
||
}
|
||
|
||
/// Pass-A OCR-only body.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn build_ocr_body(
|
||
model: &str,
|
||
max_tokens: u32,
|
||
response_format: ResponseFormat,
|
||
frequency_penalty: f64,
|
||
temperature: f64,
|
||
ocr_prompt: &str,
|
||
data_url: &str,
|
||
) -> serde_json::Value {
|
||
build_body(
|
||
model,
|
||
max_tokens,
|
||
response_format,
|
||
frequency_penalty,
|
||
temperature,
|
||
ocr_prompt,
|
||
"page_ocr",
|
||
prompt::ocr_json_schema(),
|
||
data_url,
|
||
None,
|
||
)
|
||
}
|
||
|
||
/// Pass-B grounding body — whole image + the merged OCR text as context.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn build_grounding_body(
|
||
model: &str,
|
||
max_tokens: u32,
|
||
response_format: ResponseFormat,
|
||
frequency_penalty: f64,
|
||
temperature: f64,
|
||
grounding_prompt: &str,
|
||
data_url: &str,
|
||
ocr_text: &str,
|
||
) -> serde_json::Value {
|
||
let context = format!("OCR text extracted from this page:\n{ocr_text}");
|
||
build_body(
|
||
model,
|
||
max_tokens,
|
||
response_format,
|
||
frequency_penalty,
|
||
temperature,
|
||
grounding_prompt,
|
||
"page_grounding",
|
||
prompt::grounding_json_schema(),
|
||
data_url,
|
||
Some(&context),
|
||
)
|
||
}
|
||
|
||
/// Build a chat/completions body with a system prompt, an image part, an
|
||
/// optional trailing text part, and the requested `response_format`. A
|
||
/// non-zero `frequency_penalty` is included to suppress repetition loops.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn build_body(
|
||
model: &str,
|
||
max_tokens: u32,
|
||
response_format: ResponseFormat,
|
||
frequency_penalty: f64,
|
||
temperature: f64,
|
||
system_prompt: &str,
|
||
schema_name: &str,
|
||
schema: serde_json::Value,
|
||
data_url: &str,
|
||
user_text: Option<&str>,
|
||
) -> serde_json::Value {
|
||
let mut user_content =
|
||
vec![json!({ "type": "image_url", "image_url": { "url": data_url } })];
|
||
if let Some(text) = user_text {
|
||
user_content.push(json!({ "type": "text", "text": text }));
|
||
}
|
||
let mut body = json!({
|
||
"model": model,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
"messages": [
|
||
{ "role": "system", "content": system_prompt },
|
||
{ "role": "user", "content": user_content }
|
||
]
|
||
});
|
||
if frequency_penalty.abs() > f64::EPSILON {
|
||
body["frequency_penalty"] = json!(frequency_penalty);
|
||
}
|
||
let rf = match response_format {
|
||
ResponseFormat::None => None,
|
||
ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })),
|
||
ResponseFormat::JsonSchema => Some(json!({
|
||
"type": "json_schema",
|
||
"json_schema": { "name": schema_name, "strict": true, "schema": schema }
|
||
})),
|
||
};
|
||
if let Some(rf) = rf {
|
||
body["response_format"] = rf;
|
||
}
|
||
body
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Response parsing / sanitizing
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Extract and parse the model's JSON object from a chat/completions
|
||
/// response into a (sanitized) [`VisionAnalysis`]. Lenient field defaults
|
||
/// let the same parser serve the combined, OCR-only and grounding shapes.
|
||
pub fn parse_chat_completion(value: &serde_json::Value) -> anyhow::Result<VisionAnalysis> {
|
||
let content = value
|
||
.get("choices")
|
||
.and_then(|c| c.get(0))
|
||
.and_then(|c| c.get("message"))
|
||
.and_then(|m| m.get("content"))
|
||
.and_then(|c| c.as_str())
|
||
.ok_or_else(|| anyhow!("vision response missing choices[0].message.content"))?;
|
||
let json_slice = extract_json_object(content);
|
||
let parsed: VisionAnalysis = serde_json::from_str(json_slice)
|
||
.with_context(|| format!("vision content was not valid analysis JSON: {content:?}"))?;
|
||
Ok(sanitize(parsed))
|
||
}
|
||
|
||
/// Return the substring from the first `{` to the last `}` (inclusive),
|
||
/// tolerating markdown fences or stray prose around the JSON object.
|
||
fn extract_json_object(s: &str) -> &str {
|
||
match (s.find('{'), s.rfind('}')) {
|
||
(Some(start), Some(end)) if end >= start => &s[start..=end],
|
||
_ => s.trim(),
|
||
}
|
||
}
|
||
|
||
/// Bound the model output: drop empty OCR text, truncate over-long fields,
|
||
/// clamp tags via the page-tag normalizer, and keep only recognized
|
||
/// content-warning values.
|
||
pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
||
a.ocr_results.retain(|r| !r.text.trim().is_empty());
|
||
|
||
// Collapse runaway repetition: a degenerating model emits the same line
|
||
// over and over. Drop any piece identical (normalized text + kind) to the
|
||
// one just kept, so a loop that slips past the schema caps still can't
|
||
// flood the row.
|
||
let mut prev: Option<String> = None;
|
||
a.ocr_results.retain(|r| {
|
||
let key = format!("{}\u{0}{}", r.kind, normalize_text(&r.text));
|
||
if prev.as_deref() == Some(key.as_str()) {
|
||
false
|
||
} else {
|
||
prev = Some(key);
|
||
true
|
||
}
|
||
});
|
||
|
||
a.ocr_results.truncate(MAX_OCR_PIECES);
|
||
for r in &mut a.ocr_results {
|
||
r.text = truncate_chars(r.text.trim(), MAX_OCR_TEXT_CHARS);
|
||
}
|
||
|
||
let mut seen = std::collections::HashSet::new();
|
||
let mut tags = Vec::new();
|
||
for raw in std::mem::take(&mut a.tagging_results) {
|
||
if let Ok(norm) = crate::api::page_tags::normalize_tag(&raw) {
|
||
if seen.insert(norm.clone()) {
|
||
tags.push(norm);
|
||
if tags.len() >= MAX_TAGS {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
a.tagging_results = tags;
|
||
|
||
a.scene_description = truncate_chars(a.scene_description.trim(), MAX_SCENE_CHARS);
|
||
|
||
let mut seen_w = std::collections::HashSet::new();
|
||
a.safety_flag.content_type = std::mem::take(&mut a.safety_flag.content_type)
|
||
.into_iter()
|
||
.filter_map(|w| ContentWarning::from_model_str(&w).map(|_| w.trim().to_lowercase()))
|
||
.filter(|w| seen_w.insert(w.clone()))
|
||
.collect();
|
||
|
||
a
|
||
}
|
||
|
||
fn truncate_chars(s: &str, n: usize) -> String {
|
||
s.chars().take(n).collect()
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn chat_body(content: &str) -> serde_json::Value {
|
||
json!({ "choices": [ { "message": { "content": content } } ] })
|
||
}
|
||
|
||
fn params() -> SliceParams {
|
||
SliceParams {
|
||
max_pixels: 1_000_000,
|
||
min_slice_height: 640,
|
||
overlap: 0.12,
|
||
tall_threshold: 1.6,
|
||
max_slices: 16,
|
||
}
|
||
}
|
||
|
||
fn ocr(text: &str, kind: &str) -> OcrResult {
|
||
OcrResult {
|
||
text: text.into(),
|
||
kind: kind.into(),
|
||
y: None,
|
||
}
|
||
}
|
||
|
||
fn ocr_y(text: &str, kind: &str, y: f64) -> OcrResult {
|
||
OcrResult {
|
||
text: text.into(),
|
||
kind: kind.into(),
|
||
y: Some(y),
|
||
}
|
||
}
|
||
|
||
/// A text-only slice (no positions → fuzzy fallback path).
|
||
fn tslice(pieces: Vec<OcrResult>) -> SliceOcr {
|
||
SliceOcr { y0: 0.0, y1: 1000.0, pieces }
|
||
}
|
||
|
||
fn jpeg_of(w: u32, h: u32) -> Vec<u8> {
|
||
let img = DynamicImage::ImageRgb8(image::RgbImage::from_pixel(w, h, image::Rgb([9, 9, 9])));
|
||
encode_jpeg(&img).unwrap()
|
||
}
|
||
|
||
// --- parse / sanitize (unchanged behavior) ---
|
||
|
||
#[test]
|
||
fn parses_a_clean_json_response() {
|
||
let body = chat_body(
|
||
r#"{"ocr_results":[{"text":"Hi","kind":"speech"}],
|
||
"tagging_results":["action","city"],
|
||
"scene_description":"A street.",
|
||
"safety_flag":{"is_nsfw":false,"content_type":[]}}"#,
|
||
);
|
||
let v = parse_chat_completion(&body).unwrap();
|
||
assert_eq!(v.ocr_results.len(), 1);
|
||
assert_eq!(v.tagging_results, vec!["action", "city"]);
|
||
assert!(!v.safety_flag.is_nsfw);
|
||
}
|
||
|
||
#[test]
|
||
fn ocr_only_response_parses_with_defaults() {
|
||
let body = chat_body(r#"{"ocr_results":[{"text":"Hi","kind":"speech"}]}"#);
|
||
let v = parse_chat_completion(&body).unwrap();
|
||
assert_eq!(v.ocr_results.len(), 1);
|
||
assert!(v.tagging_results.is_empty());
|
||
assert_eq!(v.scene_description, "");
|
||
}
|
||
|
||
#[test]
|
||
fn strips_markdown_fences_and_prose() {
|
||
let body = chat_body("Here:\n```json\n{\"scene_description\":\"x\"}\n```\nthanks");
|
||
assert_eq!(parse_chat_completion(&body).unwrap().scene_description, "x");
|
||
}
|
||
|
||
#[test]
|
||
fn rejects_non_json_content() {
|
||
assert!(parse_chat_completion(&chat_body("nope")).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn sanitize_filters_unknown_warnings() {
|
||
let body = chat_body(
|
||
r#"{"safety_flag":{"is_nsfw":true,"content_type":["Sexual","spicy","gore","gore"]}}"#,
|
||
);
|
||
let v = parse_chat_completion(&body).unwrap();
|
||
assert_eq!(v.safety_flag.content_type, vec!["sexual", "gore"]);
|
||
}
|
||
|
||
// --- request builders ---
|
||
|
||
#[test]
|
||
fn combined_body_uses_full_schema_and_image() {
|
||
let body = build_request_body(
|
||
"m",
|
||
100,
|
||
"data:image/png;base64,AA",
|
||
ResponseFormat::JsonSchema,
|
||
0.3,
|
||
0.0,
|
||
prompt::SYSTEM_PROMPT_DEFAULT,
|
||
);
|
||
assert_eq!(body["messages"][0]["role"], "system");
|
||
assert_eq!(
|
||
body["messages"][1]["content"][0]["image_url"]["url"],
|
||
"data:image/png;base64,AA"
|
||
);
|
||
assert_eq!(body["response_format"]["json_schema"]["name"], "page_analysis");
|
||
let req = &body["response_format"]["json_schema"]["schema"]["required"];
|
||
assert!(req.as_array().unwrap().iter().any(|v| v == "safety_flag"));
|
||
}
|
||
|
||
#[test]
|
||
fn body_carries_the_configured_prompt_and_temperature() {
|
||
let body = build_request_body(
|
||
"m",
|
||
100,
|
||
"d",
|
||
ResponseFormat::None,
|
||
0.0,
|
||
0.7,
|
||
"CUSTOM PROMPT",
|
||
);
|
||
assert_eq!(body["messages"][0]["content"], "CUSTOM PROMPT");
|
||
assert_eq!(body["temperature"], 0.7);
|
||
}
|
||
|
||
#[test]
|
||
fn ocr_body_carries_only_the_ocr_schema_and_caps() {
|
||
let body =
|
||
build_ocr_body("m", 100, ResponseFormat::JsonSchema, 0.3, 0.0, prompt::OCR_PROMPT_DEFAULT, "d");
|
||
let schema = &body["response_format"]["json_schema"]["schema"];
|
||
assert_eq!(body["response_format"]["json_schema"]["name"], "page_ocr");
|
||
assert_eq!(schema["required"], json!(["ocr_results"]));
|
||
// Hard repetition guards baked into the schema.
|
||
assert!(schema["properties"]["ocr_results"]["maxItems"].is_number());
|
||
assert!(
|
||
schema["properties"]["ocr_results"]["items"]["properties"]["text"]["maxLength"]
|
||
.is_number()
|
||
);
|
||
assert_eq!(body["messages"][1]["content"].as_array().unwrap().len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn grounding_body_includes_ocr_context_and_grounding_schema() {
|
||
let body = build_grounding_body(
|
||
"m",
|
||
100,
|
||
ResponseFormat::JsonSchema,
|
||
0.3,
|
||
0.0,
|
||
prompt::GROUNDING_PROMPT_DEFAULT,
|
||
"d",
|
||
"[speech] Hi",
|
||
);
|
||
assert_eq!(body["response_format"]["json_schema"]["name"], "page_grounding");
|
||
let parts = body["messages"][1]["content"].as_array().unwrap();
|
||
assert_eq!(parts.len(), 2, "image + text parts");
|
||
assert_eq!(parts[0]["type"], "image_url");
|
||
assert_eq!(parts[1]["type"], "text");
|
||
assert!(parts[1]["text"].as_str().unwrap().contains("[speech] Hi"));
|
||
}
|
||
|
||
#[test]
|
||
fn response_format_none_omits_the_field() {
|
||
let body = build_request_body(
|
||
"m",
|
||
100,
|
||
"d",
|
||
ResponseFormat::None,
|
||
0.3,
|
||
0.0,
|
||
prompt::SYSTEM_PROMPT_DEFAULT,
|
||
);
|
||
assert!(body.get("response_format").is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn frequency_penalty_included_only_when_nonzero() {
|
||
let on =
|
||
build_ocr_body("m", 100, ResponseFormat::None, 0.4, 0.0, prompt::OCR_PROMPT_DEFAULT, "d");
|
||
assert_eq!(on["frequency_penalty"], 0.4);
|
||
let off =
|
||
build_ocr_body("m", 100, ResponseFormat::None, 0.0, 0.0, prompt::OCR_PROMPT_DEFAULT, "d");
|
||
assert!(off.get("frequency_penalty").is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn sanitize_collapses_runaway_repetition() {
|
||
let body = chat_body(
|
||
r#"{"ocr_results":[
|
||
{"text":"I don't know","kind":"speech"},
|
||
{"text":"I don't know","kind":"speech"},
|
||
{"text":"I don't know","kind":"speech"},
|
||
{"text":"ok","kind":"speech"}
|
||
]}"#,
|
||
);
|
||
let v = parse_chat_completion(&body).unwrap();
|
||
let texts: Vec<&str> = v.ocr_results.iter().map(|o| o.text.as_str()).collect();
|
||
assert_eq!(texts, vec!["I don't know", "ok"], "consecutive repeats collapsed");
|
||
}
|
||
|
||
// --- geometry: plan_slices ---
|
||
|
||
#[test]
|
||
fn normal_aspect_page_is_a_single_call() {
|
||
assert!(matches!(
|
||
plan_slices(800, 1000, ¶ms()),
|
||
Plan::Single { .. }
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn tall_narrow_page_slices_into_portrait_bands() {
|
||
let p = params();
|
||
let Plan::Sliced { width, height, bands } = plan_slices(800, 8000, &p) else {
|
||
panic!("expected Sliced");
|
||
};
|
||
assert_eq!((width, height), (800, 8000));
|
||
assert!(bands.len() >= 2 && bands.len() <= p.max_slices);
|
||
assert_eq!(bands.first().unwrap().0, 0, "starts at top");
|
||
assert_eq!(bands.last().unwrap().1, 8000, "covers the bottom");
|
||
// Overlapping + advancing bands; each ≈ the pixel budget.
|
||
for win in bands.windows(2) {
|
||
assert!(win[1].0 < win[0].1, "adjacent bands overlap");
|
||
assert!(win[1].0 > win[0].0, "bands advance");
|
||
}
|
||
for (y0, y1) in &bands {
|
||
let px = (y1 - y0) as u64 * width as u64;
|
||
assert!(px <= 1_000_000 + 1, "band within budget, got {px}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn wide_page_yields_landscape_bands() {
|
||
// 1400 < maxSliceWidth(1562) so width is preserved; budget/1400≈714
|
||
// → bands are wider than tall.
|
||
let Plan::Sliced { width, bands, .. } = plan_slices(1400, 4000, ¶ms()) else {
|
||
panic!("expected Sliced");
|
||
};
|
||
assert_eq!(width, 1400);
|
||
for (y0, y1) in &bands {
|
||
assert!(y1 - y0 < width, "landscape band (height < width)");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn pathologically_wide_page_reduces_width() {
|
||
let p = params();
|
||
let max_slice_width = (p.max_pixels as f64 / p.min_slice_height as f64).floor() as u32;
|
||
let Plan::Sliced { width, .. } = plan_slices(3000, 9000, &p) else {
|
||
panic!("expected Sliced");
|
||
};
|
||
assert!(width <= max_slice_width && width < 3000, "width reduced to the guard");
|
||
}
|
||
|
||
#[test]
|
||
fn very_long_page_is_capped_at_max_slices() {
|
||
let p = params();
|
||
let Plan::Sliced { bands, height, .. } = plan_slices(800, 200_000, &p) else {
|
||
panic!("expected Sliced");
|
||
};
|
||
assert!(bands.len() <= p.max_slices);
|
||
assert_eq!(bands.last().unwrap().1, height, "still covers the bottom");
|
||
}
|
||
|
||
// --- merge_ocr ---
|
||
|
||
#[test]
|
||
fn merge_dedups_a_line_repeated_across_the_seam() {
|
||
let merged = merge_ocr(vec![
|
||
tslice(vec![ocr("Hello", "speech"), ocr("world", "speech")]),
|
||
tslice(vec![ocr("world", "speech"), ocr("bye", "speech")]),
|
||
]);
|
||
let texts: Vec<&str> = merged.iter().map(|o| o.text.as_str()).collect();
|
||
assert_eq!(texts, vec!["Hello", "world", "bye"]);
|
||
}
|
||
|
||
#[test]
|
||
fn merge_keeps_the_longer_transcription_without_positions() {
|
||
let merged = merge_ocr(vec![
|
||
tslice(vec![ocr("I don", "speech")]),
|
||
tslice(vec![ocr("I don't think so", "speech")]),
|
||
]);
|
||
assert_eq!(merged.len(), 1);
|
||
assert_eq!(merged[0].text, "I don't think so");
|
||
}
|
||
|
||
#[test]
|
||
fn merge_keeps_non_adjacent_repeats() {
|
||
// "boom" in slice 0 and slice 2 must both survive (not a seam dup).
|
||
let merged = merge_ocr(vec![
|
||
tslice(vec![ocr("boom", "sfx")]),
|
||
tslice(vec![ocr("crash", "sfx")]),
|
||
tslice(vec![ocr("boom", "sfx")]),
|
||
]);
|
||
let booms = merged.iter().filter(|o| o.text == "boom").count();
|
||
assert_eq!(booms, 2);
|
||
}
|
||
|
||
#[test]
|
||
fn merge_preserves_distinct_lines_and_order() {
|
||
let merged = merge_ocr(vec![
|
||
tslice(vec![ocr("one", "narration")]),
|
||
tslice(vec![ocr("two", "speech")]),
|
||
]);
|
||
assert_eq!(merged.len(), 2);
|
||
assert_eq!(merged[0].kind, "narration");
|
||
assert_eq!(merged[1].text, "two");
|
||
}
|
||
|
||
#[test]
|
||
fn merge_pairs_by_position_when_text_differs_and_keeps_the_less_cropped() {
|
||
// The same physical line sits at page-y ~920 (in the [900,1000]
|
||
// overlap). The upper slice saw it well inside; the lower slice
|
||
// saw it garbled at its very top. The two transcriptions are NOT
|
||
// text-similar — only position pairs them, and centrality keeps
|
||
// the upper (less-cropped) one.
|
||
let upper = SliceOcr {
|
||
y0: 0.0,
|
||
y1: 1000.0,
|
||
pieces: vec![ocr_y("hi", "speech", 0.2), ocr_y("Hello", "speech", 0.92)],
|
||
};
|
||
let lower = SliceOcr {
|
||
y0: 900.0,
|
||
y1: 1900.0,
|
||
pieces: vec![ocr_y("He110", "speech", 0.02), ocr_y("bye", "speech", 0.6)],
|
||
};
|
||
let merged = merge_ocr(vec![upper, lower]);
|
||
let texts: Vec<&str> = merged.iter().map(|o| o.text.as_str()).collect();
|
||
assert!(texts.contains(&"Hello"), "kept the less-cropped copy: {texts:?}");
|
||
assert!(!texts.contains(&"He110"), "dropped the garbled copy: {texts:?}");
|
||
assert_eq!(merged.len(), 3, "hi + Hello + bye");
|
||
}
|
||
|
||
// --- rendering ---
|
||
|
||
#[test]
|
||
fn render_whole_passes_small_images_through() {
|
||
let jpeg = jpeg_of(100, 80);
|
||
let img = image::load_from_memory(&jpeg).unwrap();
|
||
let out = render_whole(&img, 1_000_000).unwrap();
|
||
let dec = image::load_from_memory(&out).unwrap();
|
||
assert_eq!((dec.width(), dec.height()), (100, 80));
|
||
}
|
||
|
||
#[test]
|
||
fn render_whole_fits_oversized_to_budget() {
|
||
let jpeg = jpeg_of(2000, 1000); // 2.0 MP
|
||
let img = image::load_from_memory(&jpeg).unwrap();
|
||
let out = render_whole(&img, 500_000).unwrap();
|
||
let dec = image::load_from_memory(&out).unwrap();
|
||
let px = dec.width() as u64 * dec.height() as u64;
|
||
assert!(px <= 520_000, "fit to ~budget, got {px}");
|
||
}
|
||
|
||
#[test]
|
||
fn render_slice_crops_the_band() {
|
||
let jpeg = jpeg_of(100, 400);
|
||
let img = image::load_from_memory(&jpeg).unwrap();
|
||
let out = render_slice(&img, 50, 150, 1_000_000).unwrap();
|
||
let dec = image::load_from_memory(&out).unwrap();
|
||
assert_eq!((dec.width(), dec.height()), (100, 100));
|
||
}
|
||
|
||
fn small_params() -> SliceParams {
|
||
SliceParams {
|
||
max_pixels: 1_000_000,
|
||
min_slice_height: 100,
|
||
overlap: 0.05,
|
||
tall_threshold: 1.8,
|
||
max_slices: 6,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn prepare_analysis_single_pages_emit_one_jpeg() {
|
||
let jpeg = jpeg_of(200, 200);
|
||
match prepare_analysis(&jpeg, small_params()) {
|
||
PreparedAnalysis::Single { jpeg } => assert!(!jpeg.is_empty()),
|
||
other => panic!("expected Single, got {:?}", std::mem::discriminant(&other)),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn prepare_analysis_tall_pages_emit_slice_and_whole_jpegs() {
|
||
// Long enough that height > slice_h_budget * tall_threshold —
|
||
// with the test params (max_pixels=1M, width=200, threshold=1.8)
|
||
// the threshold is 200×5000×1.8 = 9000 px. A 200×10_000 page
|
||
// pushes us into Sliced.
|
||
let jpeg = jpeg_of(200, 10_000);
|
||
match prepare_analysis(&jpeg, small_params()) {
|
||
PreparedAnalysis::Sliced { slices, whole } => {
|
||
assert!(slices.len() >= 2, "expected at least 2 bands, got {}", slices.len());
|
||
assert!(slices.iter().all(|(_, _, j)| !j.is_empty()));
|
||
assert!(!whole.is_empty());
|
||
}
|
||
other => panic!("expected Sliced, got {:?}", std::mem::discriminant(&other)),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn prepare_analysis_garbage_bytes_yield_undecodable() {
|
||
match prepare_analysis(&[0u8, 1, 2, 3], small_params()) {
|
||
PreparedAnalysis::Undecodable => {}
|
||
other => panic!("expected Undecodable, got {:?}", std::mem::discriminant(&other)),
|
||
}
|
||
}
|
||
|
||
/// Pin the contract that the 0.87.5 fix is about: the heavy image
|
||
/// work in `analyze()` MUST go through `spawn_blocking` so the tokio
|
||
/// runtime stays responsive. If a refactor accidentally inlined
|
||
/// `prepare_analysis` back into the async function, this test would
|
||
/// catch it.
|
||
///
|
||
/// How we detect it: run `analyze()` on a `current_thread` runtime
|
||
/// (one worker thread). Schedule a concurrent counter task that
|
||
/// ticks every 5ms with `MissedTickBehavior::Skip`. The image is
|
||
/// large enough that `prepare_analysis` takes meaningfully longer
|
||
/// than the tick interval. If prep is on `spawn_blocking`, the
|
||
/// worker thread stays free to drive the counter and we observe
|
||
/// several ticks. If prep is inlined onto the runtime, the worker
|
||
/// is starved and the counter cannot advance during the prep
|
||
/// window.
|
||
///
|
||
/// **Why `MissedTickBehavior::Skip`** (rereview fix): `Burst` (the
|
||
/// default) replays every missed tick as soon as the runtime yields
|
||
/// — even a brief reqwest connect-refused yield after inlined prep
|
||
/// would burst ~20 ticks at once, faking the `>= 2` threshold and
|
||
/// hiding the regression. Skip drops the backlog; the only way to
|
||
/// observe `>= 2` ticks is for the runtime to actually have been
|
||
/// making progress during the 5ms intervals.
|
||
///
|
||
/// We also snapshot the tick count at three checkpoints (pre-prep,
|
||
/// during, post) so the assertion measures progress DURING the
|
||
/// prep window, not just the analyse-call total. This survives
|
||
/// the post-prep HTTP fail-fast that the prior version of this
|
||
/// test conflated with prep ticks.
|
||
#[tokio::test(flavor = "current_thread", start_paused = false)]
|
||
async fn analyze_dispatches_image_prep_off_runtime() {
|
||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||
use std::sync::Arc;
|
||
use std::time::Duration;
|
||
use tokio::time::MissedTickBehavior;
|
||
|
||
// VisionClient pointed at a closed local port. Reqwest fails the
|
||
// POST immediately (connection refused), but the prep work runs
|
||
// before the POST is even built.
|
||
let mut cfg = AnalysisConfig::default();
|
||
cfg.endpoint = "http://127.0.0.1:1/v1/chat/completions".into();
|
||
cfg.model = "stub".into();
|
||
cfg.request_timeout = Duration::from_secs(1);
|
||
// Tune slice geometry to match the test image's shape.
|
||
cfg.max_pixels = 1_000_000;
|
||
cfg.min_slice_height = 100;
|
||
cfg.tall_aspect_threshold = 1.8;
|
||
let http = reqwest::Client::builder()
|
||
.timeout(cfg.request_timeout)
|
||
.no_proxy()
|
||
.build()
|
||
.unwrap();
|
||
let client = VisionClient::new(http, &cfg);
|
||
|
||
// ~6 MP image so decode + JPEG encode comfortably exceeds the
|
||
// 5 ms tick interval on every realistic CI runner.
|
||
let big = jpeg_of(2000, 3000);
|
||
|
||
// Counter that ticks at 5ms with Skip semantics so a single
|
||
// post-prep yield can't burst-replay enough ticks to fake the
|
||
// assertion. If the runtime is starved during prep, this task
|
||
// makes no progress in that window.
|
||
let ticks = Arc::new(AtomicUsize::new(0));
|
||
let ticks_c = Arc::clone(&ticks);
|
||
let ticker = tokio::spawn(async move {
|
||
let mut interval = tokio::time::interval(Duration::from_millis(5));
|
||
interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
|
||
interval.tick().await; // discard immediate first tick
|
||
loop {
|
||
interval.tick().await;
|
||
ticks_c.fetch_add(1, Ordering::Relaxed);
|
||
}
|
||
});
|
||
|
||
// Snapshot tick count at three points: pre-call, mid-call
|
||
// (during prep — if `spawn_blocking` is used, this fires within
|
||
// the prep window), and post-call. The mid-call snapshot is
|
||
// what distinguishes "runtime was free during prep" from
|
||
// "runtime caught up after analyze() returned". On
|
||
// current_thread with prep on the runtime, mid-call would tick
|
||
// 0 times no matter how long the call takes.
|
||
let ticks_pre = ticks.load(Ordering::Relaxed);
|
||
let mid_check = {
|
||
let ticks_c = Arc::clone(&ticks);
|
||
tokio::spawn(async move {
|
||
// Sleep just long enough that prep is provably mid-call
|
||
// (decode + resize alone exceeds 10 ms on every real
|
||
// runner). Snapshot during, return.
|
||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||
ticks_c.load(Ordering::Relaxed)
|
||
})
|
||
};
|
||
let _ = client.analyze(&big, "image/jpeg").await;
|
||
let ticks_during = mid_check.await.unwrap();
|
||
let ticks_post = ticks.load(Ordering::Relaxed);
|
||
ticker.abort();
|
||
|
||
let during_window = ticks_during.saturating_sub(ticks_pre);
|
||
let total = ticks_post.saturating_sub(ticks_pre);
|
||
|
||
// Discrimination: the `during_window` snapshot at +30ms is
|
||
// dominated by mid_check's own scheduling latency on
|
||
// current_thread (the snapshot can't run until the analyze
|
||
// call yields), so under either strategy it sits at a similar
|
||
// small number. The `total` (post − pre) is the unambiguous
|
||
// signal — measured over the WHOLE analyze duration plus the
|
||
// post-call mid_check wait:
|
||
// * spawn_blocking → runtime stays free; counter ticks
|
||
// continuously at 5 ms cadence with Skip semantics;
|
||
// observed ~700+ on a typical runner.
|
||
// * inlined prep → runtime worker is blocked during the
|
||
// ~100 ms prep CPU work; Skip drops the entire backlog so
|
||
// only the post-prep yield ticks count; observed ~5.
|
||
// Threshold of 50 sits well above the inlined-prep ceiling and
|
||
// far below the spawn_blocking floor — robust to CPU jitter
|
||
// without sacrificing the regression signal.
|
||
let _ = ticks_during;
|
||
let _ = during_window;
|
||
assert!(
|
||
total >= 50,
|
||
"runtime ticked only {total} times over the analyze() call window \
|
||
(mid-call snapshot = {during_window}) — image prep is likely \
|
||
blocking the runtime instead of using spawn_blocking"
|
||
);
|
||
}
|
||
}
|