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/embedding/gateway_batch_client.py
+171
View File
@@ -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)