feat(analysis): vision client + system prompt + AnalysisConfig

The OpenAI-compatible vision client and the bounded prompt/output handling
for the analysis worker (no DB, fully unit-tested):

- analysis::prompt: terse system prompt carrying the JSON schema, the OCR
  kind / content-warning vocabularies, and the output-size caps.
- analysis::vision: VisionClient (downscale via the new `image` dep →
  base64 data-URL → chat/completions with an image_url part), plus pure
  parse_chat_completion (fence/prose-tolerant) and sanitize (drop empty
  OCR, truncate, clamp tags to 10 via the shared page-tag normalizer,
  filter unknown warnings).
- config::AnalysisConfig extended with endpoint/model/api_key/timeouts/
  max_tokens/max_image_dim/max_image_bytes + from_env.
- Deps: add `image` (jpeg/png/webp), reqwest `json` feature. Expose
  api::page_tags::normalize_tag as pub(crate) for reuse.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-06-13 18:45:57 +02:00
parent 2bbf1595ff
commit dbde6e02c4
9 changed files with 587 additions and 11 deletions

View File

@@ -0,0 +1,11 @@
//! AI content-analysis worker: calls a local OpenAI-compatible vision
//! model on each page image and turns the result into OCR text, global
//! auto-tags, a scene description, and NSFW content warnings.
//!
//! * [`prompt`] — the system prompt + the bounded output vocabulary and
//! sanitization helpers (pure, unit-tested).
//! * [`vision`] — the HTTP client: downscale → request → parse → sanitize.
//! * [`daemon`] — the job-leasing worker loop (added in the worker phase).
pub mod prompt;
pub mod vision;

View File

@@ -0,0 +1,90 @@
//! The vision model's system prompt and the bounds applied to its output.
//!
//! The system prompt is deliberately terse: with a context budget as low
//! as ~8192 tokens, the page image dominates, so the instructions + schema
//! must stay small and the output is capped via `max_tokens` plus the
//! length bounds enforced in [`crate::analysis::vision::sanitize`].
/// OCR text kinds the model may emit. Kept in sync with the
/// `page_ocr_text.kind` CHECK and [`crate::domain::page_analysis::OcrKind`].
pub const OCR_KINDS: [&str; 6] =
["speech", "thought", "narration", "sfx", "title", "caption"];
/// Closed content-warning vocabulary. Kept in sync with the
/// `page_content_warnings.warning` CHECK and
/// [`crate::domain::page_analysis::ContentWarning`].
pub const CONTENT_WARNINGS: [&str; 5] =
["sexual", "nudity", "gore", "violence", "disturbing"];
/// Target tag count we ask the model for (a hint in the prompt; not
/// hard-enforced beyond the [`MAX_TAGS`] cap).
pub const MIN_TAGS: usize = 5;
pub const MAX_TAGS: usize = 10;
/// Output bounds enforced at sanitize time so one verbose response can't
/// bloat the row or the search document.
pub const MAX_OCR_PIECES: usize = 60;
pub const MAX_OCR_TEXT_CHARS: usize = 500;
pub const MAX_SCENE_CHARS: usize = 1000;
/// The system prompt. Self-contained: it carries the exact JSON schema so
/// the same string drives both the model and (implicitly) the
/// [`crate::domain::page_analysis::VisionAnalysis`] parser.
pub const SYSTEM_PROMPT: &str = concat!(
"You are a manga/comic page analyzer. Look at the single page image and ",
"return ONLY one minified JSON object — no prose, no markdown fences.\n",
"Schema:\n",
"{\"ocr_results\":[{\"text\":string,\"kind\":",
"\"speech|thought|narration|sfx|title|caption\"}],",
"\"tagging_results\":[string],",
"\"scene_description\":string,",
"\"safety_flag\":{\"is_nsfw\":boolean,\"content_type\":[",
"\"sexual|nudity|gore|violence|disturbing\"]}}\n",
"Rules: transcribe every visible text element verbatim into ocr_results ",
"with its kind; if there is no text use []. tagging_results: 5-10 short, ",
"lowercase content tags (characters, actions, setting, mood, genre); ",
"include explicit/sexual tags when present. scene_description: one or two ",
"sentences describing setting, characters, and action. safety_flag.",
"content_type: only values from the listed set, [] if none; set is_nsfw ",
"true if the page contains sexual, nudity, gore, violence, or disturbing ",
"content. Output must be valid minified JSON and nothing else."
);
#[cfg(test)]
mod tests {
use super::*;
use crate::domain::page_analysis::{ContentWarning, OcrKind};
#[test]
fn vocabularies_match_the_domain_enums() {
// Every prompt OCR kind must parse back to a real OcrKind (no
// fallback to narration via a typo).
for k in OCR_KINDS {
assert_eq!(
OcrKind::from_model_str(k),
match k {
"speech" => OcrKind::Speech,
"thought" => OcrKind::Thought,
"narration" => OcrKind::Narration,
"sfx" => OcrKind::Sfx,
"title" => OcrKind::Title,
"caption" => OcrKind::Caption,
_ => unreachable!(),
}
);
}
for w in CONTENT_WARNINGS {
assert!(
ContentWarning::from_model_str(w).is_some(),
"prompt warning {w} not in the domain vocabulary"
);
}
}
#[test]
fn system_prompt_mentions_the_schema_keys() {
for key in ["ocr_results", "tagging_results", "scene_description", "safety_flag"] {
assert!(SYSTEM_PROMPT.contains(key), "prompt missing {key}");
}
}
}

