"""Tier-3 OCR processor. Routes scanned / no-text-layer PDFs (the tier-0 classifier's ``ocr`` verdict) to an OCR backend that returns per-page markdown. Two interchangeable backends, selected by ``document_ocr_provider``: * ``gateway`` -- POST the document to the Astrolabe Cloud model gateway's ``POST /v1/ocr`` (the same M2M-authenticated gateway as embeddings; NO provider keys in the pod). The platform default. * ``mistral`` -- call the Mistral OCR API directly from the pod (``MISTRAL_API_KEY``), for self-hosters / deployments without the gateway. ``auto`` prefers the gateway (if ``EMBEDDING_GATEWAY_URL`` is set) then direct Mistral (if ``MISTRAL_API_KEY``). Both return GitHub-flavoured markdown + exact ``page_boundaries``; bbox is re-derived from the PDF bytes + boundaries by ``search/pdf_highlighter``, as for the other tiers. """ import base64 import logging import time from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any import anyio import httpx from nextcloud_mcp_server.config import Settings, get_settings from .base import DocumentProcessor, ProcessingResult if TYPE_CHECKING: # Annotation-only import (the runtime import is lazy, inside # build_gateway_batch_client, to avoid a document_processors -> embedding # cycle at load). from ..embedding.gateway_batch_client import GatewayBatchOcrClient logger = logging.getLogger(__name__) # Connect timeout for the OCR backend request. The overall (read) timeout is # 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]], ) -> tuple[str, list[dict[str, Any]]]: """Join per-page markdown (ordered by index) into one string + boundaries. Pages are joined with a blank line. Boundaries are kept CONTIGUOUS (each page owns its leading ``\\n\\n`` separator) so they index exactly into the returned text and ``boundaries[-1]["end_offset"] == len(text)`` -- the ``search/pdf_highlighter`` contract. Consequence: a page's range starts at its separator, not its first glyph (the fast pypdfium2 path joins with no separator, so its ranges are glyph-tight). The 2-char offset is immaterial to page-level chunk attribution. """ sep = "\n\n" parts: list[str] = [] boundaries: list[dict[str, Any]] = [] offset = 0 for i, (index, markdown) in enumerate(sorted(pages, key=lambda p: p[0])): chunk = markdown if i == 0 else sep + markdown start = offset offset += len(chunk) parts.append(chunk) boundaries.append( {"page": index + 1, "start_offset": start, "end_offset": offset} ) 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( self, content: bytes, mime_type: str ) -> tuple[str, list[dict[str, Any]]]: ... class _GatewayOcrBackend(_OcrBackend): """Calls the model gateway's ``POST /v1/ocr`` (key-isolated, M2M-authed).""" def __init__(self, base_url: str, model: str, token_provider: Any = None): base = base_url.rstrip("/") if not base.endswith("/v1"): base = f"{base}/v1" self._url = f"{base}/ocr" self._model = model self._token_provider = token_provider async def ocr( self, content: bytes, mime_type: str ) -> tuple[str, list[dict[str, Any]]]: headers: dict[str, str] = {} if self._token_provider is not None: headers["Authorization"] = ( f"Bearer {await self._token_provider.get_token()}" ) payload = { "model": self._model, "document_b64": base64.b64encode(content).decode("ascii"), "mime_type": mime_type, } # Resolved per call (get_settings builds fresh) so test monkeypatching is # honoured; a live tenant change still needs a restart because the backend # instance itself is cached for the pod's lifetime. ocr_timeout = get_settings().document_ocr_timeout_seconds async with httpx.AsyncClient( timeout=httpx.Timeout(ocr_timeout, connect=_OCR_CONNECT_TIMEOUT_SECONDS) ) as client: resp = await client.post(self._url, json=payload, headers=headers) resp.raise_for_status() body = resp.json() pages = [(p["index"], p.get("markdown", "")) for p in body.get("pages", [])] return _pages_to_text(pages) class _MistralOcrBackend(_OcrBackend): """Calls the Mistral OCR API directly (provider key lives in the pod).""" def __init__(self, api_key: str, model: str, base_url: str | None = None): from mistralai.client import Mistral # noqa: PLC0415 -- lazy SDK import self._client = Mistral(api_key=api_key, server_url=base_url) # The gateway-namespaced "mistral/" id strips down to the bare # upstream model the SDK expects. self._model = model.split("/", 1)[-1] async def ocr( self, content: bytes, mime_type: str ) -> tuple[str, list[dict[str, Any]]]: data_url = ( f"data:{mime_type};base64,{base64.b64encode(content).decode('ascii')}" ) # Apply DOCUMENT_OCR_TIMEOUT_SECONDS uniformly with the gateway backend. # The Mistral SDK manages its own httpx client, so wrap the call in an # anyio cancel-scope timeout rather than threading a per-request timeout # through the SDK; on expiry this raises TimeoutError, which the # OcrProcessor turns into a clean parse failure. ocr_timeout = get_settings().document_ocr_timeout_seconds with anyio.fail_after(ocr_timeout): resp = await self._client.ocr.process_async( model=self._model, document={"type": "document_url", "document_url": data_url}, ) pages = [(p.index, p.markdown or "") for p in (resp.pages or [])] 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, *, model: str | None = None ) -> "GatewayBatchOcrClient | None": """Build a ``GatewayBatchOcrClient`` when the gateway is the OCR backend, else ``None`` (so batch mode falls back to sync for provider=mistral / no gateway). Batch OCR is gateway-only — Mistral's Batch API is reached *through* the gateway's batch routes, never directly from the pod. ``model`` overrides ``settings.document_ocr_model`` so a per-tier OCR rung (e.g. the in-cluster tier) submits its own provider-namespaced model id.""" if settings.document_ocr_provider not in ("gateway", "auto"): return None if not settings.embedding_gateway_url: return None from ..embedding.gateway_batch_client import GatewayBatchOcrClient # noqa: PLC0415 return GatewayBatchOcrClient( settings.embedding_gateway_url, model or settings.document_ocr_model, _build_gateway_token_provider(settings), ) def build_ocr_backend( settings: Settings, *, model: str | None = None, gateway_only: bool = False ) -> _OcrBackend | None: """Select an OCR backend from settings, or None when none is available. ``model`` overrides ``settings.document_ocr_model`` so a per-tier OCR rung binds its own provider-namespaced model id. ``gateway_only`` forces the gateway backend (never the direct Mistral fallback) — used by the **in-cluster** OCR tier, whose backend (e.g. surya on the burst GPU) is reachable ONLY through the embedding gateway over the tailnet; with no gateway URL the tier is disabled (warn) rather than misrouted to a direct backend that can't serve it. """ provider = settings.document_ocr_provider model = model or settings.document_ocr_model if provider == "none": return None if gateway_only: if settings.embedding_gateway_url: return _GatewayOcrBackend( settings.embedding_gateway_url, model, _build_gateway_token_provider(settings), ) logger.warning( "in-cluster OCR tier requires EMBEDDING_GATEWAY_URL (it routes through " "the gateway, never a direct backend); this OCR tier is disabled" ) return None if provider in ("gateway", "auto") and settings.embedding_gateway_url: return _GatewayOcrBackend( settings.embedding_gateway_url, model, _build_gateway_token_provider(settings), ) if provider in ("mistral", "auto") and settings.mistral_api_key: return _MistralOcrBackend( settings.mistral_api_key, model, settings.mistral_base_url, ) # An EXPLICIT provider that's missing its config is an operator error -- warn # loudly (once, since the backend is resolved+cached) rather than silently # disabling OCR. "auto"/"none" fall through to None quietly by design. if provider == "gateway": logger.warning( "DOCUMENT_OCR_PROVIDER=gateway but EMBEDDING_GATEWAY_URL is unset; " "OCR is disabled" ) elif provider == "mistral": logger.warning( "DOCUMENT_OCR_PROVIDER=mistral but MISTRAL_API_KEY is unset; " "OCR is disabled" ) return None class OcrProcessor(DocumentProcessor): """OCR processor for both OCR rungs — tier2 in-cluster (gateway-only GPU) and tier3 upstream (gateway or direct Mistral backend). One class, two registered instances bound to different ``(tier, model_setting, gateway_only)``.""" def __init__( self, *, name: str = "ocr-upstream", tier: str = "ocr-upstream", model_setting: str = "document_ocr_model", gateway_only: bool = False, ) -> None: # One OcrProcessor class serves BOTH OCR rungs; instances are bound to a # tier + the settings attribute holding their provider-namespaced model id # (+ whether the backend is gateway-only). The in-cluster rung is # gateway-only (its model, e.g. surya, is reachable solely via the # gateway); the upstream rung keeps the configurable gateway/mistral # selection. surya is NEVER hard-coded here — only a config default. # Fail fast on a misconfigured model_setting (a typo in a constructor call) # so it surfaces at startup, not as an AttributeError mid-OCR. The string # only ever comes from hardcoded defaults in __init__.py, never user input. if not hasattr(Settings, model_setting): raise ValueError(f"Unknown model_setting: {model_setting!r}") self._name = name self._tier = tier self._model_setting = model_setting self._gateway_only = gateway_only # Resolve the backend once and reuse it: rebuilding per call would create # a fresh GatewayTokenProvider each time (discarding its M2M-token cache # -> a token fetch per document) and a new Mistral SDK client per call. # A config change needs a pod restart anyway, so caching for the pod's # lifetime is safe. self._backend_resolved = False self._backend: _OcrBackend | None = None # Serialise first-call resolution so a burst of concurrent OCR requests # 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: GatewayBatchOcrClient | None = None self._batch_client_lock: anyio.Lock | None = None self._batch_fallback_warned = False @property def name(self) -> str: return self._name @property def tier(self) -> str: return self._tier @property def supported_mime_types(self) -> set[str]: return {"application/pdf"} async def process( self, content: bytes, content_type: str, filename: str | None = None, options: dict[str, Any] | None = None, progress_callback: ( Callable[[float, float | None, str | None], Awaitable[None]] | None ) = 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). # # A transport error from _process_batch (e.g. the gateway briefly down) # is intentionally NOT caught here: it propagates to procrastinate for a # durable retry rather than silently falling back to sync. If you've opted # into batch mode you want the retry, not an unexpected sync transcription. 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() async with self._backend_lock: if not self._backend_resolved: # double-checked self._backend = build_ocr_backend( settings, model=getattr(settings, self._model_setting), gateway_only=self._gateway_only, ) self._backend_resolved = True backend = self._backend if backend is None: logger.warning( "OCR requested for %s but no backend is configured (provider=%s)", filename or "", settings.document_ocr_provider, ) return ProcessingResult( text="", metadata={"parse_failed_reason": "unsupported"}, processor=self.name, success=False, error="no OCR backend configured", ) try: text, boundaries = await backend.ocr( content, content_type.split(";")[0].strip().lower() ) except (TimeoutError, httpx.TimeoutException): # Two timeout shapes reach here: the Mistral backend's # anyio.fail_after raises the builtin TimeoutError, while the gateway # backend's httpx.Timeout raises httpx.ReadTimeout (a # httpx.TimeoutException, NOT a TimeoutError). Catch both so a # too-low DOCUMENT_OCR_TIMEOUT_SECONDS lands in its own reason bucket # rather than being conflated with provider errors. timeout = settings.document_ocr_timeout_seconds logger.warning( "OCR timed out for %s after %.1fs", filename or "", timeout ) return ProcessingResult( text="", metadata={"parse_failed_reason": "timeout"}, processor=self.name, success=False, error=f"OCR timed out after {timeout:.1f}s", ) except Exception as e: logger.warning("OCR failed for %s: %s", filename or "", e) return ProcessingResult( text="", metadata={"parse_failed_reason": "error"}, processor=self.name, success=False, error=f"{type(e).__name__}: {e}", ) return ProcessingResult( text=text, metadata={ "page_count": len(boundaries), "page_boundaries": boundaries, "file_size": len(content), }, processor=self.name, ) async def _get_batch_client(self) -> "GatewayBatchOcrClient | None": """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. The in-cluster (``gateway_only``) rung never uses batch mode: it targets the on-demand GPU, which is synchronous/low-latency, while batch OCR is the upstream (Mistral) async-job path. So even with ``DOCUMENT_OCR_MODE=batch`` set globally, the in-cluster tier stays on the synchronous backend.""" if self._gateway_only: return None if not self._batch_client_resolved: if self._batch_client_lock is None: self._batch_client_lock = anyio.Lock() async with self._batch_client_lock: if not self._batch_client_resolved: # double-checked settings = get_settings() self._batch_client = build_gateway_batch_client( settings, model=getattr(settings, self._model_setting), ) 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: # The in-cluster (gateway_only) rung returns None here BY DESIGN — the # GPU is synchronous-only, batch is the upstream Mistral path — so don't # emit the "no gateway backend" warning (the gateway IS configured; that # warning would send an operator chasing a non-existent config problem). if not self._gateway_only: 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) — a no-op on the very first submit, # one cheap DELETE on a resubmit. 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 ) # We don't cancel the gateway-side job (there's no cancel endpoint # at this layer) — it keeps running and is reaped by the gateway's # own file purge. Dropping the row just stops us polling it. 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'}", ) if not result.is_succeeded: # Defensive: poll() maps anything that isn't "succeeded" to its raw # status, and only pending/succeeded/failed are handled above. An # unexpected terminal status (gateway version skew, a new lifecycle # state) must NOT fall through to _pages_to_text([]) -> a 0-chunk # "success" that silently indexes empty text and re-submits forever. logger.warning( "batch OCR job %s returned unexpected status %r; marking failed", job.job_id, result.status, ) return ProcessingResult( text="", metadata={"parse_failed_reason": "error"}, processor=self.name, success=False, error=f"unexpected batch status: {result.status}", ) 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 # processor reports healthy and surfaces a real failure per-document. return True