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
100
backend/Cargo.lock
generated
100
backend/Cargo.lock
generated
@@ -256,12 +256,24 @@ version = "3.20.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb"
|
||||
|
||||
[[package]]
|
||||
name = "bytemuck"
|
||||
version = "1.25.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec"
|
||||
|
||||
[[package]]
|
||||
name = "byteorder"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
|
||||
|
||||
[[package]]
|
||||
name = "byteorder-lite"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495"
|
||||
|
||||
[[package]]
|
||||
name = "bytes"
|
||||
version = "1.11.1"
|
||||
@@ -745,6 +757,15 @@ version = "2.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6"
|
||||
|
||||
[[package]]
|
||||
name = "fdeflate"
|
||||
version = "0.3.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c"
|
||||
dependencies = [
|
||||
"simd-adler32",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.9"
|
||||
@@ -1323,6 +1344,32 @@ dependencies = [
|
||||
"icu_properties",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "image"
|
||||
version = "0.25.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104"
|
||||
dependencies = [
|
||||
"bytemuck",
|
||||
"byteorder-lite",
|
||||
"image-webp",
|
||||
"moxcms",
|
||||
"num-traits",
|
||||
"png",
|
||||
"zune-core",
|
||||
"zune-jpeg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "image-webp"
|
||||
version = "0.2.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3"
|
||||
dependencies = [
|
||||
"byteorder-lite",
|
||||
"quick-error",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indexmap"
|
||||
version = "2.14.0"
|
||||
@@ -1470,7 +1517,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
|
||||
|
||||
[[package]]
|
||||
name = "mangalord"
|
||||
version = "0.66.0"
|
||||
version = "0.67.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"argon2",
|
||||
@@ -1486,6 +1533,7 @@ dependencies = [
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"http-body-util",
|
||||
"image",
|
||||
"infer",
|
||||
"mime",
|
||||
"nix 0.29.0",
|
||||
@@ -1582,6 +1630,16 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "moxcms"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
"pxfm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "multer"
|
||||
version = "3.1.0"
|
||||
@@ -2078,6 +2136,19 @@ version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
|
||||
|
||||
[[package]]
|
||||
name = "png"
|
||||
version = "0.18.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"crc32fast",
|
||||
"fdeflate",
|
||||
"flate2",
|
||||
"miniz_oxide",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "potential_utf"
|
||||
version = "0.1.5"
|
||||
@@ -2143,6 +2214,18 @@ dependencies = [
|
||||
"psl-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pxfm"
|
||||
version = "0.1.29"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
|
||||
|
||||
[[package]]
|
||||
name = "quick-error"
|
||||
version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3"
|
||||
|
||||
[[package]]
|
||||
name = "quinn"
|
||||
version = "0.11.9"
|
||||
@@ -4169,3 +4252,18 @@ name = "zmij"
|
||||
version = "1.0.21"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
|
||||
|
||||
[[package]]
|
||||
name = "zune-core"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cb8a0807f7c01457d0379ba880ba6322660448ddebc890ce29bb64da71fb40f9"
|
||||
|
||||
[[package]]
|
||||
name = "zune-jpeg"
|
||||
version = "0.5.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27bc9d5b815bc103f142aa054f561d9187d191692ec7c2d1e2b4737f8dbd7296"
|
||||
dependencies = [
|
||||
"zune-core",
|
||||
]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "mangalord"
|
||||
version = "0.66.0"
|
||||
version = "0.67.0"
|
||||
edition = "2021"
|
||||
default-run = "mangalord"
|
||||
|
||||
@@ -37,6 +37,9 @@ rand = "0.8"
|
||||
sha2 = "0.10"
|
||||
subtle = "2"
|
||||
base64 = "0.22"
|
||||
# Image decode + downscale for the analysis worker (keep the page image
|
||||
# under the local vision model's token budget). Only the manga page formats.
|
||||
image = { version = "0.25", default-features = false, features = ["jpeg", "png", "webp"] }
|
||||
axum-extra = { version = "0.9", features = ["cookie", "typed-header"] }
|
||||
time = "0.3"
|
||||
infer = "0.16"
|
||||
@@ -48,7 +51,7 @@ chromiumoxide = { version = "0.7", features = ["tokio-runtime", "_fetcher-rusttl
|
||||
sysinfo = { version = "0.32", default-features = false, features = ["system"] }
|
||||
nix = { version = "0.29", features = ["fs"] }
|
||||
scraper = "0.20"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "socks", "cookies", "stream"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "socks", "cookies", "stream", "json"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -115,7 +115,7 @@ fn is_invisible_format(c: char) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
fn normalize_tag(input: &str) -> AppResult<String> {
|
||||
pub(crate) fn normalize_tag(input: &str) -> AppResult<String> {
|
||||
// Fast-fail on absurd input before allocating in lowercase /
|
||||
// whitespace passes. Cheap and bounds worst-case work.
|
||||
if input.len() > MAX_TAG_BYTES {
|
||||
|
||||
@@ -66,27 +66,74 @@ impl Default for UploadConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// AI content-analysis worker configuration. Phase 2 only needs the
|
||||
/// `enabled` gate (so page-create paths know whether to enqueue analysis
|
||||
/// jobs); later phases extend this with the vision endpoint, model, and
|
||||
/// worker knobs.
|
||||
/// AI content-analysis worker configuration: the enable gate, the local
|
||||
/// OpenAI-compatible vision endpoint, and the worker / request knobs.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AnalysisConfig {
|
||||
/// Master switch (`ANALYSIS_ENABLED`). When `false`, no analysis jobs
|
||||
/// are enqueued and no worker runs. Defaults to `false`.
|
||||
pub enabled: bool,
|
||||
/// Number of concurrent analysis workers (`ANALYSIS_WORKERS`).
|
||||
pub workers: usize,
|
||||
/// OpenAI-compatible chat/completions URL (`ANALYSIS_VISION_URL`).
|
||||
pub endpoint: String,
|
||||
/// Model id to request (`ANALYSIS_MODEL`).
|
||||
pub model: String,
|
||||
/// Optional bearer token (`ANALYSIS_API_KEY`); local servers usually
|
||||
/// don't need one.
|
||||
pub api_key: Option<String>,
|
||||
/// Per-request HTTP timeout (`ANALYSIS_REQUEST_TIMEOUT_SECS`).
|
||||
pub request_timeout: Duration,
|
||||
/// Whole-job timeout in the worker (`ANALYSIS_JOB_TIMEOUT_SECS`).
|
||||
pub job_timeout: Duration,
|
||||
/// Output token cap sent as `max_tokens` (`ANALYSIS_MAX_TOKENS`).
|
||||
pub max_tokens: u32,
|
||||
/// Longest image edge (px) before downscaling (`ANALYSIS_MAX_IMAGE_DIM`).
|
||||
pub max_image_dim: u32,
|
||||
/// Hard cap on a page image's stored size; larger pages are skipped
|
||||
/// (`ANALYSIS_MAX_IMAGE_BYTES`).
|
||||
pub max_image_bytes: usize,
|
||||
}
|
||||
|
||||
impl Default for AnalysisConfig {
|
||||
fn default() -> Self {
|
||||
Self { enabled: false }
|
||||
Self {
|
||||
enabled: false,
|
||||
workers: 1,
|
||||
endpoint: "http://localhost:8000/v1/chat/completions".to_string(),
|
||||
model: String::new(),
|
||||
api_key: None,
|
||||
request_timeout: Duration::from_secs(120),
|
||||
job_timeout: Duration::from_secs(180),
|
||||
max_tokens: 900,
|
||||
max_image_dim: 1024,
|
||||
max_image_bytes: 8 * 1024 * 1024,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AnalysisConfig {
|
||||
pub fn from_env() -> Self {
|
||||
let d = AnalysisConfig::default();
|
||||
Self {
|
||||
enabled: env_bool("ANALYSIS_ENABLED", false),
|
||||
enabled: env_bool("ANALYSIS_ENABLED", d.enabled),
|
||||
workers: env_usize("ANALYSIS_WORKERS", d.workers).max(1),
|
||||
endpoint: std::env::var("ANALYSIS_VISION_URL").unwrap_or(d.endpoint),
|
||||
model: std::env::var("ANALYSIS_MODEL").unwrap_or(d.model),
|
||||
api_key: std::env::var("ANALYSIS_API_KEY")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty()),
|
||||
request_timeout: Duration::from_secs(env_u64(
|
||||
"ANALYSIS_REQUEST_TIMEOUT_SECS",
|
||||
d.request_timeout.as_secs(),
|
||||
)),
|
||||
job_timeout: Duration::from_secs(env_u64(
|
||||
"ANALYSIS_JOB_TIMEOUT_SECS",
|
||||
d.job_timeout.as_secs(),
|
||||
)),
|
||||
max_tokens: env_u64("ANALYSIS_MAX_TOKENS", d.max_tokens as u64) as u32,
|
||||
max_image_dim: env_u64("ANALYSIS_MAX_IMAGE_DIM", d.max_image_dim as u64) as u32,
|
||||
max_image_bytes: env_usize("ANALYSIS_MAX_IMAGE_BYTES", d.max_image_bytes),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -484,6 +531,52 @@ mod tests {
|
||||
assert_eq!(cfg.browser_restart_threshold, 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn analysis_config_defaults_when_unset() {
|
||||
let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner());
|
||||
for k in [
|
||||
"ANALYSIS_ENABLED",
|
||||
"ANALYSIS_WORKERS",
|
||||
"ANALYSIS_VISION_URL",
|
||||
"ANALYSIS_MODEL",
|
||||
"ANALYSIS_API_KEY",
|
||||
"ANALYSIS_MAX_TOKENS",
|
||||
"ANALYSIS_MAX_IMAGE_DIM",
|
||||
] {
|
||||
std::env::remove_var(k);
|
||||
}
|
||||
let cfg = AnalysisConfig::from_env();
|
||||
assert!(!cfg.enabled);
|
||||
assert_eq!(cfg.workers, 1);
|
||||
assert_eq!(cfg.max_image_dim, 1024);
|
||||
assert!(cfg.api_key.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn analysis_config_parses_from_env() {
|
||||
let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner());
|
||||
std::env::set_var("ANALYSIS_ENABLED", "true");
|
||||
std::env::set_var("ANALYSIS_WORKERS", "4");
|
||||
std::env::set_var("ANALYSIS_VISION_URL", "http://vis/v1/chat");
|
||||
std::env::set_var("ANALYSIS_MODEL", "qwen2-vl");
|
||||
std::env::set_var("ANALYSIS_MAX_IMAGE_DIM", "768");
|
||||
let cfg = AnalysisConfig::from_env();
|
||||
for k in [
|
||||
"ANALYSIS_ENABLED",
|
||||
"ANALYSIS_WORKERS",
|
||||
"ANALYSIS_VISION_URL",
|
||||
"ANALYSIS_MODEL",
|
||||
"ANALYSIS_MAX_IMAGE_DIM",
|
||||
] {
|
||||
std::env::remove_var(k);
|
||||
}
|
||||
assert!(cfg.enabled);
|
||||
assert_eq!(cfg.workers, 4);
|
||||
assert_eq!(cfg.endpoint, "http://vis/v1/chat");
|
||||
assert_eq!(cfg.model, "qwen2-vl");
|
||||
assert_eq!(cfg.max_image_dim, 768);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn private_mode_env_parses_true() {
|
||||
let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner());
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod analysis;
|
||||
pub mod api;
|
||||
pub mod app;
|
||||
pub mod auth;
|
||||
|
||||
Reference in New Issue
Block a user