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