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

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
nextcloud_mcp_server/document_processors/ocr.py
+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