feat(ocr): opt-in batch OCR mode via the gateway's async batch routes

Add DOCUMENT_OCR_MODE=sync|batch (default sync). In batch mode the tier-3 OCR
processor submits documents to the embedding gateway's async Batch OCR routes
(POST /v1/ocr/batch + GET /v1/ocr/batch/{job_id}, astrolabe-cloud-website#372)
for ~50% cheaper large-corpus backfill. The direct Mistral OCR path is left
untouched. Tracked on Deck #332.

Batch jobs run minutes-hours, so the OCR tier cannot block (the procrastinate
worker reclaims jobs in `doing` after INGEST_STALLED_JOB_SECONDS). Instead it
submits, records the gateway job id in a new per-tenant `batch_ocr_jobs` table
(procrastinate args are immutable across retries), and raises a BatchPending
signal that TieredEscalationStrategy turns into a same-queue deferred re-poll —
releasing the worker slot between polls. On completion the per-page markdown is
indexed like the sync path; a failure or a job past
DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS marks the document parse-failed.

Batch is opt-in and gateway-only: with the direct mistral backend, no gateway
URL, or the inline/memory pipeline (which can't defer), it falls back to sync.
One batch job per document (coalescing N docs/job is a follow-up).

- embedding/gateway_batch_client.py: submit/poll client (reuses GatewayTokenProvider).
- vector/batch_ocr_store.py + migration 008: job tracking (portable SQLite+PG).
- document_processors/escalation.py: BatchPending control-flow signal.
- document_processors/ocr.py: batch state machine + sync fallback.
- vector/processor.py: thread doc identity to the OCR tier; raise BatchPending
  from the pending sentinel; propagate it as control flow (not a failure).
- vector/queue/procrastinate.py: BatchPending -> same-queue retry_in, exempt
  from the transient cap (bounded by the processor's deadline).
- config + docs; tests across client/store/processor/strategy/parse-tier.

1653 unit tests 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-15 09:41:19 +02:00
co-authored by Claude Opus 4.8
parent 07ee91399b
commit 3b7e8d779b
14 changed files with 1238 additions and 59 deletions
+36
View File
@@ -572,6 +572,42 @@ A PDF larger than `DOCUMENT_MAX_PDF_SIZE_MB` fails fast with reason `oversize`
of being handed to the tiers, where a 40+ MB scan would otherwise burn the full
OCR timeout for zero recovered text.
#### OCR execution mode: synchronous vs batch (Deck #332)
The tier-3 OCR processor has two execution modes, selected by `DOCUMENT_OCR_MODE`:
```dotenv
DOCUMENT_OCR_MODE=sync # "sync" (default) | "batch"
DOCUMENT_OCR_BATCH_POLL_SECONDS=120 # re-poll cadence for a batch job (default: 120)
DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS=86400 # give up + mark timeout after this (default: 24h)
```
- **`sync`** (default) — transcribe the document inline via the backend's
synchronous path (`POST /v1/ocr` for the gateway, or the direct Mistral OCR
API). The document is parsed in a single call.
- **`batch`** — submit the document to the **gateway's async Batch OCR** job
(`POST /v1/ocr/batch`) and re-poll `GET /v1/ocr/batch/{job_id}` until it
finishes. This trades latency (a batch job runs minuteshours) for roughly
**half the OCR cost**, so it suits large-corpus backfill rather than
interactive ingest.
Batch mode is **opt-in and gateway-only**: it routes Mistral's Batch API
*through* the gateway's batch routes (no provider keys in the pod). With the
direct `mistral` backend, no `EMBEDDING_GATEWAY_URL`, or on the in-process
(`INGEST_QUEUE=memory`) pipeline — which can't defer a poll — batch transparently
**falls back to synchronous OCR**. So enabling it requires the Postgres ingest
queue (the per-tier procrastinate workers) and the gateway embedding backend.
Mechanics: the OCR tier submits the job, records its id in the `batch_ocr_jobs`
app-DB table (keyed on the document + its etag), and raises a re-poll deferral so
procrastinate re-runs the tier after `DOCUMENT_OCR_BATCH_POLL_SECONDS` — releasing
the worker slot between polls (a long batch never pins a worker or is reclaimed as
stalled). On completion the per-page markdown is indexed exactly like the sync
path; a failure or a job exceeding `DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS` marks the
document parse-failed. Each poll re-fetches + re-classifies the PDF (a known v1
inefficiency, bounded by the poll cadence); one batch job is submitted per
document (coalescing many documents per job is a planned follow-up).
### Embedding Service Configuration
The server picks an embedding provider via auto-detection. Priority order
@@ -0,0 +1,67 @@
"""Add batch_ocr_jobs table for async batch OCR job tracking.
Deck #332 / embedding-gateway batch OCR (astrolabe-cloud-website#372). When
``DOCUMENT_OCR_MODE=batch`` the OCR tier submits a document to the gateway's
async ``POST /v1/ocr/batch`` and must re-poll ``GET /v1/ocr/batch/{job_id}``
across procrastinate retries. procrastinate job args are immutable, so the
gateway ``job_id`` (and submit time, for the poll deadline) are persisted here,
keyed on the document + its content version (``etag``).
One row per in-flight job; the row is deleted once the job reaches a terminal
state. Empty + unused unless batch mode is enabled (gateway-only), so OSS/SQLite
self-hosters get an idle table and zero overhead.
Portable types only (Text + unix-epoch BigInteger timestamps, like the rest of
this schema except the CP-queried usage_events) so the same migration runs on
both self-host SQLite and cloud Postgres.
Revision ID: 008
Revises: 007
Create Date: 2026-06-15 12:00:00.000000
"""
import sqlalchemy as sa
from alembic import op
revision = "008"
down_revision = "007"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.create_table(
"batch_ocr_jobs",
# Document identity (the same keys the OCR tier receives via the
# processor ``options``). ``etag`` is the content-version key: a changed
# document (new etag) is a new job, so a stale row never serves results
# for the wrong content.
sa.Column("user_id", sa.Text(), nullable=False),
sa.Column("doc_id", sa.Text(), nullable=False),
sa.Column("doc_type", sa.Text(), nullable=False),
sa.Column("etag", sa.Text(), nullable=False),
# The gateway's namespaced batch job id ("<provider>/<batch_job_id>") —
# the only handle for polling (the gateway is stateless).
sa.Column("job_id", sa.Text(), nullable=False),
# Gateway-normalised status mirror (pending|succeeded|failed). Kept for
# observability; the live decision always comes from a fresh poll.
sa.Column("status", sa.Text(), nullable=False),
# Unix-epoch seconds. ``submitted_at`` anchors the poll deadline
# (DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS).
sa.Column("submitted_at", sa.BigInteger(), nullable=False),
sa.Column("updated_at", sa.BigInteger(), nullable=False),
# One in-flight job per (document, content version). A resubmit for a new
# etag inserts a new row; the superseded row is swept on resubmit.
sa.UniqueConstraint(
"user_id",
"doc_id",
"doc_type",
"etag",
name="uq_batch_ocr_jobs_doc",
),
)
def downgrade() -> None:
op.drop_table("batch_ocr_jobs")
+26
View File
@@ -175,6 +175,20 @@ _DEFAULTS: dict[str, Any] = {
# 20-60s; raise/lower per tenant. Configurable so a tenant isn't stuck with
# the 180s default when its gateway has its own shorter ceiling.
"document_ocr_timeout_seconds": 180.0,
# OCR execution mode (Deck #332). "sync" (default) transcribes inline via the
# backend's synchronous path. "batch" routes to the gateway's async Batch OCR
# job (~50% cheaper, minutes-hours latency) for large-corpus backfill — opt-in
# and gateway-only; with the direct mistral backend or no gateway it falls
# back to sync. The submit->defer-poll loop runs on the per-tier procrastinate
# path; the inline/memory pool can't defer, so batch falls back to sync there.
"document_ocr_mode": "sync",
# Seconds between batch-job polls (the procrastinate re-enqueue delay). Each
# poll re-runs the tier; keep it well above a few seconds.
"document_ocr_batch_poll_seconds": 120,
# Hard deadline (seconds from submit) after which a still-pending batch job is
# abandoned and the document marked parse-failed (timeout). Matches the
# gateway's 24h Batch timeout default.
"document_ocr_batch_max_wait_seconds": 86400,
# Observability
"metrics_enabled": True,
"metrics_port": 9090,
@@ -352,6 +366,11 @@ _dynaconf = Dynaconf(
Validator("DOCUMENT_CHUNK_SIZE", gte=1),
Validator("DOCUMENT_PARSE_TIMEOUT_SECONDS", gte=1),
Validator("DOCUMENT_OCR_TIMEOUT_SECONDS", gte=1),
Validator("DOCUMENT_OCR_MODE", is_in=("sync", "batch")),
# Poll cadence well above a few seconds (each poll re-runs the tier);
# deadline at least one poll interval.
Validator("DOCUMENT_OCR_BATCH_POLL_SECONDS", gte=5),
Validator("DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS", gte=60),
Validator("DOCUMENT_PARSE_MEM_LIMIT_MB", gte=128),
# 0 disables the pre-parse PDF size cap; otherwise it must be positive.
Validator("DOCUMENT_MAX_PDF_SIZE_MB", gte=0),
@@ -859,6 +878,13 @@ class Settings:
# parse timeout / httpx.Timeout; per-tenant tunable so a gateway with a
# shorter ceiling isn't masked by the 180s default.
document_ocr_timeout_seconds: float = 180.0
# OCR execution mode: "sync" | "batch" (Deck #332). batch is opt-in,
# gateway-only, and used for large-corpus backfill; it falls back to sync when
# no gateway backend resolves or the path can't defer (inline/memory pool).
document_ocr_mode: str = "sync"
# Batch-job poll cadence (procrastinate re-enqueue delay) and hard deadline.
document_ocr_batch_poll_seconds: int = 120
document_ocr_batch_max_wait_seconds: int = 86400
# OCR escalation triggers (tier-0), per-tenant tunable. A page is OCR-worthy
# if near-empty (< min_page_chars) OR low text-quality (< min_text_quality)
# OR (when detect_scanned, image-analysis only runs when OCR is enabled)
@@ -93,3 +93,26 @@ class EscalateError(Exception):
super().__init__(
f"escalate {from_tier}->{to_tier} (reason={reason})",
)
class BatchPending(Exception):
"""Raised when a tier's work is in flight on an async backend and the worker
should poll again later (Deck #332 — batch OCR).
Like :class:`EscalateError` it is a **control-flow signal, NOT a failure**:
the document's batch OCR job is still running on the gateway, so the OCR tier
submits it (or polls an existing job) and raises this to ask the procrastinate
retry strategy to re-run the SAME job on the SAME queue after ``retry_in``
seconds — releasing the worker slot meanwhile so a multi-minute/hour batch
doesn't pin a worker (and isn't reclaimed as a stalled ``doing`` job).
It must propagate untouched to the retry strategy: never swallowed by a broad
``except Exception`` on the indexing path, never counted as a drop/parse
error, and never marks the placeholder failed (the doc isn't done yet).
Unlike ``EscalateError`` it does NOT change queue — the job stays on its own
(``ocr``) tier queue and is simply deferred.
"""
def __init__(self, *, retry_in: int) -> None:
self.retry_in = retry_in
super().__init__(f"batch OCR pending (retry_in={retry_in}s)")
+240 -26
View File
@@ -18,6 +18,7 @@ Mistral (if ``MISTRAL_API_KEY``). Both return GitHub-flavoured markdown + exact
import base64
import logging
import time
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from typing import Any
@@ -35,6 +36,13 @@ logger = logging.getLogger(__name__)
# configurable via DOCUMENT_OCR_TIMEOUT_SECONDS and resolved per call.
_OCR_CONNECT_TIMEOUT_SECONDS = 10.0
# Sentinel keys on a ProcessingResult.metadata that mark "batch OCR job still in
# flight — poll again later". The processor can't raise across the registry, so
# it returns this sentinel and ``vector/processor._parse_pdf_tier`` translates it
# into a ``BatchPending`` control-flow raise (same site as ``EscalateError``).
OCR_BATCH_PENDING_KEY = "ocr_batch_pending"
OCR_BATCH_RETRY_IN_KEY = "ocr_batch_retry_in"
def _pages_to_text(
pages: list[tuple[int, str]],
@@ -64,6 +72,24 @@ def _pages_to_text(
return "".join(parts), boundaries
def _batch_identity(
options: dict[str, Any] | None,
) -> tuple[str, str, str, str] | None:
"""Extract ``(user_id, doc_id, doc_type, etag)`` from the processor options
the per-tier path threads in, or ``None`` if identity is absent (the inline
pool, which can't defer a poll). ``etag`` may be empty (a file with no etag is
still one tracked job keyed on "").
"""
if not options:
return None
user_id = options.get("user_id")
doc_id = options.get("doc_id")
doc_type = options.get("doc_type")
if not user_id or not doc_id or not doc_type:
return None
return str(user_id), str(doc_id), str(doc_type), str(options.get("etag") or "")
class _OcrBackend(ABC):
@abstractmethod
async def ocr(
@@ -141,6 +167,54 @@ class _MistralOcrBackend(_OcrBackend):
return _pages_to_text(pages)
def _build_gateway_token_provider(settings: Settings) -> Any:
"""Build the M2M ``GatewayTokenProvider`` from settings, or ``None`` when no
client-id is configured (unauthenticated gateway). Shared by the sync OCR
backend and the batch client so the M2M-triple validation lives in one place.
"""
if not settings.embedding_gateway_client_id:
return None
# Lazy import avoids a document_processors -> embedding cycle at load.
from ..embedding.gateway_client import GatewayTokenProvider # noqa: PLC0415
# Explicit (not assert -- assert is stripped under `python -O`): the M2M
# triple is all-or-nothing.
if not settings.embedding_gateway_token_url:
raise ValueError(
"EMBEDDING_GATEWAY_TOKEN_URL is required when "
"EMBEDDING_GATEWAY_CLIENT_ID is set"
)
if not settings.embedding_gateway_client_secret:
raise ValueError(
"EMBEDDING_GATEWAY_CLIENT_SECRET is required when "
"EMBEDDING_GATEWAY_CLIENT_ID is set"
)
return GatewayTokenProvider(
token_url=settings.embedding_gateway_token_url,
client_id=settings.embedding_gateway_client_id,
client_secret=settings.embedding_gateway_client_secret,
scope=settings.embedding_gateway_scope,
)
def build_gateway_batch_client(settings: Settings) -> Any:
"""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."""
if settings.document_ocr_provider not in ("gateway", "auto"):
return None
if not settings.embedding_gateway_url:
return None
from ..embedding.gateway_batch_client import GatewayBatchOcrClient # noqa: PLC0415
return GatewayBatchOcrClient(
settings.embedding_gateway_url,
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."""
provider = settings.document_ocr_provider
@@ -148,33 +222,10 @@ def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
return None
if provider in ("gateway", "auto") and settings.embedding_gateway_url:
token_provider = None
if settings.embedding_gateway_client_id:
# Lazy import avoids a document_processors -> embedding cycle at load.
from ..embedding.gateway_client import ( # noqa: PLC0415
GatewayTokenProvider,
)
# Explicit (not assert -- assert is stripped under `python -O`): the
# M2M triple is all-or-nothing.
if not settings.embedding_gateway_token_url:
raise ValueError(
"EMBEDDING_GATEWAY_TOKEN_URL is required when "
"EMBEDDING_GATEWAY_CLIENT_ID is set"
)
if not settings.embedding_gateway_client_secret:
raise ValueError(
"EMBEDDING_GATEWAY_CLIENT_SECRET is required when "
"EMBEDDING_GATEWAY_CLIENT_ID is set"
)
token_provider = GatewayTokenProvider(
token_url=settings.embedding_gateway_token_url,
client_id=settings.embedding_gateway_client_id,
client_secret=settings.embedding_gateway_client_secret,
scope=settings.embedding_gateway_scope,
)
return _GatewayOcrBackend(
settings.embedding_gateway_url, settings.document_ocr_model, token_provider
settings.embedding_gateway_url,
settings.document_ocr_model,
_build_gateway_token_provider(settings),
)
if provider in ("mistral", "auto") and settings.mistral_api_key:
@@ -215,6 +266,13 @@ class OcrProcessor(DocumentProcessor):
# doesn't each build a backend (and fetch its own M2M token). Lazy-init:
# anyio primitives must not be created at import time.
self._backend_lock: anyio.Lock | None = None
# Batch-mode (Deck #332): the gateway batch client is cached like the sync
# backend so its GatewayTokenProvider keeps its M2M-token cache across
# documents. ``_batch_fallback_warned`` rate-limits the "can't batch,
# using sync" warning to once per pod.
self._batch_client_resolved = False
self._batch_client: Any = None
self._batch_fallback_warned = False
@property
def name(self) -> str:
@@ -239,6 +297,19 @@ class OcrProcessor(DocumentProcessor):
) = None,
) -> ProcessingResult:
settings = get_settings()
# Batch mode (Deck #332): submit to the gateway's async Batch OCR job and
# poll across procrastinate retries. Returns a result (incl. the
# "pending" sentinel) when handled, or None to fall back to the
# synchronous path below (no gateway backend, or no per-doc identity — the
# inline/memory pool can't defer a poll).
if settings.document_ocr_mode == "batch":
batch_result = await self._process_batch(
content, content_type, filename, options, settings
)
if batch_result is not None:
return batch_result
if not self._backend_resolved:
if self._backend_lock is None:
self._backend_lock = anyio.Lock()
@@ -301,6 +372,149 @@ class OcrProcessor(DocumentProcessor):
processor=self.name,
)
async def _get_batch_client(self) -> Any:
"""Cached gateway batch client (or ``None`` when batch isn't applicable —
provider=mistral / no gateway). Resolved once under the backend lock so the
token provider's M2M cache survives across documents."""
if not self._batch_client_resolved:
if self._backend_lock is None:
self._backend_lock = anyio.Lock()
async with self._backend_lock:
if not self._batch_client_resolved: # double-checked
self._batch_client = build_gateway_batch_client(get_settings())
self._batch_client_resolved = True
return self._batch_client
def _batch_fallback(self, reason: str, filename: str | None) -> None:
"""Warn once that batch mode is falling back to the synchronous path."""
if not self._batch_fallback_warned:
logger.warning(
"DOCUMENT_OCR_MODE=batch but %s; falling back to synchronous OCR",
reason,
)
self._batch_fallback_warned = True
async def _process_batch(
self,
content: bytes,
content_type: str,
filename: str | None,
options: dict[str, Any] | None,
settings: Settings,
) -> ProcessingResult | None:
"""Submit + poll a one-document batch OCR job.
Returns a :class:`ProcessingResult` when batch handled the document — the
terminal success/failure result, or the *pending sentinel* (``success=False``
+ ``OCR_BATCH_PENDING_KEY`` metadata) that ``_parse_pdf_tier`` turns into a
``BatchPending`` re-poll. Returns ``None`` to fall back to synchronous OCR
(no gateway backend, or no per-doc identity — the inline pool can't defer).
"""
# Per-doc identity is threaded via ``options`` only on the per-tier
# procrastinate path; the inline/memory pool omits it and can't defer a
# poll, so batch is inapplicable there.
identity = _batch_identity(options)
if identity is None:
self._batch_fallback("no per-document identity (inline path)", filename)
return None
client = await self._get_batch_client()
if client is None:
self._batch_fallback(
"no gateway backend (provider=mistral or EMBEDDING_GATEWAY_URL unset)",
filename,
)
return None
# Lazy import: keep the vector/DB stack off the document_processors load
# path (mirrors the EscalateError lazy import in vector/processor).
from ..vector.batch_ocr_store import BatchOcrJobStore # noqa: PLC0415
user_id, doc_id, doc_type, etag = identity
store = await BatchOcrJobStore.shared()
mime = content_type.split(";")[0].strip().lower()
poll_seconds = settings.document_ocr_batch_poll_seconds
job = await store.get(
user_id=user_id, doc_id=doc_id, doc_type=doc_type, etag=etag
)
if job is None:
# New submission. Drop any superseded-version rows for this doc first
# (a re-edited file changes etag), then submit + record.
await store.delete_stale_for_doc(
user_id=user_id, doc_id=doc_id, doc_type=doc_type, keep_etag=etag
)
job_id = await client.submit(content, mime, custom_id=doc_id)
await store.insert_pending(
user_id=user_id,
doc_id=doc_id,
doc_type=doc_type,
etag=etag,
job_id=job_id,
)
logger.info(
"batch OCR job submitted for %s (job_id=%s); deferring poll",
filename or doc_id,
job_id,
)
return self._pending(poll_seconds)
# Existing job — poll the gateway.
result = await client.poll(job.job_id)
if result.is_pending:
elapsed = int(time.time()) - job.submitted_at
if elapsed >= settings.document_ocr_batch_max_wait_seconds:
await store.delete(
user_id=user_id, doc_id=doc_id, doc_type=doc_type, etag=etag
)
logger.warning(
"batch OCR job %s exceeded max wait (%ss); marking failed",
job.job_id,
settings.document_ocr_batch_max_wait_seconds,
)
return ProcessingResult(
text="",
metadata={"parse_failed_reason": "timeout"},
processor=self.name,
success=False,
error="batch OCR timed out",
)
return self._pending(poll_seconds)
# Terminal — drop the tracking row either way.
await store.delete(user_id=user_id, doc_id=doc_id, doc_type=doc_type, etag=etag)
if result.is_failed:
logger.warning(
"batch OCR job %s failed: %s", job.job_id, result.error or "unknown"
)
return ProcessingResult(
text="",
metadata={"parse_failed_reason": "error"},
processor=self.name,
success=False,
error=f"batch OCR failed: {result.error or 'unknown'}",
)
text, boundaries = _pages_to_text(result.pages)
return ProcessingResult(
text=text,
metadata={
"page_count": len(boundaries),
"page_boundaries": boundaries,
"file_size": len(content),
},
processor=self.name,
)
def _pending(self, retry_in: int) -> ProcessingResult:
"""The pending sentinel — ``_parse_pdf_tier`` raises ``BatchPending`` from
it. ``success=False`` keeps it out of the index path, and the sentinel key
keeps it out of the parse-failed path (it isn't a failure)."""
return ProcessingResult(
text="",
metadata={OCR_BATCH_PENDING_KEY: True, OCR_BATCH_RETRY_IN_KEY: retry_in},
processor=self.name,
success=False,
)
async def health_check(self) -> bool:
# Backends are resolved lazily (and configured per tenant), so there is
# nothing to probe here without making a billable upstream call -- the
@@ -0,0 +1,171 @@
"""Client for the embedding gateway's async **batch OCR** routes (Deck #332).
The gateway exposes two batch routes alongside the synchronous ``POST /v1/ocr``
(astrolabe-cloud-website#372):
- ``POST /v1/ocr/batch`` — submit N documents (each with a caller ``custom_id``)
as one Mistral Batch job; returns ``202`` + a namespaced ``job_id``
(``<provider>/<batch_job_id>``).
- ``GET /v1/ocr/batch/{job_id}`` — poll; returns the lifecycle status and, once
terminal, per-document results (per-page markdown, or a per-document error).
The gateway is a **stateless passthrough** to Mistral's Batch API — the
``job_id`` is the only handle, so the worker persists it (see
``vector/batch_ocr_store``) and re-polls across procrastinate retries.
This client submits exactly **one document per job** (the v1 unit; coalescing N
docs/job is a follow-up). Auth + ``/v1`` base-url handling mirror the synchronous
:class:`~nextcloud_mcp_server.embedding.gateway_client.GatewayProvider` /
``_GatewayOcrBackend`` — same M2M :class:`GatewayTokenProvider` bearer, no
provider keys in the pod.
"""
from __future__ import annotations
import base64
import logging
from dataclasses import dataclass
from typing import Any
import httpx
from .gateway_client import GatewayTokenProvider
logger = logging.getLogger(__name__)
# Connect timeout for the (cheap) submit/poll calls. These are control-plane-ish
# requests — a submit returns immediately with a job id and a poll is a status
# read — so they get a short, fixed timeout, NOT the document-OCR read timeout
# (which sizes a synchronous transcription).
_BATCH_CONNECT_TIMEOUT_SECONDS = 5.0
_BATCH_REQUEST_TIMEOUT_SECONDS = 30.0
# Gateway-normalised batch lifecycle (OcrBatchStatus on the gateway side).
_PENDING = "pending"
_SUCCEEDED = "succeeded"
_FAILED = "failed"
@dataclass(frozen=True)
class BatchPollResult:
"""One poll of a batch OCR job.
``status`` is the gateway-normalised lifecycle (``pending`` | ``succeeded`` |
``failed``). For a single-document job: on ``succeeded`` ``pages`` holds the
document's per-page ``(index, markdown)`` (empty + ``error`` set if that one
document errored inside an otherwise-successful job); on ``failed`` ``error``
carries the job-level failure.
"""
status: str
pages: list[tuple[int, str]]
error: str | None = None
@property
def is_pending(self) -> bool:
return self.status == _PENDING
@property
def is_succeeded(self) -> bool:
return self.status == _SUCCEEDED
@property
def is_failed(self) -> bool:
return self.status == _FAILED
class GatewayBatchOcrClient:
"""Submits + polls single-document batch OCR jobs against the gateway."""
def __init__(
self,
base_url: str,
model: str,
token_provider: GatewayTokenProvider | None = None,
) -> None:
# EMBEDDING_GATEWAY_URL is a bare origin; the batch routes live under /v1
# like the rest of the gateway API. Idempotent if already /v1-suffixed.
base = base_url.rstrip("/")
if not base.endswith("/v1"):
base = f"{base}/v1"
self._base = base
self._model = model
self._token_provider = token_provider
async def _headers(self) -> dict[str, str]:
if self._token_provider is None:
return {}
return {"Authorization": f"Bearer {await self._token_provider.get_token()}"}
async def submit(self, content: bytes, mime_type: str, custom_id: str) -> str:
"""Submit ``content`` as a one-document batch job; return the namespaced
``job_id`` to persist + poll. Raises on transport / non-2xx."""
payload = {
"model": self._model,
"documents": [
{
"custom_id": custom_id,
"mime_type": mime_type,
"document_b64": base64.b64encode(content).decode("ascii"),
}
],
}
async with httpx.AsyncClient(
timeout=httpx.Timeout(
_BATCH_REQUEST_TIMEOUT_SECONDS, connect=_BATCH_CONNECT_TIMEOUT_SECONDS
)
) as client:
resp = await client.post(
f"{self._base}/ocr/batch", json=payload, headers=await self._headers()
)
resp.raise_for_status()
body = resp.json()
job_id = body["job_id"]
logger.info(
"batch OCR submitted: job_id=%s custom_id=%s status=%s",
job_id,
custom_id,
body.get("status"),
)
return job_id
async def poll(self, job_id: str) -> BatchPollResult:
"""Poll a batch job. Raises on transport / non-2xx; maps a terminal job's
single-document result into :class:`BatchPollResult`."""
async with httpx.AsyncClient(
timeout=httpx.Timeout(
_BATCH_REQUEST_TIMEOUT_SECONDS, connect=_BATCH_CONNECT_TIMEOUT_SECONDS
)
) as client:
resp = await client.get(
f"{self._base}/ocr/batch/{job_id}", headers=await self._headers()
)
resp.raise_for_status()
body = resp.json()
status = body.get("status", _PENDING)
if status != _SUCCEEDED:
# pending: nothing to read yet. failed: surface the job-level error.
return BatchPollResult(status=status, pages=[], error=body.get("error"))
return _result_from_success(body)
def _result_from_success(body: dict[str, Any]) -> BatchPollResult:
"""Extract the single document's pages from a succeeded job's results.
Submitting one document per job means exactly one result item; defensively
take the first. A per-document error inside a succeeded job (the document
failed but the job didn't) surfaces as a failed poll so the caller marks the
doc parse-failed rather than indexing empty text.
"""
results = body.get("results") or []
if not results:
return BatchPollResult(
status=_FAILED,
pages=[],
error="batch job succeeded but returned no results",
)
item = results[0]
if item.get("error") is not None or item.get("pages") is None:
return BatchPollResult(status=_FAILED, pages=[], error=item.get("error"))
pages = [(p["index"], p.get("markdown", "")) for p in item["pages"]]
return BatchPollResult(status=_SUCCEEDED, pages=pages)
@@ -0,0 +1,121 @@
"""Tracking store for in-flight async batch OCR jobs (Deck #332).
When ``DOCUMENT_OCR_MODE=batch`` the OCR tier submits a document to the gateway's
async batch route and then re-polls across procrastinate retries. procrastinate
job args are immutable, so the gateway ``job_id`` (and submit time, for the poll
deadline) live in the ``batch_ocr_jobs`` app-DB table, keyed on the document +
its content version (``etag``).
Engine reuse mirrors :class:`~nextcloud_mcp_server.usage.store.UsageEventStore`:
rather than open its own engine this store borrows the process-wide
:class:`RefreshTokenStorage` singleton (``get_shared_storage()``) — same app DB,
dialect handling, ``?``-placeholder shim, and the guarantee that Alembic
migrations (incl. ``batch_ocr_jobs``) already ran.
"""
from __future__ import annotations
import logging
import time
from dataclasses import dataclass
import anyio
from nextcloud_mcp_server.auth.storage import RefreshTokenStorage, get_shared_storage
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class BatchOcrJob:
"""A tracked in-flight batch OCR job."""
job_id: str
status: str
submitted_at: int
class BatchOcrJobStore:
"""CRUD for the ``batch_ocr_jobs`` table (one row per in-flight job)."""
_shared_instance: "BatchOcrJobStore | None" = None
_shared_lock: anyio.Lock = anyio.Lock()
def __init__(self, storage: RefreshTokenStorage) -> None:
self._storage = storage
@classmethod
async def shared(cls) -> "BatchOcrJobStore":
"""Process-wide store backed by the storage singleton (mirrors
``UsageEventStore.shared``). Tests should construct
``BatchOcrJobStore(storage)`` directly — the cache is a process global
with no teardown hook."""
async with cls._shared_lock:
if cls._shared_instance is None:
cls._shared_instance = cls(await get_shared_storage())
return cls._shared_instance
async def get(
self, *, user_id: str, doc_id: str, doc_type: str, etag: str
) -> BatchOcrJob | None:
"""The in-flight job for this document+version, or ``None``."""
async with self._storage.acquire() as db:
async with db.execute(
"SELECT job_id, status, submitted_at FROM batch_ocr_jobs "
"WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag = ?",
(user_id, doc_id, doc_type, etag),
) as cursor:
row = await cursor.fetchone()
if row is None:
return None
return BatchOcrJob(job_id=row[0], status=row[1], submitted_at=int(row[2]))
async def insert_pending(
self,
*,
user_id: str,
doc_id: str,
doc_type: str,
etag: str,
job_id: str,
submitted_at: int | None = None,
) -> None:
"""Record a freshly-submitted job. ``ON CONFLICT DO NOTHING`` makes a
racing double-submit harmless (the first row wins; the loser's job id is
abandoned and reaped by the gateway-side file purge)."""
now = submitted_at if submitted_at is not None else int(time.time())
async with self._storage.acquire() as db:
await db.execute(
"INSERT INTO batch_ocr_jobs "
"(user_id, doc_id, doc_type, etag, job_id, status, submitted_at, updated_at) "
"VALUES (?, ?, ?, ?, ?, ?, ?, ?) "
"ON CONFLICT (user_id, doc_id, doc_type, etag) DO NOTHING",
(user_id, doc_id, doc_type, etag, job_id, "pending", now, now),
)
await db.commit()
async def delete(
self, *, user_id: str, doc_id: str, doc_type: str, etag: str
) -> None:
"""Drop the row once the job is terminal (succeeded or failed)."""
async with self._storage.acquire() as db:
await db.execute(
"DELETE FROM batch_ocr_jobs "
"WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag = ?",
(user_id, doc_id, doc_type, etag),
)
await db.commit()
async def delete_stale_for_doc(
self, *, user_id: str, doc_id: str, doc_type: str, keep_etag: str
) -> None:
"""Remove superseded-version rows for a document (any etag other than the
current one) before a resubmit, so a re-edited file doesn't leave its
old in-flight job tracked forever."""
async with self._storage.acquire() as db:
await db.execute(
"DELETE FROM batch_ocr_jobs "
"WHERE user_id = ? AND doc_id = ? AND doc_type = ? AND etag != ?",
(user_id, doc_id, doc_type, keep_etag),
)
await db.commit()
+55 -27
View File
@@ -126,6 +126,7 @@ async def _parse_pdf_tier(
filename: str | None,
tier: str,
settings: Any,
options: dict[str, Any] | None = None,
) -> "ProcessingResult":
"""Run a single extraction tier and apply the post-parse escalation gate.
@@ -141,18 +142,31 @@ async def _parse_pdf_tier(
the "OCR is an enhancement, never worse than off" invariant: a tenant who has
not enabled a higher tier (or has no processor for it) simply indexes the
cheap tier's output.
Batch OCR (Deck #332): when the OCR tier's batch job is still in flight the
processor returns a *pending sentinel* result; we translate it here into a
``BatchPending`` raise (same decision point as ``EscalateError``) so the
retry strategy re-runs this tier after a delay instead of indexing empty text.
"""
# Lazy import: keep the document stack (pymupdf/_isolation) off the module
# load path; this runs only on the per-tier worker, which needs it anyway.
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
BatchPending,
EscalateError,
)
from nextcloud_mcp_server.document_processors.ocr import ( # noqa: PLC0415
OCR_BATCH_PENDING_KEY,
OCR_BATCH_RETRY_IN_KEY,
)
# options / progress_callback are not threaded here -- the indexing caller
# passes neither today, and the inline path (registry.process) omits them
# too. Forward them if a tier processor ever needs per-call tuning (e.g. OCR
# DPI); keeping the two paths symmetric until then.
result = await registry.process_tier(content, content_type, filename, tier)
# ``options`` threads per-document identity (user_id/doc_id/doc_type/etag) to
# the OCR tier so batch mode can key its job-tracking table (Deck #332). Other
# tiers ignore it. The inline path (registry.process) passes None.
result = await registry.process_tier(
content, content_type, filename, tier, options=options
)
if result.metadata.get(OCR_BATCH_PENDING_KEY):
raise BatchPending(retry_in=int(result.metadata[OCR_BATCH_RETRY_IN_KEY]))
if result.success:
decision = registry.evaluate_escalation(
result, content, tier, settings, filename=filename
@@ -525,19 +539,21 @@ async def process_document(
re-picked on the next scan; the procrastinate path (max_retries=1) caps it at
one outer attempt (~30s) and defers. Don't stack a third retry layer here.
"""
# EscalateError is a control-flow signal that arises ONLY on the per-tier
# external path (tier set). Bind the class lazily here, and only when a tier
# is set, so the document stack is never imported at *module load* (the #877
# invariant) nor on the delete / text-doc call paths (file processing already
# imports it via get_registry regardless). When tier is None it can't be
# raised, so the guards below stay inert.
escalate_error_cls: type[BaseException] | None = None
# EscalateError and BatchPending are control-flow signals that arise ONLY on
# the per-tier external path (tier set). Bind them lazily here, and only when a
# tier is set, so the document stack is never imported at *module load* (the
# #877 invariant) nor on the delete / text-doc call paths (file processing
# already imports them via get_registry regardless). When tier is None neither
# can be raised, so the guards below stay inert. Bound as a tuple so the guards
# treat both identically: propagate untouched, never record an error/drop.
control_flow_excs: tuple[type[BaseException], ...] = ()
if tier is not None:
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
BatchPending,
EscalateError,
)
escalate_error_cls = EscalateError
control_flow_excs = (EscalateError, BatchPending)
start_time = time.time()
@@ -632,13 +648,11 @@ async def process_document(
return # Success
except Exception as e:
# An escalation signal is control flow, not a failure:
# propagate it untouched so the procrastinate retry strategy
# can hop the job to the next tier's queue. Never retry it
# A control-flow signal (escalation hop, or batch-OCR re-poll
# deferral) is not a failure: propagate it untouched so the
# procrastinate retry strategy handles it. Never retry it
# in-process and never count it as a drop.
if escalate_error_cls is not None and isinstance(
e, escalate_error_cls
):
if isinstance(e, control_flow_excs):
raise
if attempt < max_retries - 1:
logger.warning(
@@ -691,10 +705,11 @@ async def process_document(
raise
except Exception as e:
# An escalation signal must reach the procrastinate retry strategy
# un-recorded -- it is neither a processing success nor an error
# (the hop is its own event, counted via record_document_escalation).
if escalate_error_cls is not None and isinstance(e, escalate_error_cls):
# A control-flow signal must reach the procrastinate retry strategy
# un-recorded -- it is neither a processing success nor an error (an
# escalation hop is counted via record_document_escalation; a batch
# re-poll deferral is not an event at all).
if isinstance(e, control_flow_excs):
raise
# Single processing-error call site: catches exhausted-retry
# re-raises, delete failures, and setup errors (get_qdrant_client /
@@ -956,6 +971,7 @@ async def _index_document(
get_registry,
)
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
BatchPending,
EscalateError,
)
@@ -968,6 +984,15 @@ async def _index_document(
# and the in-process/memory pool (tier is None) -- runs the inline
# tiered pipeline (fast -> OCR escalation in one call).
if tier is not None and _is_pdf(content_type):
# Thread per-document identity to the OCR tier so batch mode
# (Deck #332) can key its job-tracking table; other tiers
# ignore it.
ocr_options = {
"user_id": doc_task.user_id,
"doc_id": doc_task.doc_id,
"doc_type": doc_task.doc_type,
"etag": doc_task.etag or "",
}
result = await _parse_pdf_tier(
registry,
content_bytes,
@@ -975,6 +1000,7 @@ async def _index_document(
file_path,
tier,
settings,
options=ocr_options,
)
else:
result = await registry.process(
@@ -1047,10 +1073,12 @@ async def _index_document(
)
else:
logger.debug("No page_boundaries in metadata for %s", file_path)
except EscalateError:
# Control-flow signal (per-tier path): re-raise untouched so the
# queue hops the job to the next tier. NOT a "failed to process"
# error -- don't log it as one.
except (EscalateError, BatchPending):
# Control-flow signals (per-tier path): re-raise untouched.
# EscalateError hops the job to the next tier; BatchPending defers
# a re-poll on the same tier (batch OCR still in flight, Deck
# #332). Neither is a "failed to process" error -- don't log them
# as one.
raise
except Exception as e:
logger.error("Failed to process file %s: %s", file_path, e)
@@ -328,13 +328,26 @@ class TieredEscalationStrategy(BaseRetryStrategy):
def get_retry_decision(
self, *, exception: BaseException, job: Job
) -> RetryDecision | None:
# Lazy import: EscalateError lives in the document stack, which the API
# pod (it also builds this App to defer) must not load. get_retry_decision
# runs only in the worker, where the stack is already imported.
from ...document_processors.escalation import EscalateError # noqa: PLC0415
# Lazy import: these live in the document stack, which the API pod (it
# also builds this App to defer) must not load. get_retry_decision runs
# only in the worker, where the stack is already imported.
from ...document_processors.escalation import ( # noqa: PLC0415
BatchPending,
EscalateError,
)
exc = _first_leaf(exception)
if isinstance(exc, BatchPending):
# Batch OCR job still in flight (Deck #332): defer a re-poll on the
# SAME queue after retry_in seconds. Deliberately exempt from the
# transient cap below — a batch job can take minutes-hours, so the
# poll count is unbounded here; the OCR processor's own deadline
# (DOCUMENT_OCR_BATCH_MAX_WAIT_SECONDS) is what terminates a stuck job.
# Releasing the worker between polls keeps the job out of `doing`, so
# it's never stall-reclaimed.
return RetryDecision(retry_in={"seconds": exc.retry_in})
if isinstance(exc, EscalateError):
queue = TIER_QUEUES.get(exc.to_tier)
if queue is None:
+78
View File
@@ -0,0 +1,78 @@
"""Unit tests for the batch OCR job-tracking store (Deck #332).
Runs against a real temp-SQLite ``RefreshTokenStorage`` (its ``initialize()``
applies the migrations, incl. ``batch_ocr_jobs``).
"""
import tempfile
from pathlib import Path
import pytest
from nextcloud_mcp_server.auth.storage import RefreshTokenStorage
from nextcloud_mcp_server.vector.batch_ocr_store import BatchOcrJobStore
pytestmark = pytest.mark.unit
@pytest.fixture
async def store():
with tempfile.TemporaryDirectory() as tmp:
storage = RefreshTokenStorage(db_path=str(Path(tmp) / "batch.db"))
await storage.initialize()
yield BatchOcrJobStore(storage)
_DOC = dict(user_id="u1", doc_id="d1", doc_type="file", etag="v1")
async def test_get_missing_returns_none(store):
assert await store.get(**_DOC) is None
async def test_insert_then_get(store):
await store.insert_pending(**_DOC, job_id="mistral/j1")
job = await store.get(**_DOC)
assert job is not None
assert job.job_id == "mistral/j1"
assert job.status == "pending"
assert job.submitted_at > 0
async def test_insert_is_idempotent_on_conflict(store):
await store.insert_pending(**_DOC, job_id="mistral/j1", submitted_at=100)
# A racing re-submit must not overwrite the first row's job id.
await store.insert_pending(**_DOC, job_id="mistral/j2", submitted_at=200)
job = await store.get(**_DOC)
assert job.job_id == "mistral/j1"
assert job.submitted_at == 100
async def test_delete(store):
await store.insert_pending(**_DOC, job_id="mistral/j1")
await store.delete(**_DOC)
assert await store.get(**_DOC) is None
async def test_delete_stale_for_doc_keeps_current_etag(store):
await store.insert_pending(
user_id="u1", doc_id="d1", doc_type="file", etag="old", job_id="mistral/old"
)
await store.insert_pending(
user_id="u1", doc_id="d1", doc_type="file", etag="new", job_id="mistral/new"
)
await store.delete_stale_for_doc(
user_id="u1", doc_id="d1", doc_type="file", keep_etag="new"
)
# Old version row gone; current one kept.
assert (
await store.get(user_id="u1", doc_id="d1", doc_type="file", etag="old")
) is None
kept = await store.get(user_id="u1", doc_id="d1", doc_type="file", etag="new")
assert kept is not None and kept.job_id == "mistral/new"
async def test_rows_are_scoped_per_document(store):
await store.insert_pending(**_DOC, job_id="mistral/j1")
other = await store.get(user_id="u1", doc_id="d2", doc_type="file", etag="v1")
assert other is None # different doc_id
+141
View File
@@ -0,0 +1,141 @@
"""Unit tests for the gateway batch OCR client (Deck #332).
HTTP is exercised via an ``httpx.MockTransport`` injected by monkeypatching
``httpx.AsyncClient`` (the repo has no respx dependency).
"""
from typing import Any, cast
import httpx
import pytest
from nextcloud_mcp_server.embedding import gateway_batch_client as gbc
pytestmark = pytest.mark.unit
def _patch_transport(monkeypatch, handler) -> list[httpx.Request]:
"""Route the client's httpx calls through ``handler``; return a list that
captures each issued request for assertions."""
seen: list[httpx.Request] = []
real = httpx.AsyncClient
def factory(*args: Any, **kwargs: Any) -> httpx.AsyncClient:
def _wrapped(request: httpx.Request) -> httpx.Response:
seen.append(request)
return handler(request)
kwargs["transport"] = httpx.MockTransport(_wrapped)
return real(*args, **kwargs)
monkeypatch.setattr(httpx, "AsyncClient", factory)
return seen
def test_base_url_normalization():
assert gbc.GatewayBatchOcrClient("http://gw", "m")._base == "http://gw/v1"
assert gbc.GatewayBatchOcrClient("http://gw/v1/", "m")._base == "http://gw/v1"
async def test_submit_posts_one_document_and_returns_job_id(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
202, json={"job_id": "mistral/job-1", "status": "pending"}
)
seen = _patch_transport(monkeypatch, handler)
client = gbc.GatewayBatchOcrClient("http://gw", "mistral/mistral-ocr-latest")
job_id = await client.submit(b"%PDF-1.7", "application/pdf", custom_id="doc-9")
assert job_id == "mistral/job-1"
req = seen[0]
assert req.method == "POST" and req.url.path == "/v1/ocr/batch"
import json
body = json.loads(req.content)
assert body["model"] == "mistral/mistral-ocr-latest"
assert len(body["documents"]) == 1
assert body["documents"][0]["custom_id"] == "doc-9"
assert body["documents"][0]["mime_type"] == "application/pdf"
assert body["documents"][0]["document_b64"] # base64 present
async def test_submit_sends_bearer_when_token_provider(monkeypatch):
class _Tok:
async def get_token(self) -> str:
return "tok-abc"
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(202, json={"job_id": "mistral/j", "status": "pending"})
seen = _patch_transport(monkeypatch, handler)
# _Tok duck-types get_token; cast for the type checker (the client only awaits
# get_token()).
client = gbc.GatewayBatchOcrClient(
"http://gw", "m", token_provider=cast(Any, _Tok())
)
await client.submit(b"x", "application/pdf", custom_id="d")
assert seen[0].headers["Authorization"] == "Bearer tok-abc"
async def test_poll_pending(monkeypatch):
_patch_transport(
monkeypatch,
lambda r: httpx.Response(200, json={"status": "pending", "total": 1}),
)
result = await gbc.GatewayBatchOcrClient("http://gw", "m").poll("mistral/j")
assert result.is_pending and result.pages == []
async def test_poll_succeeded_maps_pages(monkeypatch):
body = {
"status": "succeeded",
"results": [
{
"custom_id": "d",
"pages": [
{"index": 1, "markdown": "two"},
{"index": 0, "markdown": "one"},
],
}
],
}
_patch_transport(monkeypatch, lambda r: httpx.Response(200, json=body))
result = await gbc.GatewayBatchOcrClient("http://gw", "m").poll("mistral/j")
assert result.is_succeeded
# Order is preserved as returned; _pages_to_text sorts downstream.
assert result.pages == [(1, "two"), (0, "one")]
async def test_poll_failed_surfaces_error(monkeypatch):
_patch_transport(
monkeypatch,
lambda r: httpx.Response(200, json={"status": "failed", "error": "quota"}),
)
result = await gbc.GatewayBatchOcrClient("http://gw", "m").poll("mistral/j")
assert result.is_failed and result.error == "quota"
async def test_poll_succeeded_with_per_document_error_is_failed(monkeypatch):
body = {"status": "succeeded", "results": [{"custom_id": "d", "error": "bad page"}]}
_patch_transport(monkeypatch, lambda r: httpx.Response(200, json=body))
result = await gbc.GatewayBatchOcrClient("http://gw", "m").poll("mistral/j")
assert result.is_failed and result.error == "bad page"
async def test_poll_succeeded_no_results_is_failed(monkeypatch):
_patch_transport(
monkeypatch,
lambda r: httpx.Response(200, json={"status": "succeeded", "results": []}),
)
result = await gbc.GatewayBatchOcrClient("http://gw", "m").poll("mistral/j")
assert result.is_failed
async def test_poll_raises_on_http_error(monkeypatch):
_patch_transport(
monkeypatch, lambda r: httpx.Response(503, json={"detail": "down"})
)
with pytest.raises(httpx.HTTPStatusError):
await gbc.GatewayBatchOcrClient("http://gw", "m").poll("mistral/j")
+198 -2
View File
@@ -16,6 +16,9 @@ def _settings(**kw) -> Any: # a Settings stand-in (only the read fields matter)
document_ocr_provider="auto",
document_ocr_model="mistral/mistral-ocr-latest",
document_ocr_timeout_seconds=180.0,
document_ocr_mode="sync",
document_ocr_batch_poll_seconds=120,
document_ocr_batch_max_wait_seconds=86400,
embedding_gateway_url=None,
embedding_gateway_client_id=None,
embedding_gateway_client_secret=None,
@@ -146,7 +149,7 @@ async def test_processor_timeout_returns_timeout_reason(monkeypatch):
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is False
assert r.metadata["parse_failed_reason"] == "timeout"
assert "timed out" in r.error
assert "timed out" in (r.error or "")
async def test_gateway_httpx_timeout_maps_to_timeout_reason(monkeypatch):
@@ -165,7 +168,7 @@ async def test_gateway_httpx_timeout_maps_to_timeout_reason(monkeypatch):
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
assert r.success is False
assert r.metadata["parse_failed_reason"] == "timeout"
assert "timed out" in r.error
assert "timed out" in (r.error or "")
async def test_gateway_backend_uses_configured_timeout(mocker, monkeypatch):
@@ -218,3 +221,196 @@ async def test_mistral_backend_applies_timeout(mocker, monkeypatch):
with pytest.raises(TimeoutError):
await backend.ocr(b"%PDF-1.7", "application/pdf")
# --- batch mode (Deck #332) --------------------------------------------------
from nextcloud_mcp_server.embedding.gateway_batch_client import ( # noqa: E402
BatchPollResult,
)
from nextcloud_mcp_server.vector import batch_ocr_store as _bos # noqa: E402
_IDENTITY = {"user_id": "u1", "doc_id": "d1", "doc_type": "file", "etag": "v1"}
class _FakeStore:
"""In-memory stand-in for BatchOcrJobStore keyed like the real table."""
def __init__(self, preset=None):
self.rows: dict[tuple, Any] = {}
self.deleted: list[tuple] = []
self.stale_swept: list[tuple] = []
if preset is not None:
self.rows[("u1", "d1", "file", "v1")] = preset
async def get(self, *, user_id, doc_id, doc_type, etag):
return self.rows.get((user_id, doc_id, doc_type, etag))
async def insert_pending(
self, *, user_id, doc_id, doc_type, etag, job_id, submitted_at=None
):
self.rows[(user_id, doc_id, doc_type, etag)] = SimpleNamespace(
job_id=job_id, status="pending", submitted_at=submitted_at or 1000
)
async def delete(self, *, user_id, doc_id, doc_type, etag):
self.deleted.append((user_id, doc_id, doc_type, etag))
self.rows.pop((user_id, doc_id, doc_type, etag), None)
async def delete_stale_for_doc(self, *, user_id, doc_id, doc_type, keep_etag):
self.stale_swept.append((user_id, doc_id, doc_type, keep_etag))
class _FakeBatchClient:
def __init__(self, *, submit_job="mistral/job-1", poll=None):
self._submit_job = submit_job
self._poll = poll or BatchPollResult(status="pending", pages=[])
self.submitted: list[tuple] = []
self.polled: list[str] = []
async def submit(self, content, mime_type, custom_id):
self.submitted.append((content, mime_type, custom_id))
return self._submit_job
async def poll(self, job_id):
self.polled.append(job_id)
return self._poll
def _wire_batch(monkeypatch, *, client, store, settings=None):
settings = settings or _settings(
document_ocr_mode="batch",
document_ocr_provider="gateway",
embedding_gateway_url="http://gw",
)
monkeypatch.setattr(ocr, "get_settings", lambda: settings)
monkeypatch.setattr(ocr, "build_gateway_batch_client", lambda s: client)
async def _shared(cls):
return store
monkeypatch.setattr(_bos.BatchOcrJobStore, "shared", classmethod(_shared))
async def test_batch_first_run_submits_and_returns_pending_sentinel(monkeypatch):
client = _FakeBatchClient()
store = _FakeStore()
_wire_batch(monkeypatch, client=client, store=store)
r = await ocr.OcrProcessor().process(
b"%PDF-1.7", "application/pdf", options=dict(_IDENTITY)
)
assert r.success is False
assert r.metadata[ocr.OCR_BATCH_PENDING_KEY] is True
assert r.metadata[ocr.OCR_BATCH_RETRY_IN_KEY] == 120
# submitted with the doc id as custom_id, recorded a pending row, swept stale
assert client.submitted and client.submitted[0][2] == "d1"
assert store.rows[("u1", "d1", "file", "v1")].job_id == "mistral/job-1"
assert store.stale_swept == [("u1", "d1", "file", "v1")]
async def test_batch_existing_pending_polls_and_defers(monkeypatch):
preset = SimpleNamespace(job_id="mistral/j", status="pending", submitted_at=1000)
client = _FakeBatchClient(poll=BatchPollResult(status="pending", pages=[]))
store = _FakeStore(preset=preset)
# submitted just now -> deadline not reached
monkeypatch.setattr(ocr.time, "time", lambda: 1000.0)
_wire_batch(monkeypatch, client=client, store=store)
r = await ocr.OcrProcessor().process(
b"%PDF", "application/pdf", options=dict(_IDENTITY)
)
assert client.polled == ["mistral/j"]
assert r.metadata[ocr.OCR_BATCH_PENDING_KEY] is True
assert client.submitted == [] # did NOT resubmit
async def test_batch_succeeded_returns_indexed_result(monkeypatch):
preset = SimpleNamespace(job_id="mistral/j", status="pending", submitted_at=1000)
client = _FakeBatchClient(
poll=BatchPollResult(status="succeeded", pages=[(0, "# One"), (1, "## Two")])
)
store = _FakeStore(preset=preset)
_wire_batch(monkeypatch, client=client, store=store)
r = await ocr.OcrProcessor().process(
b"%PDF", "application/pdf", options=dict(_IDENTITY)
)
assert r.success is True
assert r.text == "# One\n\n## Two"
assert r.metadata["page_count"] == 2
assert ("u1", "d1", "file", "v1") in store.deleted # row cleaned up
async def test_batch_failed_marks_parse_error(monkeypatch):
preset = SimpleNamespace(job_id="mistral/j", status="pending", submitted_at=1000)
client = _FakeBatchClient(
poll=BatchPollResult(status="failed", pages=[], error="x")
)
store = _FakeStore(preset=preset)
_wire_batch(monkeypatch, client=client, store=store)
r = await ocr.OcrProcessor().process(
b"%PDF", "application/pdf", options=dict(_IDENTITY)
)
assert r.success is False
assert r.metadata["parse_failed_reason"] == "error"
assert ("u1", "d1", "file", "v1") in store.deleted
async def test_batch_deadline_exceeded_marks_timeout(monkeypatch):
preset = SimpleNamespace(job_id="mistral/j", status="pending", submitted_at=1000)
client = _FakeBatchClient(poll=BatchPollResult(status="pending", pages=[]))
store = _FakeStore(preset=preset)
# now far past submitted_at + max_wait (86400)
monkeypatch.setattr(ocr.time, "time", lambda: 1000.0 + 90000)
_wire_batch(monkeypatch, client=client, store=store)
r = await ocr.OcrProcessor().process(
b"%PDF", "application/pdf", options=dict(_IDENTITY)
)
assert r.success is False
assert r.metadata["parse_failed_reason"] == "timeout"
assert ("u1", "d1", "file", "v1") in store.deleted
async def test_batch_falls_back_to_sync_when_no_gateway(monkeypatch):
class _FakeBackend:
async def ocr(self, content, mime_type):
return "sync text", [{"page": 1, "start_offset": 0, "end_offset": 9}]
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())
r = await ocr.OcrProcessor().process(
b"%PDF", "application/pdf", options=dict(_IDENTITY)
)
assert r.success is True and r.text == "sync text"
async def test_batch_falls_back_to_sync_when_no_identity(monkeypatch):
class _FakeBackend:
async def ocr(self, content, mime_type):
return "sync text", [{"page": 1, "start_offset": 0, "end_offset": 9}]
client = _FakeBatchClient()
settings = _settings(
document_ocr_mode="batch",
document_ocr_provider="gateway",
embedding_gateway_url="http://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())
# No options -> inline path -> batch inapplicable -> sync fallback.
r = await ocr.OcrProcessor().process(b"%PDF", "application/pdf", options=None)
assert r.success is True and r.text == "sync text"
assert client.submitted == [] # never attempted batch
+37
View File
@@ -97,3 +97,40 @@ async def test_hard_failure_returns_result_without_escalating(monkeypatch):
reg.evaluate_escalation.assert_not_called()
rec.assert_not_called()
sup.assert_not_called()
async def test_ocr_batch_pending_sentinel_raises_batch_pending():
"""Batch OCR (Deck #332): the OCR tier's pending sentinel result is turned
into a BatchPending raise (same decision point as EscalateError), carrying
the processor's retry_in, and the escalation gate is never consulted."""
from nextcloud_mcp_server.document_processors.escalation import BatchPending
from nextcloud_mcp_server.document_processors.ocr import (
OCR_BATCH_PENDING_KEY,
OCR_BATCH_RETRY_IN_KEY,
)
result = ProcessingResult(
text="",
metadata={OCR_BATCH_PENDING_KEY: True, OCR_BATCH_RETRY_IN_KEY: 90},
processor="ocr",
success=False,
)
reg = _registry(result, decision=None)
with pytest.raises(BatchPending) as ei:
await processor._parse_pdf_tier(
reg, b"%PDF", "application/pdf", "scan.pdf", "ocr", settings=object()
)
assert ei.value.retry_in == 90
reg.evaluate_escalation.assert_not_called()
async def test_options_threaded_to_process_tier():
"""The OCR identity options are forwarded to process_tier (batch needs them)."""
result = ProcessingResult(text="clean", metadata={}, processor="ocr")
reg = _registry(result, decision=None)
opts = {"user_id": "u", "doc_id": "d", "doc_type": "file", "etag": "v"}
await processor._parse_pdf_tier(
reg, b"%PDF", "application/pdf", "f.pdf", "ocr", settings=object(), options=opts
)
# process_tier(content, content_type, filename, tier, options=...)
assert reg.process_tier.await_args.kwargs["options"] == opts
@@ -14,6 +14,7 @@ from procrastinate.jobs import Job
import nextcloud_mcp_server.vector.queue.procrastinate as pq
from nextcloud_mcp_server.document_processors.escalation import (
TIER_LADDER,
BatchPending,
EscalateError,
next_tier,
)
@@ -120,3 +121,30 @@ class TestTieredEscalationStrategy:
exception=ValueError("permanent"), job=_job(attempts=1)
)
assert decision is None
def test_batch_pending_defers_same_queue(self):
# 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)
)
after = datetime.now(timezone.utc)
assert decision is not None
assert decision.queue is None # stays on its own tier queue
assert decision.retry_at is not None
lo = (decision.retry_at - after).total_seconds()
hi = (decision.retry_at - before).total_seconds()
assert lo <= 120 <= hi
def test_batch_pending_exempt_from_transient_cap(self):
# A batch can take hours -> many polls; the transient cap must NOT stop it
# (the OCR processor's own deadline terminates a stuck job instead).
decision = self._strategy(max_transient=5).get_retry_decision(
exception=BatchPending(retry_in=60), job=_job(attempts=999)
)
assert decision is not None and decision.retry_at is not None
def test_batch_pending_unwraps_exception_group(self):
group = ExceptionGroup("wrapped", [BatchPending(retry_in=60)])
decision = self._strategy().get_retry_decision(exception=group, job=_job())
assert decision is not None and decision.retry_at is not None