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
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