Files
Mangalord/backend/src/analysis/vision.rs
MechaCat02 80f4819fad perf(analysis): run image decode/resize/encode on the blocking pool (0.87.5)
`VisionClient::analyze` ran `image::load_from_memory`, `DynamicImage::resize_exact(Triangle)`,
and per-band JPEG encodes directly on the tokio task. Each call is
hundreds of ms of pure CPU per page. With multiple workers, the runtime
threads got starved — axum handlers, SSE streams, the crawler daemon's
async timers, and other tasks sharing the runtime all stalled while
analysis was active.

Extract a `prepare_analysis` helper that does ALL the CPU work
(decode + plan + width-reduce + slice + JPEG encode) and returns a
`PreparedAnalysis { Single | Sliced | Undecodable }` of already-encoded
byte payloads. `analyze` calls it inside `tokio::task::spawn_blocking`,
then the async HTTP loop only iterates over the prepared bytes — no
image-crate operations happen on the runtime any more.

Single page → one spawn_blocking → one HTTP call.
Tall page → one spawn_blocking → N OCR calls + 1 grounding call.

Memory peak rises slightly for the Sliced path (all bands held
encoded at the same time instead of one-at-a-time) — for a typical
manga page split into 4 slices at ~200 KB each, that's ~600 KB of
extra peak. Negligible vs. the runtime-starvation cost.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-06-22 21:23:34 +02:00

1144 lines
41 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 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.
fn prepare_analysis(image: &[u8], params: SliceParams) -> PreparedAnalysis {
let Some(img) = image::load_from_memory(image).ok() else {
return PreparedAnalysis::Undecodable;
};
match plan_slices(img.width(), img.height(), &params) {
Plan::Single { .. } => match render_whole(&img, params.max_pixels) {
Some(jpeg) => PreparedAnalysis::Single { jpeg },
None => 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.
return PreparedAnalysis::Undecodable;
};
slices.push((*y0, *y1, jpeg));
}
let Some(whole) = render_whole(&img, params.max_pixels) else {
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, &params()),
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, &params()) 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)),
}
}
}