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:
co-authored by
Claude Opus 4.8
parent
07ee91399b
commit
3b7e8d779b
@@ -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
|
of being handed to the tiers, where a 40+ MB scan would otherwise burn the full
|
||||||
OCR timeout for zero recovered text.
|
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 minutes–hours) 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
|
### Embedding Service Configuration
|
||||||
|
|
||||||
The server picks an embedding provider via auto-detection. Priority order
|
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")
|
||||||
@@ -175,6 +175,20 @@ _DEFAULTS: dict[str, Any] = {
|
|||||||
# 20-60s; raise/lower per tenant. Configurable so a tenant isn't stuck with
|
# 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.
|
# the 180s default when its gateway has its own shorter ceiling.
|
||||||
"document_ocr_timeout_seconds": 180.0,
|
"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
|
# Observability
|
||||||
"metrics_enabled": True,
|
"metrics_enabled": True,
|
||||||
"metrics_port": 9090,
|
"metrics_port": 9090,
|
||||||
@@ -352,6 +366,11 @@ _dynaconf = Dynaconf(
|
|||||||
Validator("DOCUMENT_CHUNK_SIZE", gte=1),
|
Validator("DOCUMENT_CHUNK_SIZE", gte=1),
|
||||||
Validator("DOCUMENT_PARSE_TIMEOUT_SECONDS", gte=1),
|
Validator("DOCUMENT_PARSE_TIMEOUT_SECONDS", gte=1),
|
||||||
Validator("DOCUMENT_OCR_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),
|
Validator("DOCUMENT_PARSE_MEM_LIMIT_MB", gte=128),
|
||||||
# 0 disables the pre-parse PDF size cap; otherwise it must be positive.
|
# 0 disables the pre-parse PDF size cap; otherwise it must be positive.
|
||||||
Validator("DOCUMENT_MAX_PDF_SIZE_MB", gte=0),
|
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
|
# parse timeout / httpx.Timeout; per-tenant tunable so a gateway with a
|
||||||
# shorter ceiling isn't masked by the 180s default.
|
# shorter ceiling isn't masked by the 180s default.
|
||||||
document_ocr_timeout_seconds: float = 180.0
|
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
|
# 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)
|
# 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)
|
# OR (when detect_scanned, image-analysis only runs when OCR is enabled)
|
||||||
|
|||||||
@@ -93,3 +93,26 @@ class EscalateError(Exception):
|
|||||||
super().__init__(
|
super().__init__(
|
||||||
f"escalate {from_tier}->{to_tier} (reason={reason})",
|
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)")
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ Mistral (if ``MISTRAL_API_KEY``). Both return GitHub-flavoured markdown + exact
|
|||||||
|
|
||||||
import base64
|
import base64
|
||||||
import logging
|
import logging
|
||||||
|
import time
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -35,6 +36,13 @@ logger = logging.getLogger(__name__)
|
|||||||
# configurable via DOCUMENT_OCR_TIMEOUT_SECONDS and resolved per call.
|
# configurable via DOCUMENT_OCR_TIMEOUT_SECONDS and resolved per call.
|
||||||
_OCR_CONNECT_TIMEOUT_SECONDS = 10.0
|
_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(
|
def _pages_to_text(
|
||||||
pages: list[tuple[int, str]],
|
pages: list[tuple[int, str]],
|
||||||
@@ -64,6 +72,24 @@ def _pages_to_text(
|
|||||||
return "".join(parts), boundaries
|
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):
|
class _OcrBackend(ABC):
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
async def ocr(
|
async def ocr(
|
||||||
@@ -141,6 +167,54 @@ class _MistralOcrBackend(_OcrBackend):
|
|||||||
return _pages_to_text(pages)
|
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:
|
def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
|
||||||
"""Select an OCR backend from settings, or None when none is available."""
|
"""Select an OCR backend from settings, or None when none is available."""
|
||||||
provider = settings.document_ocr_provider
|
provider = settings.document_ocr_provider
|
||||||
@@ -148,33 +222,10 @@ def build_ocr_backend(settings: Settings) -> _OcrBackend | None:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
if provider in ("gateway", "auto") and settings.embedding_gateway_url:
|
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(
|
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:
|
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:
|
# doesn't each build a backend (and fetch its own M2M token). Lazy-init:
|
||||||
# anyio primitives must not be created at import time.
|
# anyio primitives must not be created at import time.
|
||||||
self._backend_lock: anyio.Lock | None = None
|
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
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
@@ -239,6 +297,19 @@ class OcrProcessor(DocumentProcessor):
|
|||||||
) = None,
|
) = None,
|
||||||
) -> ProcessingResult:
|
) -> ProcessingResult:
|
||||||
settings = get_settings()
|
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 not self._backend_resolved:
|
||||||
if self._backend_lock is None:
|
if self._backend_lock is None:
|
||||||
self._backend_lock = anyio.Lock()
|
self._backend_lock = anyio.Lock()
|
||||||
@@ -301,6 +372,149 @@ class OcrProcessor(DocumentProcessor):
|
|||||||
processor=self.name,
|
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:
|
async def health_check(self) -> bool:
|
||||||
# Backends are resolved lazily (and configured per tenant), so there is
|
# Backends are resolved lazily (and configured per tenant), so there is
|
||||||
# nothing to probe here without making a billable upstream call -- the
|
# 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()
|
||||||
@@ -126,6 +126,7 @@ async def _parse_pdf_tier(
|
|||||||
filename: str | None,
|
filename: str | None,
|
||||||
tier: str,
|
tier: str,
|
||||||
settings: Any,
|
settings: Any,
|
||||||
|
options: dict[str, Any] | None = None,
|
||||||
) -> "ProcessingResult":
|
) -> "ProcessingResult":
|
||||||
"""Run a single extraction tier and apply the post-parse escalation gate.
|
"""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
|
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
|
not enabled a higher tier (or has no processor for it) simply indexes the
|
||||||
cheap tier's output.
|
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
|
# 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.
|
# load path; this runs only on the per-tier worker, which needs it anyway.
|
||||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||||
|
BatchPending,
|
||||||
EscalateError,
|
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
|
# ``options`` threads per-document identity (user_id/doc_id/doc_type/etag) to
|
||||||
# passes neither today, and the inline path (registry.process) omits them
|
# the OCR tier so batch mode can key its job-tracking table (Deck #332). Other
|
||||||
# too. Forward them if a tier processor ever needs per-call tuning (e.g. OCR
|
# tiers ignore it. The inline path (registry.process) passes None.
|
||||||
# DPI); keeping the two paths symmetric until then.
|
result = await registry.process_tier(
|
||||||
result = await registry.process_tier(content, content_type, filename, 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:
|
if result.success:
|
||||||
decision = registry.evaluate_escalation(
|
decision = registry.evaluate_escalation(
|
||||||
result, content, tier, settings, filename=filename
|
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
|
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.
|
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
|
# EscalateError and BatchPending are control-flow signals that arise ONLY on
|
||||||
# external path (tier set). Bind the class lazily here, and only when a tier
|
# the per-tier external path (tier set). Bind them lazily here, and only when a
|
||||||
# is set, so the document stack is never imported at *module load* (the #877
|
# tier is set, so the document stack is never imported at *module load* (the
|
||||||
# invariant) nor on the delete / text-doc call paths (file processing already
|
# #877 invariant) nor on the delete / text-doc call paths (file processing
|
||||||
# imports it via get_registry regardless). When tier is None it can't be
|
# already imports them via get_registry regardless). When tier is None neither
|
||||||
# raised, so the guards below stay inert.
|
# can be raised, so the guards below stay inert. Bound as a tuple so the guards
|
||||||
escalate_error_cls: type[BaseException] | None = None
|
# treat both identically: propagate untouched, never record an error/drop.
|
||||||
|
control_flow_excs: tuple[type[BaseException], ...] = ()
|
||||||
if tier is not None:
|
if tier is not None:
|
||||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||||
|
BatchPending,
|
||||||
EscalateError,
|
EscalateError,
|
||||||
)
|
)
|
||||||
|
|
||||||
escalate_error_cls = EscalateError
|
control_flow_excs = (EscalateError, BatchPending)
|
||||||
|
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
@@ -632,13 +648,11 @@ async def process_document(
|
|||||||
return # Success
|
return # Success
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
# An escalation signal is control flow, not a failure:
|
# A control-flow signal (escalation hop, or batch-OCR re-poll
|
||||||
# propagate it untouched so the procrastinate retry strategy
|
# deferral) is not a failure: propagate it untouched so the
|
||||||
# can hop the job to the next tier's queue. Never retry it
|
# procrastinate retry strategy handles it. Never retry it
|
||||||
# in-process and never count it as a drop.
|
# in-process and never count it as a drop.
|
||||||
if escalate_error_cls is not None and isinstance(
|
if isinstance(e, control_flow_excs):
|
||||||
e, escalate_error_cls
|
|
||||||
):
|
|
||||||
raise
|
raise
|
||||||
if attempt < max_retries - 1:
|
if attempt < max_retries - 1:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
@@ -691,10 +705,11 @@ async def process_document(
|
|||||||
raise
|
raise
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
# An escalation signal must reach the procrastinate retry strategy
|
# A control-flow signal must reach the procrastinate retry strategy
|
||||||
# un-recorded -- it is neither a processing success nor an error
|
# un-recorded -- it is neither a processing success nor an error (an
|
||||||
# (the hop is its own event, counted via record_document_escalation).
|
# escalation hop is counted via record_document_escalation; a batch
|
||||||
if escalate_error_cls is not None and isinstance(e, escalate_error_cls):
|
# re-poll deferral is not an event at all).
|
||||||
|
if isinstance(e, control_flow_excs):
|
||||||
raise
|
raise
|
||||||
# Single processing-error call site: catches exhausted-retry
|
# Single processing-error call site: catches exhausted-retry
|
||||||
# re-raises, delete failures, and setup errors (get_qdrant_client /
|
# re-raises, delete failures, and setup errors (get_qdrant_client /
|
||||||
@@ -956,6 +971,7 @@ async def _index_document(
|
|||||||
get_registry,
|
get_registry,
|
||||||
)
|
)
|
||||||
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
from nextcloud_mcp_server.document_processors.escalation import ( # noqa: PLC0415
|
||||||
|
BatchPending,
|
||||||
EscalateError,
|
EscalateError,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -968,6 +984,15 @@ async def _index_document(
|
|||||||
# and the in-process/memory pool (tier is None) -- runs the inline
|
# and the in-process/memory pool (tier is None) -- runs the inline
|
||||||
# tiered pipeline (fast -> OCR escalation in one call).
|
# tiered pipeline (fast -> OCR escalation in one call).
|
||||||
if tier is not None and _is_pdf(content_type):
|
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(
|
result = await _parse_pdf_tier(
|
||||||
registry,
|
registry,
|
||||||
content_bytes,
|
content_bytes,
|
||||||
@@ -975,6 +1000,7 @@ async def _index_document(
|
|||||||
file_path,
|
file_path,
|
||||||
tier,
|
tier,
|
||||||
settings,
|
settings,
|
||||||
|
options=ocr_options,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
result = await registry.process(
|
result = await registry.process(
|
||||||
@@ -1047,10 +1073,12 @@ async def _index_document(
|
|||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
logger.debug("No page_boundaries in metadata for %s", file_path)
|
logger.debug("No page_boundaries in metadata for %s", file_path)
|
||||||
except EscalateError:
|
except (EscalateError, BatchPending):
|
||||||
# Control-flow signal (per-tier path): re-raise untouched so the
|
# Control-flow signals (per-tier path): re-raise untouched.
|
||||||
# queue hops the job to the next tier. NOT a "failed to process"
|
# EscalateError hops the job to the next tier; BatchPending defers
|
||||||
# error -- don't log it as one.
|
# 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
|
raise
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error("Failed to process file %s: %s", file_path, e)
|
logger.error("Failed to process file %s: %s", file_path, e)
|
||||||
|
|||||||
@@ -328,13 +328,26 @@ class TieredEscalationStrategy(BaseRetryStrategy):
|
|||||||
def get_retry_decision(
|
def get_retry_decision(
|
||||||
self, *, exception: BaseException, job: Job
|
self, *, exception: BaseException, job: Job
|
||||||
) -> RetryDecision | None:
|
) -> RetryDecision | None:
|
||||||
# Lazy import: EscalateError lives in the document stack, which the API
|
# Lazy import: these live in the document stack, which the API pod (it
|
||||||
# pod (it also builds this App to defer) must not load. get_retry_decision
|
# also builds this App to defer) must not load. get_retry_decision runs
|
||||||
# runs only in the worker, where the stack is already imported.
|
# only in the worker, where the stack is already imported.
|
||||||
from ...document_processors.escalation import EscalateError # noqa: PLC0415
|
from ...document_processors.escalation import ( # noqa: PLC0415
|
||||||
|
BatchPending,
|
||||||
|
EscalateError,
|
||||||
|
)
|
||||||
|
|
||||||
exc = _first_leaf(exception)
|
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):
|
if isinstance(exc, EscalateError):
|
||||||
queue = TIER_QUEUES.get(exc.to_tier)
|
queue = TIER_QUEUES.get(exc.to_tier)
|
||||||
if queue is None:
|
if queue is None:
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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")
|
||||||
@@ -16,6 +16,9 @@ def _settings(**kw) -> Any: # a Settings stand-in (only the read fields matter)
|
|||||||
document_ocr_provider="auto",
|
document_ocr_provider="auto",
|
||||||
document_ocr_model="mistral/mistral-ocr-latest",
|
document_ocr_model="mistral/mistral-ocr-latest",
|
||||||
document_ocr_timeout_seconds=180.0,
|
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_url=None,
|
||||||
embedding_gateway_client_id=None,
|
embedding_gateway_client_id=None,
|
||||||
embedding_gateway_client_secret=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")
|
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||||
assert r.success is False
|
assert r.success is False
|
||||||
assert r.metadata["parse_failed_reason"] == "timeout"
|
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):
|
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")
|
r = await ocr.OcrProcessor().process(b"%PDF-1.7", "application/pdf")
|
||||||
assert r.success is False
|
assert r.success is False
|
||||||
assert r.metadata["parse_failed_reason"] == "timeout"
|
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):
|
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):
|
with pytest.raises(TimeoutError):
|
||||||
await backend.ocr(b"%PDF-1.7", "application/pdf")
|
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
|
||||||
|
|||||||
@@ -97,3 +97,40 @@ async def test_hard_failure_returns_result_without_escalating(monkeypatch):
|
|||||||
reg.evaluate_escalation.assert_not_called()
|
reg.evaluate_escalation.assert_not_called()
|
||||||
rec.assert_not_called()
|
rec.assert_not_called()
|
||||||
sup.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
|
import nextcloud_mcp_server.vector.queue.procrastinate as pq
|
||||||
from nextcloud_mcp_server.document_processors.escalation import (
|
from nextcloud_mcp_server.document_processors.escalation import (
|
||||||
TIER_LADDER,
|
TIER_LADDER,
|
||||||
|
BatchPending,
|
||||||
EscalateError,
|
EscalateError,
|
||||||
next_tier,
|
next_tier,
|
||||||
)
|
)
|
||||||
@@ -120,3 +121,30 @@ class TestTieredEscalationStrategy:
|
|||||||
exception=ValueError("permanent"), job=_job(attempts=1)
|
exception=ValueError("permanent"), job=_job(attempts=1)
|
||||||
)
|
)
|
||||||
assert decision is None
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user