feat(analysis): add ocrs OCR backend and OCR-driven page text search
All checks were successful
deploy / test-backend (push) Successful in 28m40s
deploy / test-frontend (push) Successful in 10m36s
deploy / build-and-push (push) Successful in 11m8s
deploy / deploy (push) Successful in 12s

Adds an in-process ocrs OCR backend as the active analysis engine
(vision left dormant), enables OCR text search on the page-tag
aggregation endpoints, and reshapes the admin analysis/settings UI to
present the OCR-only surface. Bumps version 0.89.0 -> 0.90.0.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-06-30 19:52:43 +02:00
parent 4fb98e4a1e
commit 83c2899373
24 changed files with 1335 additions and 443 deletions

View File

@@ -9,5 +9,6 @@
pub mod daemon;
pub mod events;
pub mod ocr;
pub mod prompt;
pub mod vision;

201
backend/src/analysis/ocr.rs Normal file
View File

@@ -0,0 +1,201 @@
//! The in-process OCR analysis backend (the `ocrs` engine).
//!
//! A lightweight alternative to [`crate::analysis::vision`]: instead of a slow
//! local LLM, each page is run through `ocrs` — a pure-Rust detect→recognize
//! OCR pipeline on the `rten` runtime. It extracts **text only** (no tags,
//! scene description or NSFW flags — those stay the vision backend's job), then
//! reuses [`repo::page_analysis::persist_analysis`] so the OCR lines land in
//! `page_ocr_text` and the weighted `search_doc` tsvector exactly as the vision
//! path produces them. That makes the existing text-search surfaces
//! (`/v1/me/page-search` and the tag aggregations) work with no further wiring.
//!
//! The engine is split behind the [`OcrEngine`] trait so the dispatcher is
//! unit-testable without shipping the (multi-MB) `.rten` model files: tests use
//! [`test_support::StubOcrEngine`], production uses [`OcrsEngine`].
use std::sync::Arc;
use async_trait::async_trait;
use sqlx::PgPool;
use uuid::Uuid;
use crate::analysis::daemon::AnalyzeDispatcher;
use crate::domain::page_analysis::{OcrResult, SafetyFlag, VisionAnalysis};
use crate::repo;
use crate::storage::Storage;
/// The `model` label stamped onto `page_analysis` rows written by this backend.
pub const OCR_MODEL_LABEL: &str = "ocrs";
/// Extracts text lines from a decoded-or-encoded page image. The production
/// impl ([`OcrsEngine`]) decodes the bytes itself; the trait takes the raw
/// stored image bytes so the dispatcher stays engine-agnostic.
pub trait OcrEngine: Send + Sync {
/// Run OCR over one page image (the bytes as stored, e.g. PNG/JPEG/WebP).
/// Returns the recognized text lines in reading order (top→bottom).
fn recognize(&self, image: &[u8]) -> anyhow::Result<Vec<String>>;
}
/// Turn the OCR engine's ordered text lines into the [`VisionAnalysis`] shape
/// that [`repo::page_analysis::persist_analysis`] consumes. OCR-only: the tag,
/// scene and safety fields are left empty/default. The line `kind` is left
/// blank — `persist_analysis` maps an empty kind to the neutral mid-weight
/// `OcrKind::Narration` bucket (classifying speech/sfx/… is the deferred
/// vision backend's job).
pub fn lines_to_analysis(lines: Vec<String>) -> VisionAnalysis {
let ocr_results = lines
.into_iter()
.map(|text| OcrResult { text, kind: String::new(), y: None })
.collect();
VisionAnalysis {
ocr_results,
tagging_results: Vec::new(),
scene_description: String::new(),
safety_flag: SafetyFlag::default(),
}
}
/// Production OCR engine: an `ocrs` detect+recognize pipeline with the two
/// `.rten` models loaded once at startup. Cheap to share across workers — the
/// recognize path borrows `&self`.
pub struct OcrsEngine {
engine: ocrs::OcrEngine,
}
impl OcrsEngine {
/// Load the detection + recognition models from disk and build the engine.
/// Fails (at startup) if either model file is missing or unreadable, so a
/// misconfigured path is a loud boot error rather than a per-page failure.
pub fn from_model_paths(detection: &str, recognition: &str) -> anyhow::Result<Self> {
use anyhow::Context;
let detection_model = rten::Model::load_file(detection)
.with_context(|| format!("load ocrs detection model {detection}"))?;
let recognition_model = rten::Model::load_file(recognition)
.with_context(|| format!("load ocrs recognition model {recognition}"))?;
let engine = ocrs::OcrEngine::new(ocrs::OcrEngineParams {
detection_model: Some(detection_model),
recognition_model: Some(recognition_model),
..Default::default()
})
.context("construct ocrs engine")?;
Ok(Self { engine })
}
}
impl OcrEngine for OcrsEngine {
fn recognize(&self, image: &[u8]) -> anyhow::Result<Vec<String>> {
use anyhow::Context;
// Decode to RGB8 so `ImageSource` gets a known channel layout.
let rgb = image::load_from_memory(image)
.context("decode page image for OCR")?
.into_rgb8();
let source = ocrs::ImageSource::from_bytes(rgb.as_raw(), rgb.dimensions())
.map_err(|e| anyhow::anyhow!("build OCR image source: {e}"))?;
let input = self.engine.prepare_input(source)?;
// detect words → group into lines → recognize each line. Mirrors
// `OcrEngine::get_text`, but keeps the lines as a Vec instead of
// joining them, so each becomes its own `page_ocr_text` row.
let words = self.engine.detect_words(&input)?;
let line_rects = self.engine.find_text_lines(&input, &words);
let lines = self
.engine
.recognize_text(&input, &line_rects)?
.into_iter()
.filter_map(|line| line.map(|l| l.to_string()))
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
Ok(lines)
}
}
/// Production dispatcher for the OCR backend: load the page, read its image
/// from storage, run OCR on the blocking pool, and persist the lines. Mirrors
/// [`crate::analysis::daemon::RealAnalyzeDispatcher`] but with no network I/O.
pub struct OcrAnalyzeDispatcher {
pub db: PgPool,
pub storage: Arc<dyn Storage>,
pub engine: Arc<dyn OcrEngine>,
pub max_image_bytes: usize,
}
#[async_trait]
impl AnalyzeDispatcher for OcrAnalyzeDispatcher {
async fn dispatch(&self, page_id: Uuid) -> anyhow::Result<()> {
let Some(page) = repo::page::find_by_id(&self.db, page_id).await? else {
// Page was deleted between enqueue and dispatch — nothing to do.
return Ok(());
};
let bytes = self
.storage
.get(&page.storage_key)
.await
.map_err(|e| anyhow::anyhow!("read page image {}: {e}", page.storage_key))?;
if bytes.len() > self.max_image_bytes {
anyhow::bail!(
"page image {} is {} bytes, over the {} cap",
page.storage_key,
bytes.len(),
self.max_image_bytes
);
}
// OCR inference is CPU-bound and synchronous — keep it off the async
// worker's runtime thread.
let engine = Arc::clone(&self.engine);
let lines = tokio::task::spawn_blocking(move || engine.recognize(&bytes))
.await
.map_err(|e| anyhow::anyhow!("OCR task join error: {e}"))??;
let analysis = lines_to_analysis(lines);
repo::page_analysis::persist_analysis(&self.db, page_id, &analysis, OCR_MODEL_LABEL).await?;
Ok(())
}
}
/// Stubs for the OCR dispatcher's integration tests. Public because the tests
/// live in the `tests/` dir (a separate crate).
pub mod test_support {
use super::*;
/// An [`OcrEngine`] that returns a fixed set of lines regardless of input,
/// so the dispatcher's storage→persist path can be tested without models.
pub struct StubOcrEngine {
pub lines: Vec<String>,
}
impl StubOcrEngine {
pub fn new(lines: &[&str]) -> Arc<Self> {
Arc::new(Self { lines: lines.iter().map(|s| s.to_string()).collect() })
}
}
impl OcrEngine for StubOcrEngine {
fn recognize(&self, _image: &[u8]) -> anyhow::Result<Vec<String>> {
Ok(self.lines.clone())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn lines_to_analysis_maps_lines_in_order_and_leaves_rest_empty() {
let v = lines_to_analysis(vec!["Hello".to_string(), "world!".to_string()]);
assert_eq!(v.ocr_results.len(), 2);
assert_eq!(v.ocr_results[0].text, "Hello");
assert_eq!(v.ocr_results[1].text, "world!");
// OCR-only: kind blank (→ Narration at persist), no tags/scene/safety.
assert!(v.ocr_results.iter().all(|r| r.kind.is_empty()));
assert!(v.tagging_results.is_empty());
assert_eq!(v.scene_description, "");
assert!(!v.safety_flag.is_nsfw);
assert!(v.safety_flag.content_type.is_empty());
}
#[test]
fn lines_to_analysis_handles_no_text() {
let v = lines_to_analysis(Vec::new());
assert!(v.ocr_results.is_empty());
}
}

View File

@@ -392,10 +392,9 @@ pub struct AggregateParams {
pub limit: i64,
#[serde(default)]
pub offset: i64,
/// Reserved for the planned OCR text-search input. Accepted on
/// the wire so adding OCR later won't break the API shape, but
/// rejected with 501 `text_search_not_yet_supported` if non-empty
/// until the backend supports it.
/// OCR text filter. When non-empty, only pages whose analysis
/// `search_doc` matches the query (`plainto_tsquery`) are aggregated.
/// Blank/absent ⇒ tag-only aggregation.
#[serde(default)]
pub text: Option<String>,
}
@@ -411,32 +410,17 @@ fn parse_order(raw: Option<&str>) -> AppResult<Order> {
}
}
fn ensure_text_unsupported(text: Option<&str>) -> AppResult<()> {
// Future OCR search will plug in here. Until then, return a
// distinct code (`text_search_not_yet_supported`) so clients can
// detect "feature pending" vs. a generic 4xx — the code is the
// wire contract, not the message.
if text.is_some_and(|s| !s.trim().is_empty()) {
return Err(AppError::NotImplemented {
code: "text_search_not_yet_supported",
message: "text search is reserved for the planned OCR input but not yet supported",
});
}
Ok(())
}
async fn list_chapters_for_tag(
State(state): State<AppState>,
CurrentUser(user): CurrentUser,
Query(params): Query<AggregateParams>,
) -> AppResult<Json<PagedResponse<TaggedChapterAggregate>>> {
ensure_text_unsupported(params.text.as_deref())?;
let tag = normalize_tag(&params.tag)?;
let order = parse_order(params.order.as_deref())?;
let limit = params.limit.clamp(1, 200);
let offset = params.offset.max(0);
let (items, total) = repo::page_tag::aggregate_chapters_for_tag(
&state.db, user.id, &tag, order, limit, offset,
&state.db, user.id, &tag, order, limit, offset, params.text.as_deref(),
)
.await?;
Ok(Json(PagedResponse::with_total(items, limit, offset, total)))
@@ -447,13 +431,12 @@ async fn list_mangas_for_tag(
CurrentUser(user): CurrentUser,
Query(params): Query<AggregateParams>,
) -> AppResult<Json<PagedResponse<TaggedMangaAggregate>>> {
ensure_text_unsupported(params.text.as_deref())?;
let tag = normalize_tag(&params.tag)?;
let order = parse_order(params.order.as_deref())?;
let limit = params.limit.clamp(1, 200);
let offset = params.offset.max(0);
let (items, total) = repo::page_tag::aggregate_mangas_for_tag(
&state.db, user.id, &tag, order, limit, offset,
&state.db, user.id, &tag, order, limit, offset, params.text.as_deref(),
)
.await?;
Ok(Json(PagedResponse::with_total(items, limit, offset, total)))

View File

@@ -429,46 +429,74 @@ async fn spawn_analysis_daemon(
),
Err(e) => tracing::warn!(?e, "analysis: reclaim_orphaned at startup failed"),
}
let http = reqwest::Client::builder()
.timeout(cfg.request_timeout)
// Refuse to honour ambient HTTP_PROXY / HTTPS_PROXY container env.
// The vision call carries an env-managed bearer token + page image
// bytes; a stray upstream proxy would exfiltrate both. Mirrors the
// crawler client's `.no_proxy()` (see `spawn_crawler_daemon`).
.no_proxy()
.build()
.context("build analysis http client")?;
let vision = crate::analysis::vision::VisionClient::new(http, cfg);
let dispatcher = Arc::new(crate::analysis::daemon::RealAnalyzeDispatcher {
db: db.clone(),
storage,
vision,
model: cfg.model.clone(),
max_image_bytes: cfg.max_image_bytes,
});
// When a readiness URL is configured, gate leasing on it so an
// autoscaler that idle-stops the vision container never lets a job burn
// its retries. A dedicated short-timeout client keeps the probe snappy
// and independent of the (long) per-request analysis timeout.
let readiness: Option<Arc<dyn crate::analysis::daemon::VisionReadiness>> =
match &cfg.vision_health_url {
Some(url) if !url.is_empty() => {
let probe = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(5))
// Same reasoning as the main analysis client: do not
// honour ambient HTTP_PROXY env. The readiness probe is
// unauthenticated but a hostile upstream still gets a
// useful side-channel on backend uptime + vision health.
.no_proxy()
.build()
.context("build vision readiness http client")?;
Some(Arc::new(crate::analysis::daemon::HttpVisionReadiness {
http: probe,
health_url: url.clone(),
}))
}
_ => None,
};
// Pick the engine. The OCR backend runs in-process (no network, no
// readiness gate); the vision backend talks to a local LLM server.
let (dispatcher, readiness): (
Arc<dyn crate::analysis::daemon::AnalyzeDispatcher>,
Option<Arc<dyn crate::analysis::daemon::VisionReadiness>>,
) = match cfg.effective_backend() {
crate::config::AnalysisBackend::Ocr => {
// Load the `.rten` models once; a bad path is a loud boot error.
let engine = crate::analysis::ocr::OcrsEngine::from_model_paths(
&cfg.ocr_detection_model,
&cfg.ocr_recognition_model,
)
.context("build ocrs engine")?;
let dispatcher = Arc::new(crate::analysis::ocr::OcrAnalyzeDispatcher {
db: db.clone(),
storage,
engine: Arc::new(engine),
max_image_bytes: cfg.max_image_bytes,
});
// In-process engine is always ready — no gate.
(dispatcher, None)
}
crate::config::AnalysisBackend::Vision => {
let http = reqwest::Client::builder()
.timeout(cfg.request_timeout)
// Refuse to honour ambient HTTP_PROXY / HTTPS_PROXY container
// env. The vision call carries an env-managed bearer token +
// page image bytes; a stray upstream proxy would exfiltrate
// both. Mirrors the crawler client's `.no_proxy()` (see
// `spawn_crawler_daemon`).
.no_proxy()
.build()
.context("build analysis http client")?;
let vision = crate::analysis::vision::VisionClient::new(http, cfg);
let dispatcher = Arc::new(crate::analysis::daemon::RealAnalyzeDispatcher {
db: db.clone(),
storage,
vision,
model: cfg.model.clone(),
max_image_bytes: cfg.max_image_bytes,
});
// When a readiness URL is configured, gate leasing on it so an
// autoscaler that idle-stops the vision container never lets a job
// burn its retries. A dedicated short-timeout client keeps the
// probe snappy and independent of the (long) per-request timeout.
let readiness: Option<Arc<dyn crate::analysis::daemon::VisionReadiness>> =
match &cfg.vision_health_url {
Some(url) if !url.is_empty() => {
let probe = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(5))
// Same reasoning as the main analysis client: do
// not honour ambient HTTP_PROXY env. The readiness
// probe is unauthenticated but a hostile upstream
// still gets a useful side-channel on backend
// uptime + vision health.
.no_proxy()
.build()
.context("build vision readiness http client")?;
Some(Arc::new(crate::analysis::daemon::HttpVisionReadiness {
http: probe,
health_url: url.clone(),
}))
}
_ => None,
};
(dispatcher, readiness)
}
};
let handle = crate::analysis::daemon::spawn(
db,
CancellationToken::new(),
@@ -480,7 +508,21 @@ async fn spawn_analysis_daemon(
readiness,
},
);
tracing::info!(workers = cfg.workers, model = %cfg.model, "analysis worker daemon started");
// Log the *effective* backend (what the worker actually dispatches
// through), not the raw `cfg.backend` — with vision dormant they diverge
// when an operator requests `vision`, and a misleading line here was a
// real observability footgun. The model label tracks the effective engine.
let effective_backend = cfg.effective_backend();
let effective_model: &str = match effective_backend {
crate::config::AnalysisBackend::Ocr => crate::analysis::ocr::OCR_MODEL_LABEL,
crate::config::AnalysisBackend::Vision => &cfg.model,
};
tracing::info!(
workers = cfg.workers,
backend = ?effective_backend,
model = %effective_model,
"analysis worker daemon started"
);
Ok(handle)
}
@@ -1440,25 +1482,30 @@ mod tests {
// Bind to a local so the TempDir lives for the rest of the test.
// `tempfile::tempdir().unwrap().path()` would drop the TempDir
// at end-of-expression and `LocalStorage` would hold a path to
// a deleted directory. (Today this is fine because the dispatch
// fails before storage is touched, but it makes the test fragile
// to any future code rearrangement.)
// a deleted directory.
let storage_dir = tempfile::tempdir().unwrap();
let storage: Arc<dyn Storage> =
Arc::new(LocalStorage::new(storage_dir.path()));
let mut cfg = crate::config::AnalysisConfig::default();
// The worker always runs OCR now (vision is dormant — see
// `effective_backend`), and the `.rten` models aren't shipped to unit
// CI. Point the engine at a path that can't exist so the *engine build*
// fails deterministically. Reclaim runs at the very top of
// `spawn_analysis_daemon`, before — and independently of — engine
// readiness, so the row must still be reclaimed even though spawn
// returns Err. That's exactly the regression this test guards (an
// analysis-only deploy must reclaim orphaned leases at startup).
cfg.ocr_detection_model = "/nonexistent/text-detection.rten".to_string();
cfg.ocr_recognition_model = "/nonexistent/text-recognition.rten".to_string();
cfg.workers = 1;
cfg.job_timeout = Duration::from_secs(1);
let events = Arc::new(crate::analysis::events::AnalysisEvents::new());
let handle = spawn_analysis_daemon(pool.clone(), storage, &cfg, events)
.await
.expect("spawn");
// Immediately shut down — we're only here to prove reclaim ran.
// The workers may briefly pick up the now-pending row; that's
// fine, but we cancel before any real dispatch (the LocalStorage
// key doesn't exist, so a dispatch would fail anyway).
handle.shutdown().await;
let spawned = spawn_analysis_daemon(pool.clone(), storage, &cfg, events).await;
assert!(
spawned.is_err(),
"engine build must fail with a missing model path"
);
// The reclaim must have moved the row back to pending with the
// attempt refunded (attempts goes from 1 → 0). Reaching it on the

View File

@@ -113,6 +113,32 @@ impl ResponseFormat {
}
}
/// Which engine the analysis worker dispatches each page through. A
/// deploy-time choice (the engine is either installed or not), so it lives in
/// env only and is not part of the admin-editable `AnalysisSettings`.
///
/// * `Ocr` — in-process [`crate::analysis::ocr`] (the `ocrs` engine): fast,
/// CPU-only, English text. Writes OCR text only (no tags/scene/safety).
/// * `Vision` — the local OpenAI-compatible LLM in [`crate::analysis::vision`]:
/// full OCR + tags + scene + safety, but heavy.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AnalysisBackend {
Ocr,
Vision,
}
impl AnalysisBackend {
/// Lenient env parse. Defaults to `Ocr` (the Pi-friendly path) for any
/// unset or unrecognized value; `vision` opts back into the LLM engine.
fn from_str(s: &str) -> AnalysisBackend {
match s.trim().to_lowercase().as_str() {
"vision" | "llm" => AnalysisBackend::Vision,
// Default (incl. "ocr", "ocrs", and anything unrecognized).
_ => AnalysisBackend::Ocr,
}
}
}
/// AI content-analysis worker configuration: the enable gate, the local
/// OpenAI-compatible vision endpoint, and the worker / request knobs.
#[derive(Clone, Debug)]
@@ -120,6 +146,16 @@ pub struct AnalysisConfig {
/// Master switch (`ANALYSIS_ENABLED`). When `false`, no analysis jobs
/// are enqueued and no worker runs. Defaults to `false`.
pub enabled: bool,
/// Which engine the worker dispatches through (`ANALYSIS_BACKEND`):
/// `ocr` (default, the in-process `ocrs` engine) or `vision` (the local
/// LLM). Deploy-time, env-only — see [`AnalysisBackend`].
pub backend: AnalysisBackend,
/// Path to the ocrs text-*detection* `.rten` model
/// (`OCRS_DETECTION_MODEL`). Only read when `backend == Ocr`.
pub ocr_detection_model: String,
/// Path to the ocrs text-*recognition* `.rten` model
/// (`OCRS_RECOGNITION_MODEL`). Only read when `backend == Ocr`.
pub ocr_recognition_model: String,
/// Number of concurrent analysis workers (`ANALYSIS_WORKERS`).
pub workers: usize,
/// OpenAI-compatible chat/completions URL (`ANALYSIS_VISION_URL`).
@@ -195,6 +231,9 @@ impl Default for AnalysisConfig {
fn default() -> Self {
Self {
enabled: false,
backend: AnalysisBackend::Ocr,
ocr_detection_model: "/models/text-detection.rten".to_string(),
ocr_recognition_model: "/models/text-recognition.rten".to_string(),
workers: 1,
endpoint: "http://localhost:8000/v1/chat/completions".to_string(),
vision_health_url: None,
@@ -223,10 +262,39 @@ impl Default for AnalysisConfig {
}
impl AnalysisConfig {
/// The backend the worker actually dispatches through.
///
/// Vision is **temporarily disabled**: the engine code (`analysis::vision`,
/// `RealAnalyzeDispatcher`, the readiness probe) is kept intact but never
/// selected. Until it's re-enabled, this returns [`AnalysisBackend::Ocr`]
/// regardless of the parsed `backend`, logging a warning if `vision` was
/// requested so an env override isn't silently ignored. Re-enabling vision
/// is then a one-line change here (return `self.backend`).
pub fn effective_backend(&self) -> AnalysisBackend {
if self.backend == AnalysisBackend::Vision {
tracing::warn!(
"ANALYSIS_BACKEND=vision requested but the vision backend is temporarily \
disabled; running OCR instead"
);
}
AnalysisBackend::Ocr
}
pub fn from_env() -> Self {
let d = AnalysisConfig::default();
Self {
enabled: env_bool("ANALYSIS_ENABLED", d.enabled),
backend: std::env::var("ANALYSIS_BACKEND")
.map(|s| AnalysisBackend::from_str(&s))
.unwrap_or(d.backend),
ocr_detection_model: std::env::var("OCRS_DETECTION_MODEL")
.ok()
.filter(|s| !s.is_empty())
.unwrap_or(d.ocr_detection_model),
ocr_recognition_model: std::env::var("OCRS_RECOGNITION_MODEL")
.ok()
.filter(|s| !s.is_empty())
.unwrap_or(d.ocr_recognition_model),
workers: env_usize("ANALYSIS_WORKERS", d.workers).max(1),
endpoint: std::env::var("ANALYSIS_VISION_URL").unwrap_or(d.endpoint),
vision_health_url: std::env::var("ANALYSIS_VISION_HEALTH_URL")
@@ -734,11 +802,18 @@ mod tests {
"ANALYSIS_MAX_SLICES",
"ANALYSIS_RESPONSE_FORMAT",
"ANALYSIS_FREQUENCY_PENALTY",
"ANALYSIS_BACKEND",
"OCRS_DETECTION_MODEL",
"OCRS_RECOGNITION_MODEL",
] {
std::env::remove_var(k);
}
let cfg = AnalysisConfig::from_env();
assert!(!cfg.enabled);
// OCR is the default engine (the Pi-friendly path).
assert_eq!(cfg.backend, AnalysisBackend::Ocr);
assert_eq!(cfg.ocr_detection_model, "/models/text-detection.rten");
assert_eq!(cfg.ocr_recognition_model, "/models/text-recognition.rten");
assert_eq!(cfg.workers, 1);
assert_eq!(cfg.max_pixels, 1_000_000);
assert_eq!(cfg.min_slice_height, 640);
@@ -765,6 +840,49 @@ mod tests {
std::env::remove_var("ANALYSIS_RESPONSE_FORMAT");
}
#[test]
fn analysis_backend_parses_and_defaults_to_ocr() {
let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner());
for (raw, want) in [
("ocr", AnalysisBackend::Ocr),
("ocrs", AnalysisBackend::Ocr),
("vision", AnalysisBackend::Vision),
("llm", AnalysisBackend::Vision),
("anything-else", AnalysisBackend::Ocr),
] {
std::env::set_var("ANALYSIS_BACKEND", raw);
assert_eq!(AnalysisConfig::from_env().backend, want, "raw={raw}");
}
// Unset → OCR.
std::env::remove_var("ANALYSIS_BACKEND");
assert_eq!(AnalysisConfig::from_env().backend, AnalysisBackend::Ocr);
}
#[test]
fn effective_backend_is_ocr_even_when_vision_requested() {
// Vision is temporarily disabled: the worker always runs OCR no matter
// what `ANALYSIS_BACKEND` parsed to. `backend` still reflects the raw
// request (so the override is visible/loggable), but `effective_backend`
// is the value the daemon actually dispatches through.
let mut cfg = AnalysisConfig::default();
cfg.backend = AnalysisBackend::Vision;
assert_eq!(cfg.effective_backend(), AnalysisBackend::Ocr);
cfg.backend = AnalysisBackend::Ocr;
assert_eq!(cfg.effective_backend(), AnalysisBackend::Ocr);
}
#[test]
fn ocr_model_paths_parse_from_env() {
let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner());
std::env::set_var("OCRS_DETECTION_MODEL", "/opt/det.rten");
std::env::set_var("OCRS_RECOGNITION_MODEL", "/opt/rec.rten");
let cfg = AnalysisConfig::from_env();
std::env::remove_var("OCRS_DETECTION_MODEL");
std::env::remove_var("OCRS_RECOGNITION_MODEL");
assert_eq!(cfg.ocr_detection_model, "/opt/det.rten");
assert_eq!(cfg.ocr_recognition_model, "/opt/rec.rten");
}
#[test]
fn analysis_config_parses_from_env() {
let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner());

View File

@@ -39,10 +39,10 @@ pub enum AppError {
details: serde_json::Value,
},
/// 501 — the wire shape is accepted but the feature isn't built yet.
/// Carries a `&'static str` snake_case code so clients can detect
/// the specific pending feature (`text_search_not_yet_supported`,
/// etc.) without parsing the message. Used today by the `?text=`
/// reservation on the page-tag aggregation endpoints.
/// Carries a `&'static str` snake_case code so clients can detect the
/// specific pending feature without parsing the message. Generic; kept
/// for future reservations (the `?text=` page-tag reservation that first
/// used it now performs real OCR search).
#[error("not implemented: {code}")]
NotImplemented {
code: &'static str,

View File

@@ -247,6 +247,11 @@ pub async fn distinct_tags_for_user(
/// storage keys (page-number ascending) so the row can render a
/// thumbnail strip without a follow-up fetch.
///
/// When `text` is non-blank, results are further restricted to pages whose
/// analysis `search_doc` matches the query (OCR text search), and both
/// `match_count` and the sample thumbnails reflect that filtered set —
/// i.e. `match_count` counts tagged pages whose OCR also matches `text`.
///
/// `order` is inlined via `format!()` — the enum value space is
/// closed (`ASC` / `DESC`) so this is not a SQL-injection vector.
pub async fn aggregate_chapters_for_tag(
@@ -256,7 +261,30 @@ pub async fn aggregate_chapters_for_tag(
order: Order,
limit: i64,
offset: i64,
text: Option<&str>,
) -> AppResult<(Vec<TaggedChapterAggregate>, i64)> {
// OCR text filter: when present, additionally require the page's analysis
// `search_doc` to match the query. Reuses the precomputed tsvector exactly
// like `repo::page_analysis::page_search`. `None`/blank ⇒ tag-only.
let text = text.map(str::trim).filter(|s| !s.is_empty());
let (text_join, text_where) = match text {
// `$5` in the main query, `$3` in the count query (see binds below).
Some(_) => (
"JOIN page_analysis pa ON pa.page_id = p.id",
"AND pa.search_doc @@ plainto_tsquery('simple', {n})",
),
None => ("", ""),
};
// Same filter inside the correlated sample-thumbnail subquery (its page is
// aliased `p`), so the thumbnails match what the text search matched. The
// subquery lives in the main query, so it reuses the `$5` text bind.
let (sample_join, sample_where) = match text {
Some(_) => (
"JOIN page_analysis pa2 ON pa2.page_id = p.id",
"AND pa2.search_doc @@ plainto_tsquery('simple', $5)",
),
None => ("", ""),
};
let sql = format!(
r#"
SELECT
@@ -274,9 +302,11 @@ pub async fn aggregate_chapters_for_tag(
FROM pages p
JOIN page_tags pt2 ON pt2.page_id = p.id
JOIN tags t2 ON t2.id = pt2.tag_id
{sample_join}
WHERE p.chapter_id = ch.id
AND pt2.user_id = $1
AND lower(t2.name) = $2
{sample_where}
ORDER BY p.page_number ASC
LIMIT 3
) p2
@@ -288,43 +318,60 @@ pub async fn aggregate_chapters_for_tag(
JOIN pages p ON p.id = pt.page_id
JOIN chapters ch ON ch.id = p.chapter_id
JOIN mangas m ON m.id = ch.manga_id
{text_join}
WHERE pt.user_id = $1
AND lower(t.name) = $2
{text_where}
GROUP BY ch.id, ch.manga_id, m.title, ch.number, ch.title
ORDER BY match_count {dir}, ch.id
LIMIT $3 OFFSET $4
"#,
dir = order.as_sql(),
text_join = text_join,
text_where = text_where.replace("{n}", "$5"),
sample_join = sample_join,
sample_where = sample_where,
);
let rows = sqlx::query_as::<_, TaggedChapterAggregate>(&sql)
let mut q = sqlx::query_as::<_, TaggedChapterAggregate>(&sql)
.bind(user_id)
.bind(tag)
.bind(limit)
.bind(offset)
.fetch_all(pool)
.await?;
.bind(offset);
if let Some(text) = text {
q = q.bind(text);
}
let rows = q.fetch_all(pool).await?;
let (total,): (i64,) = sqlx::query_as(
let count_sql = format!(
r#"
SELECT count(*) FROM (
SELECT 1
FROM page_tags pt
JOIN tags t ON t.id = pt.tag_id
JOIN pages p ON p.id = pt.page_id
{text_join}
WHERE pt.user_id = $1 AND lower(t.name) = $2
{text_where}
GROUP BY p.chapter_id
) c
"#,
)
.bind(user_id)
.bind(tag)
.fetch_one(pool)
.await?;
text_join = text_join,
text_where = text_where.replace("{n}", "$3"),
);
let mut cq = sqlx::query_as::<_, (i64,)>(&count_sql).bind(user_id).bind(tag);
if let Some(text) = text {
cq = cq.bind(text);
}
let (total,) = cq.fetch_one(pool).await?;
Ok((rows, total))
}
/// Paged list of mangas containing pages tagged `tag` for `user_id`,
/// ranked by `match_count` summed across all their chapters.
///
/// `text` behaves as in [`aggregate_chapters_for_tag`]: non-blank restricts to
/// pages whose OCR `search_doc` matches, and both `match_count` and the sample
/// thumbnails reflect that filtered set.
pub async fn aggregate_mangas_for_tag(
pool: &PgPool,
user_id: Uuid,
@@ -332,7 +379,26 @@ pub async fn aggregate_mangas_for_tag(
order: Order,
limit: i64,
offset: i64,
text: Option<&str>,
) -> AppResult<(Vec<TaggedMangaAggregate>, i64)> {
// OCR text filter — see `aggregate_chapters_for_tag` for the rationale.
let text = text.map(str::trim).filter(|s| !s.is_empty());
let (text_join, text_where) = match text {
Some(_) => (
"JOIN page_analysis pa ON pa.page_id = p.id",
"AND pa.search_doc @@ plainto_tsquery('simple', {n})",
),
None => ("", ""),
};
// Same filter inside the sample-thumbnail subquery (page aliased `p`),
// reusing the main query's `$5` text bind.
let (sample_join, sample_where) = match text {
Some(_) => (
"JOIN page_analysis pa2 ON pa2.page_id = p.id",
"AND pa2.search_doc @@ plainto_tsquery('simple', $5)",
),
None => ("", ""),
};
let sql = format!(
r#"
SELECT
@@ -349,9 +415,11 @@ pub async fn aggregate_mangas_for_tag(
JOIN chapters ch2 ON ch2.id = p.chapter_id
JOIN page_tags pt2 ON pt2.page_id = p.id
JOIN tags t2 ON t2.id = pt2.tag_id
{sample_join}
WHERE ch2.manga_id = m.id
AND pt2.user_id = $1
AND lower(t2.name) = $2
{sample_where}
ORDER BY p.page_number ASC
LIMIT 3
) p2
@@ -363,23 +431,31 @@ pub async fn aggregate_mangas_for_tag(
JOIN pages p ON p.id = pt.page_id
JOIN chapters ch ON ch.id = p.chapter_id
JOIN mangas m ON m.id = ch.manga_id
{text_join}
WHERE pt.user_id = $1
AND lower(t.name) = $2
{text_where}
GROUP BY m.id, m.title, m.cover_image_path
ORDER BY match_count {dir}, m.id
LIMIT $3 OFFSET $4
"#,
dir = order.as_sql(),
text_join = text_join,
text_where = text_where.replace("{n}", "$5"),
sample_join = sample_join,
sample_where = sample_where,
);
let rows = sqlx::query_as::<_, TaggedMangaAggregate>(&sql)
let mut q = sqlx::query_as::<_, TaggedMangaAggregate>(&sql)
.bind(user_id)
.bind(tag)
.bind(limit)
.bind(offset)
.fetch_all(pool)
.await?;
.bind(offset);
if let Some(text) = text {
q = q.bind(text);
}
let rows = q.fetch_all(pool).await?;
let (total,): (i64,) = sqlx::query_as(
let count_sql = format!(
r#"
SELECT count(*) FROM (
SELECT 1
@@ -387,15 +463,20 @@ pub async fn aggregate_mangas_for_tag(
JOIN tags t ON t.id = pt.tag_id
JOIN pages p ON p.id = pt.page_id
JOIN chapters ch ON ch.id = p.chapter_id
{text_join}
WHERE pt.user_id = $1 AND lower(t.name) = $2
{text_where}
GROUP BY ch.manga_id
) m
"#,
)
.bind(user_id)
.bind(tag)
.fetch_one(pool)
.await?;
text_join = text_join,
text_where = text_where.replace("{n}", "$3"),
);
let mut cq = sqlx::query_as::<_, (i64,)>(&count_sql).bind(user_id).bind(tag);
if let Some(text) = text {
cq = cq.bind(text);
}
let (total,) = cq.fetch_one(pool).await?;
Ok((rows, total))
}

View File

@@ -25,7 +25,7 @@ use serde::{Deserialize, Serialize};
use crate::analysis::prompt::{
GROUNDING_PROMPT_DEFAULT, OCR_PROMPT_DEFAULT, SYSTEM_PROMPT_DEFAULT,
};
use crate::config::{AnalysisConfig, CrawlerConfig, ResponseFormat};
use crate::config::{AnalysisBackend, AnalysisConfig, CrawlerConfig, ResponseFormat};
use crate::crawler::safety::DownloadAllowlist;
/// `app_settings.key` for the crawler group.
@@ -355,6 +355,15 @@ impl AnalysisSettings {
if self.workers < 1 {
errs.push("workers", "must be at least 1");
}
// The endpoint/model are vision-only knobs — the OCR backend never
// dials a URL or sends a model id. While vision is dormant
// (`base.backend == Ocr`), skip the live-worker SSRF gate and the
// model-required check so an OCR operator can enable the worker
// without a vision endpoint/model (the OCR settings UI doesn't even
// expose them). The basic malformed-URL sanity check below stays
// unconditional. Re-enabling vision restores the full gate via the
// `base.backend == Vision` predicate.
let vision_active = base.backend == AnalysisBackend::Vision;
let trimmed_endpoint = self.endpoint.trim();
if trimmed_endpoint.is_empty() {
errs.push("endpoint", "must be a valid absolute URL");
@@ -362,7 +371,7 @@ impl AnalysisSettings {
// Always reject obviously-malformed URLs (matches prior behaviour).
let _ = e;
errs.push("endpoint", "must be a valid absolute URL");
} else if self.enabled {
} else if self.enabled && vision_active {
// SSRF + API-key exfiltration defence — only enforced when the
// worker is actually live. Worker attaches an env-managed bearer
// token to every call; without this check, an admin (or one
@@ -379,7 +388,7 @@ impl AnalysisSettings {
errs.push("endpoint", url_safety_message(&e));
}
}
if self.enabled && self.model.trim().is_empty() {
if self.enabled && vision_active && self.model.trim().is_empty() {
errs.push("model", "required when analysis is enabled");
}
if self.max_tokens < 1 {
@@ -459,6 +468,11 @@ impl AnalysisSettings {
api_key: base.api_key.clone(),
// Env-only readiness probe URL preserved from the base.
vision_health_url: base.vision_health_url.clone(),
// Deploy-time engine selection + ocrs model paths: env-only, not
// admin-tunable, so carry them through from the base unchanged.
backend: base.backend,
ocr_detection_model: base.ocr_detection_model.clone(),
ocr_recognition_model: base.ocr_recognition_model.clone(),
})
}
}
@@ -696,7 +710,8 @@ mod tests {
// a hostile/CSRF-able admin must NOT be able to point endpoint at
// cloud metadata, loopback services, or RFC1918 hosts — when the
// worker is enabled (toggling enabled=true later re-runs this gate).
let base = AnalysisConfig::default();
let mut base = AnalysisConfig::default();
base.backend = AnalysisBackend::Vision;
for url in [
"http://169.254.169.254/v1/chat/completions",
"http://127.0.0.1:5432/",
@@ -720,7 +735,8 @@ mod tests {
// The documented default — docker DNS name resolving to a private IP
// at runtime — must still validate, because the bearer recipient
// identity is the operator-chosen hostname, not the underlying IP.
let base = AnalysisConfig::default();
let mut base = AnalysisConfig::default();
base.backend = AnalysisBackend::Vision;
for url in [
"http://mangalord-vision:8000/v1/chat/completions",
"https://api.openai.com/v1/chat/completions",
@@ -751,7 +767,10 @@ mod tests {
#[test]
fn analysis_requires_model_only_when_enabled() {
let base = AnalysisConfig::default();
// Vision base: the model id is required only when the vision worker
// is actually live (see the OCR carve-out below).
let mut base = AnalysisConfig::default();
base.backend = AnalysisBackend::Vision;
let disabled = AnalysisSettings {
enabled: false,
model: "".to_string(),
@@ -767,6 +786,28 @@ mod tests {
assert!(errs.errors.iter().any(|e| e.field == "model"));
}
#[test]
fn analysis_ocr_backend_enable_skips_vision_endpoint_and_model_validation() {
// With the OCR backend active the worker never dials the vision
// endpoint nor sends a model id, so enabling analysis must NOT
// validate those vision-only fields. The default base carries the
// dev-localhost endpoint (which the SSRF gate would otherwise reject)
// and an empty model — under OCR, enabling is still valid. Without
// this carve-out an OCR operator can't turn the worker on, since the
// OCR settings UI doesn't expose endpoint/model to fix.
let base = AnalysisConfig::default(); // backend = Ocr
let dto = AnalysisSettings {
enabled: true,
endpoint: "http://localhost:8000/v1/chat/completions".to_string(),
model: "".to_string(),
..AnalysisSettings::from_config(&base)
};
assert!(
dto.to_config(&base).is_ok(),
"OCR-enabled config must not fail on vision endpoint/model"
);
}
#[test]
fn dtos_serialize_to_json_and_back() {
let c = CrawlerSettings::default();