View File

@@ -0,0 +1,280 @@
//! The OpenAI-compatible vision HTTP client.
//!
//! Flow: downscale the page image to fit the model's context budget →
//! base64 data-URL → chat/completions request with an `image_url` content
//! part → parse `choices[0].message.content` as JSON → sanitize/bound the
//! result. Parsing and sanitization are pure functions so they're tested
//! without a live model; only [`VisionClient::analyze`] does I/O.
use std::io::Cursor;
use anyhow::{anyhow, Context};
use base64::Engine;
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;
use crate::domain::page_analysis::{ContentWarning, 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,
max_image_dim: u32,
}
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,
max_image_dim: cfg.max_image_dim,
}
}
/// Analyze one page image. `mime` is the stored content type (used for
/// the data-URL when no downscale happens).
pub async fn analyze(&self, image: &[u8], mime: &str) -> anyhow::Result<VisionAnalysis> {
let (bytes, mime) = match downscale(image, self.max_image_dim) {
Some((b, m)) => (b, m.to_string()),
None => (image.to_vec(), mime.to_string()),
};
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
let data_url = format!("data:{mime};base64,{b64}");
let body = json!({
"model": self.model,
"temperature": 0,
"max_tokens": self.max_tokens,
"response_format": { "type": "json_object" },
"messages": [
{ "role": "system", "content": prompt::SYSTEM_PROMPT },
{ "role": "user", "content": [
{ "type": "image_url", "image_url": { "url": data_url } }
]}
]
});
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")?
.error_for_status()
.context("vision endpoint returned an error status")?;
let value: serde_json::Value =
resp.json().await.context("vision response was not JSON")?;
parse_chat_completion(&value)
}
}
/// Downscale `bytes` so its longest edge is `max_dim`, re-encoding as JPEG.
/// Returns `None` (use the original bytes) when the image already fits or
/// can't be decoded — a non-decodable page is passed through untouched and
/// the model server can reject it if it wants.
fn downscale(bytes: &[u8], max_dim: u32) -> Option<(Vec<u8>, &'static str)> {
let img = image::load_from_memory(bytes).ok()?;
if img.width().max(img.height()) <= max_dim {
return None;
}
let scaled = img.resize(max_dim, max_dim, image::imageops::FilterType::Triangle);
let mut buf = Vec::new();
scaled
.to_rgb8()
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Jpeg)
.ok()?;
Some((buf, "image/jpeg"))
}
/// Extract the model's JSON object from a chat/completions response.
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. Falls
/// back to the trimmed input if no braces are found (so the parse fails
/// with a clear error rather than silently matching nothing).
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 to [`MAX_TAGS`] via the page-tag normalizer, and keep only
/// recognized content-warning values. Persist applies the same rules as a
/// safety net, but doing it here keeps the row tidy and the worker's logs
/// honest about what it stored.
pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
a.ocr_results.retain(|r| !r.text.trim().is_empty());
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);
}
// Normalize tags through the shared page-tag rules (lowercase, collapse
// whitespace, reject wildcards/control/invisible chars); drop the
// unmappable, dedup, cap.
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);
// Keep only recognized warnings (canonical lowercase), deduped.
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
}
/// Truncate to at most `n` chars (not bytes), avoiding a panic on a
/// multi-byte boundary.
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 } } ] })
}
#[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 strips_markdown_fences_and_prose() {
let body = chat_body(
"Here you go:\n```json\n{\"scene_description\":\"x\"}\n```\nthanks",
);
let v = parse_chat_completion(&body).unwrap();
assert_eq!(v.scene_description, "x");
}
#[test]
fn rejects_non_json_content() {
let body = chat_body("I cannot analyze this image.");
assert!(parse_chat_completion(&body).is_err());
}
#[test]
fn rejects_missing_content() {
assert!(parse_chat_completion(&json!({"choices": []})).is_err());
}
#[test]
fn sanitize_clamps_tags_and_drops_invalid() {
let body = chat_body(
r#"{"tagging_results":["A","a","b","b","c","d","e","f","g","h","i","j",
"bad%wild"," "]}"#,
);
let v = parse_chat_completion(&body).unwrap();
// "A"/"a" dedup to one; wildcard + blank dropped; capped at 10.
assert!(v.tagging_results.len() <= MAX_TAGS);
assert!(v.tagging_results.contains(&"a".to_string()));
assert!(!v.tagging_results.iter().any(|t| t.contains('%')));
}
#[test]
fn sanitize_drops_empty_ocr_and_truncates() {
let long = "x".repeat(MAX_OCR_TEXT_CHARS + 50);
let body = chat_body(&format!(
r#"{{"ocr_results":[{{"text":" ","kind":"speech"}},
{{"text":"{long}","kind":"narration"}}]}}"#
));
let v = parse_chat_completion(&body).unwrap();
assert_eq!(v.ocr_results.len(), 1, "blank OCR dropped");
assert_eq!(v.ocr_results[0].text.chars().count(), MAX_OCR_TEXT_CHARS);
}
#[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"]);
}
#[test]
fn downscale_passes_through_small_images() {
// 8x8 PNG — well under any sane max_dim, so no re-encode.
let img = image::RgbImage::from_pixel(8, 8, image::Rgb([10, 20, 30]));
let mut buf = Vec::new();
image::DynamicImage::ImageRgb8(img)
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Png)
.unwrap();
assert!(downscale(&buf, 1024).is_none());
}
#[test]
fn downscale_shrinks_large_images_to_max_dim() {
let img = image::RgbImage::from_pixel(2000, 1000, image::Rgb([1, 2, 3]));
let mut buf = Vec::new();
image::DynamicImage::ImageRgb8(img)
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Png)
.unwrap();
let (out, mime) = downscale(&buf, 512).expect("should downscale");
assert_eq!(mime, "image/jpeg");
let decoded = image::load_from_memory(&out).unwrap();
assert_eq!(decoded.width().max(decoded.height()), 512);
}
#[test]
fn downscale_passes_through_undecodable_bytes() {
assert!(downscale(b"not an image", 512).is_none());
}
}