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:
11
backend/src/analysis/mod.rs
Normal file
11
backend/src/analysis/mod.rs
Normal 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;
|
||||
90
backend/src/analysis/prompt.rs
Normal file
90
backend/src/analysis/prompt.rs
Normal 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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
280
backend/src/analysis/vision.rs
Normal file
280
backend/src/analysis/vision.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user