admin/mcp-nextcloud
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ingest): split OCR into tier2 in-cluster (GPU, gateway-only) + tier3 upstream

Browse Source 
Insert a configurable in-cluster OCR rung into the escalation ladder (Deck #353):
a tier2-eligible doc is OCR'd on the on-demand burst GPU before falling through to
paid upstream OCR. The in-cluster backend is reached ONLY via the embedding gateway
(model prefix routes to the GPU over the tailnet) and is a config value (default
surya/surya-ocr-2, swappable to e.g. lightonocr) — never hard-coded.

Ladder: fast -> structured -> ocr-incluster -> ocr-upstream
(queues ingest-ocr-incluster / ingest-ocr-upstream).

- escalation.py: 4-tier ladder; in-cluster flag folded into the dead-letter signature.
- ocr.py: OcrProcessor(name, tier, model_setting, gateway_only); build_ocr_backend(
  ..., model=, gateway_only=) — gateway_only forces the gateway backend (never the
  direct Mistral fallback), disabling the tier with a warning if no gateway URL.
- registry.py: per-rung enable map; scanned docs target minimum="ocr-incluster";
  inline path runs the cheapest available OCR rung.
- procrastinate.py: two OCR queues; legacy ingest-ocr kept as a drain target.
- config.py: DOCUMENT_OCR_INCLUSTER_ENABLED (off) + DOCUMENT_OCR_INCLUSTER_MODEL.
- __init__.py: register the two OCR instances; vector/processor.py: pages_ocr
  metered for the upstream (paid) rung only; cli.py: new --tier choices + legacy drain.
- metrics.py: zero the legacy ingest-ocr queue gauge during rollout.
- tests: migrated to the split ladder + new tests (gateway-only forcing, per-tier
  model incl. lightonocr override, no-hard-coded-surya guard). 1792 pass; ruff + ty green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2026-06-17 23:28:29 +02:00
co-authored by Claude Opus 4.8
parent 060084029f
commit c21804fbbc
16 changed files with 384 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nextcloud_mcp_server/cli.py
+14 -4
View File
@@ -335,7 +335,7 @@ def _init_worker_observability(settings: Settings) -> None:
)
@click.option(
"--tier",
type=click.Choice(["fast", "structured", "ocr"]),
type=click.Choice(["fast", "structured", "ocr-incluster", "ocr-upstream"]),
default=None,
help=(
"Run only this extraction tier's queue (Deck #323). Omit to drain ALL "
@@ -386,21 +386,31 @@ def worker(concurrency: int | None, tier: str | None):
ALL_INGEST_QUEUES,
INGEST_QUEUE_MAINTENANCE,
LEGACY_INGEST_QUEUE,
LEGACY_INGEST_QUEUE_OCR,
TIER_QUEUES,
apply_ingest_queue_schema,
get_procrastinate_app,
)
# Which queues this process drains. A single tier -> just its queue; no tier
# -> every tier queue PLUS the legacy single queue, so a rolling upgrade
# never strands jobs deferred under the pre-#323 name. Every worker also
# -> every tier queue PLUS the legacy queues, so a rolling upgrade never
# strands jobs deferred under the pre-#323 single queue or the pre-#353 single
# OCR queue. The upstream OCR worker also drains the pre-split ``ingest-ocr``
# queue (the old single OCR tier defaulted to upstream/Mistral). Every worker
# drains the maintenance queue so the periodic stalled-job reclaim fires
# regardless of which tier(s) are scaled up (procrastinate dedups the
# periodic, so multiple drainers don't multiply the reclaim).
if tier is not None:
queues = [TIER_QUEUES[tier], INGEST_QUEUE_MAINTENANCE]
if tier == "ocr-upstream":
queues.insert(1, LEGACY_INGEST_QUEUE_OCR)
else:
queues = [*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE, INGEST_QUEUE_MAINTENANCE]
queues = [
*ALL_INGEST_QUEUES,
LEGACY_INGEST_QUEUE,
LEGACY_INGEST_QUEUE_OCR,
INGEST_QUEUE_MAINTENANCE,
]
# This is the consumer side of the distributed (postgres) ingest backend.
# Unlike the in-process anyio pool, the worker talks to procrastinate's App
nextcloud_mcp_server/config.py
+14
View File
@@ -162,6 +162,13 @@ _DEFAULTS: dict[str, Any] = {
# Provider-namespaced OCR model id (gateway routes on the prefix; the direct
# mistral backend strips it).
"document_ocr_model": "mistral/mistral-ocr-latest",
# In-cluster OCR tier (tier2, Deck #353): the on-demand burst GPU, reached
# ONLY via the embedding gateway. Off + per-tenant by default; tried before the
# paid upstream OCR rung. The model is a config value (the gateway routes on its
# "<provider>/" prefix) — surya is the current pick but swappable (e.g.
# "lightonocr/...") and is NEVER hard-coded, only this default.
"document_ocr_incluster_enabled": False,
"document_ocr_incluster_model": "surya/surya-ocr-2",
# OCR escalation triggers (tier-0). A page is OCR-worthy when its text is
# near-empty (< min_page_chars) OR low-quality (< min_text_quality) OR (when
# detect_scanned) mostly a raster image; a doc escalates when the OCR-worthy
@@ -884,6 +891,11 @@ class Settings:
# gateway routes on the "<provider>/" prefix; the direct mistral backend
# strips it.
document_ocr_model: str = "mistral/mistral-ocr-latest"
# In-cluster OCR tier (tier2, Deck #353): on-demand burst GPU reached ONLY via
# the embedding gateway; off + per-tenant, tried before the upstream rung. The
# model is a swappable config value (surya is the current pick, never hardcoded).
document_ocr_incluster_enabled: bool = False
document_ocr_incluster_model: str = "surya/surya-ocr-2"
# OCR backend HTTP request timeout (seconds). float for parity with the
# parse timeout / httpx.Timeout; per-tenant tunable so a gateway with a
# shorter ceiling isn't masked by the 180s default.
@@ -1535,6 +1547,8 @@ def get_settings() -> Settings:
"document_ocr_enabled": "DOCUMENT_OCR_ENABLED",
"document_ocr_provider": "DOCUMENT_OCR_PROVIDER",
"document_ocr_model": "DOCUMENT_OCR_MODEL",
"document_ocr_incluster_enabled": "DOCUMENT_OCR_INCLUSTER_ENABLED",
"document_ocr_incluster_model": "DOCUMENT_OCR_INCLUSTER_MODEL",
"document_ocr_timeout_seconds": "DOCUMENT_OCR_TIMEOUT_SECONDS",
"document_ocr_mode": "DOCUMENT_OCR_MODE",
"document_ocr_batch_poll_seconds": "DOCUMENT_OCR_BATCH_POLL_SECONDS",
nextcloud_mcp_server/document_processors/__init__.py
+23 -4
View File
@@ -8,13 +8,32 @@ from .registry import ProcessorRegistry, get_registry
# Register processors at module initialization. The tiered PDF pipeline selects
# by tier (not priority): Pypdfium2FastProcessor is the ``fast`` tier,
# PyMuPDFProcessor the ``structured`` rollback, and OcrProcessor the ``ocr``
# escalation target (reached only when document_ocr_enabled). OcrProcessor gets
# the lowest priority so it is never the non-tiered default for PDFs.
# PyMuPDFProcessor the ``structured`` rollback, and TWO OcrProcessor instances are
# the OCR rungs — ``ocr-incluster`` (the on-demand burst GPU, gateway-only, reached
# via the embedding gateway over the tailnet; e.g. surya) tried before
# ``ocr-upstream`` (paid Mistral). Each is reached only when its own opt-in flag is
# set. OCR gets the lowest priorities so it's never the non-tiered default for PDFs.
_registry = get_registry()
_registry.register(Pypdfium2FastProcessor(), priority=20)
_registry.register(PyMuPDFProcessor(), priority=10)
_registry.register(OcrProcessor(), priority=1)
_registry.register(
OcrProcessor(
name="ocr-incluster",
tier="ocr-incluster",
model_setting="document_ocr_incluster_model",
gateway_only=True,
),
priority=2,
)
_registry.register(
OcrProcessor(
name="ocr-upstream",
tier="ocr-upstream",
model_setting="document_ocr_model",
gateway_only=False,
),
priority=1,
)
__all__ = [
"DocumentProcessor",
nextcloud_mcp_server/document_processors/escalation.py
+6 -3
View File
@@ -24,9 +24,11 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Literal
# Cheapest-first. ``llm`` is reserved (see base.DocumentProcessor.tier) and not
# wired yet, so it is intentionally absent from the live ladder.
TIER_LADDER: tuple[str, ...] = ("fast", "structured", "ocr")
# Cheapest-first. OCR is split into two rungs: ``ocr-incluster`` (the on-demand
# burst GPU, e.g. surya, reached via the embedding gateway over the tailnet) tried
# BEFORE ``ocr-upstream`` (paid Mistral). ``llm`` is reserved (see
# base.DocumentProcessor.tier) and not wired yet.
TIER_LADDER: tuple[str, ...] = ("fast", "structured", "ocr-incluster", "ocr-upstream")
def escalation_tiers_signature(settings: Any) -> str:
@@ -55,6 +57,7 @@ def escalation_tiers_signature(settings: Any) -> str:
"""
return (
f"ocr={int(bool(settings.document_ocr_enabled))};"
f"ocric={int(bool(settings.document_ocr_incluster_enabled))};"
f"t1={settings.document_tier1_engine}"
)
nextcloud_mcp_server/document_processors/ocr.py
+65 -12
View File
@@ -203,11 +203,16 @@ def _build_gateway_token_provider(settings: Settings) -> Any:
)
def build_gateway_batch_client(settings: Settings) -> "GatewayBatchOcrClient | None":
def build_gateway_batch_client(
settings: Settings, *, model: str | None = None
) -> "GatewayBatchOcrClient | None":
"""Build a ``GatewayBatchOcrClient`` when the gateway is the OCR backend, else
``None`` (so batch mode falls back to sync for provider=mistral / no gateway).
Batch OCR is gateway-only — Mistral's Batch API is reached *through* the
gateway's batch routes, never directly from the pod."""
gateway's batch routes, never directly from the pod.
``model`` overrides ``settings.document_ocr_model`` so a per-tier OCR rung
(e.g. the in-cluster tier) submits its own provider-namespaced model id."""
if settings.document_ocr_provider not in ("gateway", "auto"):
return None
if not settings.embedding_gateway_url:
@@ -216,28 +221,52 @@ def build_gateway_batch_client(settings: Settings) -> "GatewayBatchOcrClient | N
return GatewayBatchOcrClient(
settings.embedding_gateway_url,
settings.document_ocr_model,
model or settings.document_ocr_model,
_build_gateway_token_provider(settings),
)
def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
"""Select an OCR backend from settings, or None when none is available."""
def build_ocr_backend(
settings: Settings, *, model: str | None = None, gateway_only: bool = False
) -> _OcrBackend | None:
"""Select an OCR backend from settings, or None when none is available.
``model`` overrides ``settings.document_ocr_model`` so a per-tier OCR rung
binds its own provider-namespaced model id. ``gateway_only`` forces the
gateway backend (never the direct Mistral fallback) — used by the **in-cluster**
OCR tier, whose backend (e.g. surya on the burst GPU) is reachable ONLY through
the embedding gateway over the tailnet; with no gateway URL the tier is
disabled (warn) rather than misrouted to a direct backend that can't serve it.
"""
provider = settings.document_ocr_provider
model = model or settings.document_ocr_model
if provider == "none":
return None
if gateway_only:
if settings.embedding_gateway_url:
return _GatewayOcrBackend(
settings.embedding_gateway_url,
model,
_build_gateway_token_provider(settings),
)
logger.warning(
"in-cluster OCR tier requires EMBEDDING_GATEWAY_URL (it routes through "
"the gateway, never a direct backend); this OCR tier is disabled"
)
return None
if provider in ("gateway", "auto") and settings.embedding_gateway_url:
return _GatewayOcrBackend(
settings.embedding_gateway_url,
settings.document_ocr_model,
model,
_build_gateway_token_provider(settings),
)
if provider in ("mistral", "auto") and settings.mistral_api_key:
return _MistralOcrBackend(
settings.mistral_api_key,
settings.document_ocr_model,
model,
settings.mistral_base_url,
)
@@ -260,7 +289,24 @@ def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
class OcrProcessor(DocumentProcessor):
"""Tier-3 OCR processor (gateway or direct Mistral backend)."""
def __init__(self) -> None:
def __init__(
self,
*,
name: str = "ocr-upstream",
tier: str = "ocr-upstream",
model_setting: str = "document_ocr_model",
gateway_only: bool = False,
) -> None:
# One OcrProcessor class serves BOTH OCR rungs; instances are bound to a
# tier + the settings attribute holding their provider-namespaced model id
# (+ whether the backend is gateway-only). The in-cluster rung is
# gateway-only (its model, e.g. surya, is reachable solely via the
# gateway); the upstream rung keeps the configurable gateway/mistral
# selection. surya is NEVER hard-coded here — only a config default.
self._name = name
self._tier = tier
self._model_setting = model_setting
self._gateway_only = gateway_only
# Resolve the backend once and reuse it: rebuilding per call would create
# a fresh GatewayTokenProvider each time (discarding its M2M-token cache
# -> a token fetch per document) and a new Mistral SDK client per call.
@@ -283,11 +329,11 @@ class OcrProcessor(DocumentProcessor):
@property
def name(self) -> str:
return "ocr"
return self._name
@property
def tier(self) -> str:
return "ocr"
return self._tier
@property
def supported_mime_types(self) -> set[str]:
@@ -327,7 +373,11 @@ class OcrProcessor(DocumentProcessor):
self._backend_lock = anyio.Lock()
async with self._backend_lock:
if not self._backend_resolved: # double-checked
self._backend = build_ocr_backend(settings)
self._backend = build_ocr_backend(
settings,
model=getattr(settings, self._model_setting),
gateway_only=self._gateway_only,
)
self._backend_resolved = True
backend = self._backend
if backend is None:
@@ -393,7 +443,10 @@ class OcrProcessor(DocumentProcessor):
self._batch_client_lock = anyio.Lock()
async with self._batch_client_lock:
if not self._batch_client_resolved: # double-checked
self._batch_client = build_gateway_batch_client(get_settings())
self._batch_client = build_gateway_batch_client(
get_settings(),
model=getattr(get_settings(), self._model_setting),
)
self._batch_client_resolved = True
return self._batch_client
nextcloud_mcp_server/document_processors/registry.py
+30 -14
View File
@@ -349,10 +349,18 @@ class ProcessorRegistry:
and not structured_failed
and classification.recommended_tier in ("ocr", "structured")
and classification.page_count > 0
and settings.document_ocr_enabled
and (
settings.document_ocr_enabled or settings.document_ocr_incluster_enabled
)
):
ocr = self._pdf_processor_for_tier("ocr")
if ocr is not None:
# Inline (memory pool) path: no queues to hop, so pick the cheapest
# available OCR rung (in-cluster GPU before paid upstream) via the same
# availability walk the queue path uses.
ocr_tier = self.next_available_tier(
from_tier, settings, minimum="ocr-incluster"
)
ocr = self._pdf_processor_for_tier(ocr_tier) if ocr_tier else None
if ocr is not None and ocr_tier is not None:
reason = (
"corrupt_glyphs"
if classification.recommended_tier == "structured"
@@ -360,11 +368,12 @@ class ProcessorRegistry:
if classification.total_chars == 0
else "low_confidence"
)
record_document_escalation(from_tier, "ocr", reason)
record_document_escalation(from_tier, ocr_tier, reason)
logger.info(
"Escalating %s %s->ocr (reason=%s)",
"Escalating %s %s->%s (reason=%s)",
filename or "<bytes>",
from_tier,
ocr_tier,
reason,
)
ocr_result = await self._run_processor(
@@ -379,13 +388,14 @@ class ProcessorRegistry:
# OCR is an enhancement, not a gate: if it can't run (no backend
# configured / API down) or returns nothing, keep the tier-1
# result rather than failing the document. Otherwise an operator
# who sets DOCUMENT_OCR_ENABLED=true without credentials would
# make scanned docs fail entirely -- strictly worse than off.
# who enables OCR without credentials would make scanned docs fail
# entirely -- strictly worse than off.
if ocr_result.success:
return ocr_result
logger.warning(
"OCR escalation did not succeed for %s (%s); keeping the "
"OCR escalation to %s did not succeed for %s (%s); keeping the "
"tier-1 result",
ocr_tier,
filename or "<bytes>",
ocr_result.metadata.get("parse_failed_reason", "error"),
)
@@ -503,11 +513,14 @@ class ProcessorRegistry:
"""
if self._pdf_processor_for_tier(tier) is None:
return False
if (
not ignore_ocr_enabled
and tier == "ocr"
and not settings.document_ocr_enabled
):
# Each OCR rung has its own opt-in flag (in-cluster vs upstream); a rung is
# unavailable when its flag is off (unless we're computing the what-if
# ideal target). Non-OCR tiers have no enabled gate.
ocr_enable = {
"ocr-incluster": settings.document_ocr_incluster_enabled,
"ocr-upstream": settings.document_ocr_enabled,
}
if not ignore_ocr_enabled and tier in ocr_enable and not ocr_enable[tier]:
return False
return True
@@ -645,7 +658,10 @@ class ProcessorRegistry:
minimum = "structured"
reason = "corrupt_glyphs"
elif classification.total_chars == 0:
minimum = "ocr"
# Scanned / no text layer: target the cheapest OCR rung (in-cluster
# GPU); next_available_tier then falls through to the upstream rung if
# in-cluster is disabled/unregistered.
minimum = "ocr-incluster"
reason = "empty_text"
else:
minimum = None
nextcloud_mcp_server/observability/metrics.py
+2 -1
View File
@@ -708,9 +708,10 @@ def update_ingest_queue_depth(by_queue: dict[str, dict[str, int]] | None) -> Non
from nextcloud_mcp_server.vector.queue.procrastinate import ( # noqa: PLC0415
ALL_INGEST_QUEUES,
LEGACY_INGEST_QUEUE,
LEGACY_INGEST_QUEUE_OCR,
)
for queue in (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE):
for queue in (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE, LEGACY_INGEST_QUEUE_OCR):
for status in _INGEST_DEPTH_STATUSES:
ingest_queue_depth.labels(queue=queue, status=status).set(0)
for queue, per_status in by_queue.items():
nextcloud_mcp_server/vector/processor.py
+7 -4
View File
@@ -353,10 +353,13 @@ async def record_indexing_usage(
)
# Paid-OCR pages are metered as a SEPARATE line (Deck #323) so the
# expensive tier's cost is billable independently of CPU-cheap parsing
# -- pages_embedded counts all parsed pages, pages_ocr only the OCR
# tier's. Gated on the tier so it's emitted exactly when the doc was
# actually OCR'd; the same page_count guard above applies.
if pipeline_tier == "ocr":
# -- pages_embedded counts all parsed pages, pages_ocr only the paid
# OCR tier's. After the OCR split (Deck #353) "paid" = the UPSTREAM
# (Mistral) rung; the in-cluster GPU rung's cost is recovered via the
# burst lifecycle, not per-page, so it is NOT metered here (separate
# per-tier OCR metering is a billing follow-up). Gated on the tier so
# it's emitted exactly when the doc hit upstream OCR.
if pipeline_tier == "ocr-upstream":
await store.record_usage_event(
metric="pages_ocr",
value=page_count,
nextcloud_mcp_server/vector/queue/procrastinate.py
+18 -8
View File
@@ -66,24 +66,30 @@ logger = logging.getLogger(__name__)
# network-bound ``ocr`` fleet scale (and fail) independently.
INGEST_QUEUE_FAST = "ingest-fast"
INGEST_QUEUE_STRUCTURED = "ingest-structured"
INGEST_QUEUE_OCR = "ingest-ocr"
# OCR split into two rungs (Deck #353): in-cluster (burst GPU via the gateway,
# the queue the GPU autoscaler counts) tried before upstream (paid Mistral).
INGEST_QUEUE_OCR_INCLUSTER = "ingest-ocr-incluster"
INGEST_QUEUE_OCR_UPSTREAM = "ingest-ocr-upstream"
# tier -> queue. The producer always defers onto the cheapest tier's queue; a
# low-quality parse hops the job up the ladder via the retry strategy below.
TIER_QUEUES: dict[str, str] = {
"fast": INGEST_QUEUE_FAST,
"structured": INGEST_QUEUE_STRUCTURED,
"ocr": INGEST_QUEUE_OCR,
"ocr-incluster": INGEST_QUEUE_OCR_INCLUSTER,
"ocr-upstream": INGEST_QUEUE_OCR_UPSTREAM,
}
_QUEUE_TIERS: dict[str, str] = {queue: tier for tier, queue in TIER_QUEUES.items()}
ALL_INGEST_QUEUES: tuple[str, ...] = tuple(TIER_QUEUES.values())
# New jobs start here; ``ocr`` is reached only by escalation.
# New jobs start here; the OCR rungs are reached only by escalation.
DEFAULT_INGEST_QUEUE = INGEST_QUEUE_FAST
# Legacy single-queue name (pre-#323). A rolling upgrade may still have jobs
# parked on it; a worker can be told to drain it alongside the tier queues, and
# the job-count / reclaim helpers include it so nothing is stranded.
# Legacy single-queue name (pre-#323) and the pre-split single OCR queue
# (pre-#353). A rolling upgrade may still have jobs parked on either; a worker
# can drain them alongside the tier queues, and the job-count / reclaim helpers
# include them so nothing is stranded.
LEGACY_INGEST_QUEUE = "ingest"
LEGACY_INGEST_QUEUE_OCR = "ingest-ocr"
# Back-compat alias for callers that imported the old single-queue constant.
INGEST_QUEUE_NAME = DEFAULT_INGEST_QUEUE
@@ -95,8 +101,12 @@ INGEST_QUEUE_NAME = DEFAULT_INGEST_QUEUE
# queues so tier isolation (which fleet processes which docs) is preserved.
INGEST_QUEUE_MAINTENANCE = "ingest-maintenance"
# Queues the job-count + reclaim helpers sweep (tier queues + the legacy one).
_MANAGED_QUEUES: tuple[str, ...] = (*ALL_INGEST_QUEUES, LEGACY_INGEST_QUEUE)
# Queues the job-count + reclaim helpers sweep (tier queues + the legacy ones).
_MANAGED_QUEUES: tuple[str, ...] = (
*ALL_INGEST_QUEUES,
LEGACY_INGEST_QUEUE,
LEGACY_INGEST_QUEUE_OCR,
)
# Blueprint namespace → registered task names are prefixed ``ingest:``.
_NAMESPACE = "ingest"
tests/unit/test_escalation_signature.py
+8 -2
View File
@@ -19,8 +19,14 @@ from nextcloud_mcp_server.document_processors.escalation import (
pytestmark = pytest.mark.unit
def _settings(*, ocr: bool, engine: str = "pypdfium2") -> SimpleNamespace:
return SimpleNamespace(document_ocr_enabled=ocr, document_tier1_engine=engine)
def _settings(
*, ocr: bool, ocr_incluster: bool = False, engine: str = "pypdfium2"
) -> SimpleNamespace:
return SimpleNamespace(
document_ocr_enabled=ocr,
document_ocr_incluster_enabled=ocr_incluster,
document_tier1_engine=engine,
)
def test_signature_is_stable_for_same_config() -> None:
tests/unit/test_ocr_processor.py
+107 -11
View File
@@ -99,6 +99,100 @@ def test_build_gateway_batch_client_gateway_only(kw, expect_client):
assert (client is not None) is expect_client
# --- in-cluster (tier2) backend: gateway-forced + configurable model -----------
def test_build_backend_gateway_only_forces_gateway():
"""The in-cluster tier is gateway-only: even with provider=mistral + a key it
builds the gateway backend, NEVER the direct Mistral fallback (the GPU is
reachable solely through the gateway)."""
b = ocr.build_ocr_backend(
_settings(
document_ocr_provider="mistral",
mistral_api_key="k",
embedding_gateway_url="http://gw",
),
gateway_only=True,
)
assert isinstance(b, ocr._GatewayOcrBackend)
def test_build_backend_gateway_only_no_url_disabled():
"""Gateway-only with no gateway URL -> disabled (None), never a direct backend."""
assert (
ocr.build_ocr_backend(_settings(mistral_api_key="k"), gateway_only=True) is None
)
def test_build_backend_model_override_is_not_hardcoded():
"""The per-tier model is whatever config passes -- surya by default, but fully
swappable (e.g. lightonocr) with no code change."""
surya = ocr.build_ocr_backend(
_settings(embedding_gateway_url="http://gw"),
model="surya/surya-ocr-2",
gateway_only=True,
)
assert isinstance(surya, ocr._GatewayOcrBackend)
assert surya._model == "surya/surya-ocr-2"
lit = ocr.build_ocr_backend(
_settings(embedding_gateway_url="http://gw"),
model="lightonocr/lightonocr-1b",
gateway_only=True,
)
assert isinstance(lit, ocr._GatewayOcrBackend)
assert lit._model == "lightonocr/lightonocr-1b" # config-driven, not hardcoded
async def test_incluster_processor_resolves_its_model_gateway_only(monkeypatch):
"""An OcrProcessor bound to the in-cluster rung builds its backend with its own
configured model (document_ocr_incluster_model) and gateway_only=True."""
captured: dict[str, Any] = {}
class _FakeBackend:
async def ocr(self, content, mime_type):
return "ocr text", [{"page": 1, "start_offset": 0, "end_offset": 8}]
def _spy(settings, *, model=None, gateway_only=False):
captured["model"] = model
captured["gateway_only"] = gateway_only
return _FakeBackend()
monkeypatch.setattr(ocr, "build_ocr_backend", _spy)
monkeypatch.setattr(
ocr,
"get_settings",
lambda: _settings(
document_ocr_incluster_model="surya/surya-ocr-2",
embedding_gateway_url="http://gw",
),
)
proc = ocr.OcrProcessor(
name="ocr-incluster",
tier="ocr-incluster",
model_setting="document_ocr_incluster_model",
gateway_only=True,
)
await proc.process(b"%PDF", "application/pdf", "x.pdf")
assert captured == {"model": "surya/surya-ocr-2", "gateway_only": True}
def test_no_surya_string_literal_in_document_processors():
"""surya must be a CONFIG default only -- never a hard-coded behavioural literal
in the worker (it's swappable, e.g. lightonocr). Comments may mention it; code
string literals may not (the default lives in config.py, a different module)."""
import pathlib # noqa: PLC0415
assert ocr.__file__ is not None
pkg = pathlib.Path(ocr.__file__).parent
offenders = [
f"{p.name}: {ln.strip()}"
for p in pkg.glob("*.py")
for ln in p.read_text().splitlines()
if '"surya' in ln.split("#", 1)[0] or "'surya" in ln.split("#", 1)[0]
]
assert not offenders, offenders
def test_build_backend_gateway_missing_m2m_raises():
# client_id set but token_url/secret missing -> explicit ValueError (not a
# stripped assert), surfaced on backend resolution.
@@ -126,7 +220,7 @@ async def test_processor_unsupported_when_no_backend(monkeypatch):
monkeypatch.setattr(
ocr, "get_settings", lambda: _settings(document_ocr_provider="none")
)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: None)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: None)
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is False
assert r.metadata["parse_failed_reason"] == "unsupported"
@@ -138,12 +232,12 @@ async def test_processor_success(monkeypatch):
return "hello world", [{"page": 1, "start_offset": 0, "end_offset": 11}]
monkeypatch.setattr(ocr, "get_settings", lambda: _settings())
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _FakeBackend())
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _FakeBackend())
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is True
assert r.text == "hello world"
assert r.metadata["page_count"] == 1
assert r.processor == "ocr"
assert r.processor == "ocr-upstream"
async def test_processor_backend_error_returns_success_false(monkeypatch):
@@ -152,7 +246,7 @@ async def test_processor_backend_error_returns_success_false(monkeypatch):
raise RuntimeError("api down")
monkeypatch.setattr(ocr, "get_settings", lambda: _settings())
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _BoomBackend())
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _BoomBackend())
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is False
assert r.metadata["parse_failed_reason"] == "error"
@@ -168,7 +262,7 @@ async def test_processor_timeout_returns_timeout_reason(monkeypatch):
monkeypatch.setattr(
ocr, "get_settings", lambda: _settings(document_ocr_timeout_seconds=5.0)
)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _TimeoutBackend())
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _TimeoutBackend())
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is False
assert r.metadata["parse_failed_reason"] == "timeout"
@@ -187,7 +281,9 @@ async def test_gateway_httpx_timeout_maps_to_timeout_reason(monkeypatch):
monkeypatch.setattr(
ocr, "get_settings", lambda: _settings(document_ocr_timeout_seconds=5.0)
)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _HttpxTimeoutBackend())
monkeypatch.setattr(
ocr, "build_ocr_backend", lambda s, **kw: _HttpxTimeoutBackend()
)
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is False
assert r.metadata["parse_failed_reason"] == "timeout"
@@ -331,7 +427,7 @@ def _wire_batch(monkeypatch, *, client, store, settings=None):
embedding_gateway_url="https://gw",
)
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: client)
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s, **kw: client)
async def _shared(cls):
return store
@@ -451,8 +547,8 @@ async def test_batch_falls_back_to_sync_when_no_gateway(monkeypatch):
settings = _settings(document_ocr_mode="batch", document_ocr_provider="mistral")
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: None)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _FakeBackend())
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s, **kw: None)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _FakeBackend())
r = await ocr.OcrProcessor().process(
b"%PDF", "application/pdf", options=dict(_IDENTITY)
@@ -472,8 +568,8 @@ async def test_batch_falls_back_to_sync_when_no_identity(monkeypatch):
embedding_gateway_url="https://gw",
)
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: client)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s: _FakeBackend())
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s, **kw: client)
monkeypatch.setattr(ocr, "build_ocr_backend", lambda s, **kw: _FakeBackend())
# No options -> inline path -> batch inapplicable -> sync fallback.
r = await ocr.OcrProcessor().process(b"%PDF", "application/pdf", options=None)
tests/unit/test_processor_metering.py
+2 -2
View File
@@ -191,7 +191,7 @@ async def test_ocr_tier_records_pages_ocr(store_spy):
token_count=900,
total_chars=40000,
page_count=8,
pipeline_tier="ocr",
pipeline_tier="ocr-upstream",
)
by_metric = {
c.kwargs["metric"]: c.kwargs["value"]
@@ -205,7 +205,7 @@ async def test_ocr_tier_records_pages_ocr(store_spy):
}
# pipeline_tier is threaded into the billing metadata for CP attribution.
for c in store_spy.record_usage_event.await_args_list:
assert c.kwargs["metadata"]["pipeline_tier"] == "ocr"
assert c.kwargs["metadata"]["pipeline_tier"] == "ocr-upstream"
@pytest.mark.unit
tests/unit/test_registry_tiering.py
+70 -45
View File
@@ -74,6 +74,7 @@ class _Settings:
engine="pypdfium2",
classify=True,
ocr=False,
ocr_incluster=False,
min_text_quality=0.5,
page_fraction=0.5,
min_page_chars=16,
@@ -85,7 +86,13 @@ class _Settings:
):
self.document_tier1_engine = engine
self.document_classify_enabled = classify
# ``ocr`` enables the UPSTREAM rung (document_ocr_enabled); ``ocr_incluster``
# the in-cluster GPU rung (tried first). The model attrs let
# build_ocr_backend resolve a per-tier model id.
self.document_ocr_enabled = ocr
self.document_ocr_incluster_enabled = ocr_incluster
self.document_ocr_model = "mistral/mistral-ocr-latest"
self.document_ocr_incluster_model = "surya/surya-ocr-2"
self.document_ocr_min_text_quality = min_text_quality
self.document_ocr_page_fraction = page_fraction
self.document_ocr_min_page_chars = min_page_chars
@@ -195,10 +202,10 @@ async def test_ocr_escalation_on_empty_text(monkeypatch):
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
r = _registry(
(_Fake("fast", "fast", text=""), 20),
(_Fake("ocr", "ocr", text="ocr text"), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="ocr text"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "ocr"
assert res.processor == "ocr-upstream"
esc.assert_called_once()
@@ -210,7 +217,7 @@ async def test_zero_page_pdf_does_not_escalate(monkeypatch):
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
r = _registry(
(_Fake("fast", "fast", text="", pages=0), 20),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "fast"
@@ -231,7 +238,7 @@ async def test_ocr_failure_falls_back_to_fast(monkeypatch):
monkeypatch.setattr(reg_mod, "record_document_escalation", MagicMock())
r = _registry(
(_Fake("fast", "fast", text=""), 20),
(_Fake("ocr", "ocr", text="", success=False), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="", success=False), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "fast"
@@ -242,7 +249,7 @@ async def test_no_ocr_escalation_when_disabled(monkeypatch):
monkeypatch.setattr(reg_mod, "get_settings", lambda: _Settings(ocr=False))
r = _registry(
(_Fake("fast", "fast", text=""), 20),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
# Fast tier is terminal when OCR is disabled.
@@ -289,11 +296,11 @@ async def test_glyph_corrupt_no_structured_falls_through_to_ocr(monkeypatch):
monkeypatch.setattr(reg_mod, "record_document_escalation", esc)
r = _registry(
(_Fake("fast", "fast", text=_GLYPH), 20),
(_Fake("ocr", "ocr", text="ocr recovered text"), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="ocr recovered text"), 5),
) # no structured registered
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "ocr"
esc.assert_called_once_with("fast", "ocr", "corrupt_glyphs")
assert res.processor == "ocr-upstream"
esc.assert_called_once_with("fast", "ocr-upstream", "corrupt_glyphs")
async def test_inline_lowconf_tries_structured_before_ocr(monkeypatch):
@@ -305,7 +312,7 @@ async def test_inline_lowconf_tries_structured_before_ocr(monkeypatch):
r = _registry(
(_Fake("fast", "fast", text="x" * 40), 20), # one long token -> quality ~0
(_Fake("structured", "structured", text="clean recovered prose text here"), 10),
(_Fake("ocr", "ocr", text="ocr text"), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="ocr text"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "structured"
@@ -321,11 +328,11 @@ async def test_inline_empty_skips_structured_straight_to_ocr(monkeypatch):
r = _registry(
(_Fake("fast", "fast", text=""), 20),
(_Fake("structured", "structured", text="should not run"), 10),
(_Fake("ocr", "ocr", text="ocr text"), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="ocr text"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "ocr"
esc.assert_called_once_with("fast", "ocr", "empty_text")
assert res.processor == "ocr-upstream"
esc.assert_called_once_with("fast", "ocr-upstream", "empty_text")
async def test_inline_fast_structured_ocr_cascade(monkeypatch):
@@ -339,13 +346,13 @@ async def test_inline_fast_structured_ocr_cascade(monkeypatch):
r = _registry(
(_Fake("fast", "fast", text="x" * 40), 20), # quality ~0, non-empty
(_Fake("structured", "structured", text=""), 10), # re-extract empty
(_Fake("ocr", "ocr", text="ocr recovered text"), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="ocr recovered text"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "ocr"
assert res.processor == "ocr-upstream"
assert esc.call_args_list == [
call("fast", "structured", "low_confidence"),
call("structured", "ocr", "empty_text"),
call("structured", "ocr-upstream", "empty_text"),
]
@@ -359,13 +366,13 @@ async def test_inline_structured_still_corrupt_escalates_to_ocr(monkeypatch):
r = _registry(
(_Fake("fast", "fast", text=_GLYPH), 20),
(_Fake("structured", "structured", text=_GLYPH), 10), # still corrupt
(_Fake("ocr", "ocr", text="ocr recovered text"), 5),
(_Fake("ocr-upstream", "ocr-upstream", text="ocr recovered text"), 5),
)
res = await r.process(b"%PDF-1.7", "application/pdf")
assert res.processor == "ocr"
assert res.processor == "ocr-upstream"
assert esc.call_args_list == [
call("fast", "structured", "corrupt_glyphs"),
call("structured", "ocr", "corrupt_glyphs"),
call("structured", "ocr-upstream", "corrupt_glyphs"),
]
@@ -375,7 +382,7 @@ def test_evaluate_escalation_glyph_corrupt_goes_structured(monkeypatch):
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = ProcessingResult(
text=_GLYPH,
@@ -399,7 +406,7 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_falls_through_to_ocr(
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
) # no structured registered
res = ProcessingResult(
text=_GLYPH,
@@ -412,7 +419,7 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_falls_through_to_ocr(
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=True))
assert decision == EscalationDecision("hop", "ocr", "corrupt_glyphs")
assert decision == EscalationDecision("hop", "ocr-upstream", "corrupt_glyphs")
def test_evaluate_escalation_glyph_corrupt_no_structured_ocr_disabled_suppressed(
@@ -424,7 +431,7 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_ocr_disabled_suppressed
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
) # structured not registered; ocr registered but disabled below
res = ProcessingResult(
text=_GLYPH,
@@ -437,7 +444,9 @@ def test_evaluate_escalation_glyph_corrupt_no_structured_ocr_disabled_suppressed
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
assert decision == EscalationDecision("suppressed", "ocr", "corrupt_glyphs")
assert decision == EscalationDecision(
"suppressed", "ocr-upstream", "corrupt_glyphs"
)
# --- Per-tier external path (Deck #323) -------------------------------------
@@ -449,7 +458,7 @@ async def test_process_tier_runs_named_tier(monkeypatch):
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = await r.process_tier(b"%PDF-1.7", "application/pdf", "f.pdf", "structured")
assert res.processor == "structured"
@@ -469,7 +478,7 @@ async def test_process_tier_oversize_fails_fast(monkeypatch):
monkeypatch.setattr(
reg_mod, "get_settings", lambda: _Settings(max_pdf_size_mb=0.001)
)
r = _registry((_Fake("ocr", "ocr"), 5))
r = _registry((_Fake("ocr-upstream", "ocr-upstream"), 5))
res = await r.process_tier(b"x" * 4096, "application/pdf", "big.pdf", "ocr")
assert res.success is False
assert res.metadata["parse_failed_reason"] == "oversize"
@@ -479,7 +488,7 @@ def test_next_available_tier_walks_ladder():
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
# ocr disabled -> structured is the only target above fast.
s = _Settings(ocr=False)
@@ -487,21 +496,25 @@ def test_next_available_tier_walks_ladder():
assert r.next_available_tier("structured", s) is None # ocr gated off
# ocr enabled -> reachable; minimum skips the structured rung.
s_ocr = _Settings(ocr=True)
assert r.next_available_tier("structured", s_ocr) == "ocr"
assert r.next_available_tier("fast", s_ocr, minimum="ocr") == "ocr"
assert r.next_available_tier("structured", s_ocr) == "ocr-upstream"
assert (
r.next_available_tier("fast", s_ocr, minimum="ocr-upstream") == "ocr-upstream"
)
def test_next_available_tier_skips_unregistered():
# No structured processor -> fast escalates straight to ocr.
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
assert r.next_available_tier("fast", _Settings(ocr=True)) == "ocr"
r = _registry(
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
)
assert r.next_available_tier("fast", _Settings(ocr=True)) == "ocr-upstream"
def test_evaluate_escalation_good_text_indexes(monkeypatch):
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
r = _registry(
(_Fake("fast", "fast", text="This is clean readable prose text."), 20),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = ProcessingResult(
text="This is clean readable prose text.",
@@ -520,7 +533,7 @@ def test_evaluate_escalation_empty_jumps_to_ocr(monkeypatch):
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = ProcessingResult(
text="",
@@ -531,7 +544,7 @@ def test_evaluate_escalation_empty_jumps_to_ocr(monkeypatch):
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=True))
assert decision == EscalationDecision("hop", "ocr", "empty_text")
assert decision == EscalationDecision("hop", "ocr-upstream", "empty_text")
def test_evaluate_escalation_lowconf_goes_to_structured(monkeypatch):
@@ -541,7 +554,7 @@ def test_evaluate_escalation_lowconf_goes_to_structured(monkeypatch):
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = ProcessingResult(
text=junk,
@@ -559,7 +572,9 @@ def test_evaluate_escalation_lowconf_goes_to_structured(monkeypatch):
def test_evaluate_escalation_failure_not_escalated(monkeypatch):
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
r = _registry(
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
)
res = ProcessingResult(
text="",
metadata={"parse_failed_reason": "error"},
@@ -590,7 +605,9 @@ def test_evaluate_escalation_terminal_when_no_higher_tier(monkeypatch):
def test_evaluate_escalation_zero_page_does_not_escalate(monkeypatch):
"""A zero-page (empty/corrupt) PDF never escalates on the external path."""
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
r = _registry(
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
)
res = ProcessingResult(
text="",
metadata={"page_count": 0, "page_boundaries": []},
@@ -604,7 +621,9 @@ def test_evaluate_escalation_lowconf_to_ocr_when_no_structured(monkeypatch):
unregistered structured rung), not to None."""
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
junk = "z" * 40 # non-empty but junk -> recommended ocr, total_chars > 0
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
r = _registry(
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
)
res = ProcessingResult(
text=junk,
metadata={
@@ -616,14 +635,16 @@ def test_evaluate_escalation_lowconf_to_ocr_when_no_structured(monkeypatch):
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=True))
assert decision == EscalationDecision("hop", "ocr", "low_confidence")
assert decision == EscalationDecision("hop", "ocr-upstream", "low_confidence")
def test_evaluate_escalation_suppressed_when_ocr_disabled(monkeypatch):
"""OCR off: a scanned doc does NOT hop to ocr; it returns a 'suppressed'
decision (the what-if-OCR signal) so the caller indexes at the current tier."""
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
r = _registry(
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
)
res = ProcessingResult(
text="",
metadata={
@@ -633,7 +654,7 @@ def test_evaluate_escalation_suppressed_when_ocr_disabled(monkeypatch):
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
assert decision == EscalationDecision("suppressed", "ocr", "empty_text")
assert decision == EscalationDecision("suppressed", "ocr-upstream", "empty_text")
def test_evaluate_escalation_lowconf_suppressed_when_only_ocr_disabled(monkeypatch):
@@ -641,7 +662,9 @@ def test_evaluate_escalation_lowconf_suppressed_when_only_ocr_disabled(monkeypat
the would-be hop is suppressed (not a structured hop, which isn't registered)."""
monkeypatch.setattr(reg_mod, "record_document_classification", MagicMock())
junk = "q" * 40
r = _registry((_Fake("fast", "fast"), 20), (_Fake("ocr", "ocr"), 5))
r = _registry(
(_Fake("fast", "fast"), 20), (_Fake("ocr-upstream", "ocr-upstream"), 5)
)
res = ProcessingResult(
text=junk,
metadata={
@@ -653,7 +676,9 @@ def test_evaluate_escalation_lowconf_suppressed_when_only_ocr_disabled(monkeypat
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
assert decision == EscalationDecision("suppressed", "ocr", "low_confidence")
assert decision == EscalationDecision(
"suppressed", "ocr-upstream", "low_confidence"
)
def test_evaluate_escalation_structured_hop_not_suppressed_when_ocr_off(monkeypatch):
@@ -664,7 +689,7 @@ def test_evaluate_escalation_structured_hop_not_suppressed_when_ocr_off(monkeypa
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10),
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = ProcessingResult(
text=junk,
@@ -706,7 +731,7 @@ def test_evaluate_escalation_empty_suppressed_even_when_structured_registered(
r = _registry(
(_Fake("fast", "fast"), 20),
(_Fake("structured", "structured"), 10), # registered but skipped for empty
(_Fake("ocr", "ocr"), 5),
(_Fake("ocr-upstream", "ocr-upstream"), 5),
)
res = ProcessingResult(
text="",
@@ -717,4 +742,4 @@ def test_evaluate_escalation_empty_suppressed_even_when_structured_registered(
processor="fast",
)
decision = r.evaluate_escalation(res, b"%PDF", "fast", _Settings(ocr=False))
assert decision == EscalationDecision("suppressed", "ocr", "empty_text")
assert decision == EscalationDecision("suppressed", "ocr-upstream", "empty_text")
tests/unit/vector/test_processor_dead_letter.py
+1
View File
@@ -29,6 +29,7 @@ pytestmark = pytest.mark.unit
def _settings(*, ocr_enabled: bool) -> SimpleNamespace:
return SimpleNamespace(
document_ocr_enabled=ocr_enabled,
document_ocr_incluster_enabled=False,
document_tier1_engine="pypdfium2",
get_collection_name=lambda: "c",
)
tests/unit/vector/test_procrastinate_producer.py
+2 -2
View File
@@ -115,7 +115,7 @@ class TestProcessDocumentTask:
# Calling the Task runs its wrapped function in-process. The job is on the
# ocr queue, so the queue-aware task must derive tier="ocr".
await pq.process_document_task(
_ctx(pq.INGEST_QUEUE_OCR),
_ctx(pq.INGEST_QUEUE_OCR_UPSTREAM),
user_id="alice",
doc_id="42",
doc_type="note",
@@ -131,7 +131,7 @@ class TestProcessDocumentTask:
# Worker disables the in-process retry loop; durable retry is the queue's.
assert captured["max_retries"] == 1
# Tier is derived from the job's queue (escalation enabled by default).
assert captured["tier"] == "ocr"
assert captured["tier"] == "ocr-upstream"
fake_client.close.assert_awaited_once()
async def test_pipeline_error_propagates_and_closes_client(self, monkeypatch):
tests/unit/vector/test_tiered_escalation_strategy.py
+15 -9
View File
@@ -36,15 +36,16 @@ def _job(queue: str = pq.INGEST_QUEUE_FAST, attempts: int = 1) -> Job:
class TestLadder:
def test_next_tier_ordering(self):
assert next_tier("fast") == "structured"
assert next_tier("structured") == "ocr"
assert next_tier("ocr") is None # terminal
assert next_tier("structured") == "ocr-incluster"
assert next_tier("ocr-incluster") == "ocr-upstream"
assert next_tier("ocr-upstream") is None # terminal
assert next_tier("unknown") is None
def test_ladder_is_cheapest_first(self):
assert TIER_LADDER == ("fast", "structured", "ocr")
assert TIER_LADDER == ("fast", "structured", "ocr-incluster", "ocr-upstream")
def test_tier_for_queue(self):
assert pq.tier_for_queue(pq.INGEST_QUEUE_OCR) == "ocr"
assert pq.tier_for_queue(pq.INGEST_QUEUE_OCR_UPSTREAM) == "ocr-upstream"
assert pq.tier_for_queue(pq.INGEST_QUEUE_STRUCTURED) == "structured"
# Legacy / unknown / None all fall back to the cheapest tier.
assert pq.tier_for_queue(pq.LEGACY_INGEST_QUEUE) == "fast"
@@ -56,10 +57,12 @@ class TestTieredEscalationStrategy:
return pq.TieredEscalationStrategy(max_transient_attempts=max_transient)
def test_escalate_hops_to_target_queue(self):
exc = EscalateError(from_tier="fast", to_tier="ocr", reason="empty_text")
exc = EscalateError(
from_tier="fast", to_tier="ocr-upstream", reason="empty_text"
)
decision = self._strategy().get_retry_decision(exception=exc, job=_job())
assert decision is not None
assert decision.queue == pq.INGEST_QUEUE_OCR
assert decision.queue == pq.INGEST_QUEUE_OCR_UPSTREAM
def test_escalate_to_structured(self):
exc = EscalateError(
@@ -75,11 +78,13 @@ class TestTieredEscalationStrategy:
assert decision is None
def test_escalate_unwraps_exception_group(self):
exc = EscalateError(from_tier="fast", to_tier="ocr", reason="empty_text")
exc = EscalateError(
from_tier="fast", to_tier="ocr-upstream", reason="empty_text"
)
group = ExceptionGroup("wrapped", [exc])
decision = self._strategy().get_retry_decision(exception=group, job=_job())
assert decision is not None
assert decision.queue == pq.INGEST_QUEUE_OCR
assert decision.queue == pq.INGEST_QUEUE_OCR_UPSTREAM
def test_transient_retries_same_queue_under_cap(self):
decision = self._strategy(max_transient=5).get_retry_decision(
@@ -126,7 +131,8 @@ class TestTieredEscalationStrategy:
# Batch OCR re-poll (Deck #332): same-queue deferral after retry_in.
before = datetime.now(timezone.utc)
decision = self._strategy().get_retry_decision(
exception=BatchPending(retry_in=120), job=_job(queue=pq.INGEST_QUEUE_OCR)
exception=BatchPending(retry_in=120),
job=_job(queue=pq.INGEST_QUEUE_OCR_UPSTREAM),
)
after = datetime.now(timezone.utc)
assert decision is not None