231 lines
8.8 KiB
Python
231 lines
8.8 KiB
Python
"""Extract worker handler: download blob → extract/OCR → upload text → publish.
|
|
|
|
The handler is intentionally separate from the consumer so it can be tested
|
|
in-process without spinning up a real RabbitMQ consumer. It owns:
|
|
- idempotency checks against documents.status
|
|
- status transitions (queued → extracting)
|
|
- MinIO download/upload
|
|
- extraction/OCR via core.extraction
|
|
- publishing DocumentExtracted to analyze.q
|
|
- updating the jobs row, recording failure class, and refund-on-DLQ.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
from typing import TYPE_CHECKING
|
|
|
|
from src.contract_check.core.billing.quota import release_document_slot
|
|
from src.contract_check.core.config import get_settings
|
|
from src.contract_check.core.credits import refund_credit
|
|
from src.contract_check.core.db.enums import DOC_TERMINAL, FailureClass
|
|
from src.contract_check.core.db.repositories import DocumentRepository, JobRepository
|
|
from src.contract_check.core.extraction import (
|
|
ExtractionFailedError,
|
|
UnsupportedFormatError,
|
|
detect_format,
|
|
extract_document,
|
|
)
|
|
from src.contract_check.core.logging import get_logger
|
|
from src.contract_check.core.metrics import (
|
|
extract_duration,
|
|
extraction_total,
|
|
mq_failed,
|
|
mq_published,
|
|
)
|
|
from src.contract_check.core.mq.messages import (
|
|
DocumentExtracted,
|
|
DocumentUploaded,
|
|
PrescreenRequested,
|
|
)
|
|
from src.contract_check.core.mq.publisher import Publisher
|
|
from src.contract_check.core.mq.topology import RK_ANALYZE, RK_PRESCREEN
|
|
from src.contract_check.core.s3 import extracted_key
|
|
from src.contract_check.core.s3.minio_storage import MinioStorage
|
|
|
|
if TYPE_CHECKING:
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from src.contract_check.core.extraction import ExtractedDocument
|
|
|
|
log = get_logger(__name__)
|
|
|
|
|
|
class ExtractHandler:
|
|
"""Business logic for worker-extract."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
session_factory: async_sessionmaker[AsyncSession],
|
|
publish_routing_key: str = RK_ANALYZE,
|
|
) -> None:
|
|
self._session_factory = session_factory
|
|
self._publish_routing_key = publish_routing_key
|
|
self._settings = get_settings()
|
|
self._storage = MinioStorage.from_endpoint_url(
|
|
endpoint_url=self._settings.s3_endpoint_url,
|
|
access_key=self._settings.s3_access_key,
|
|
secret_key=self._settings.s3_secret_key,
|
|
bucket=self._settings.s3_bucket,
|
|
region=self._settings.s3_region,
|
|
)
|
|
self._publisher: Publisher | None = None
|
|
|
|
async def _publisher_instance(self) -> Publisher:
|
|
if self._publisher is None:
|
|
self._publisher = Publisher(self._settings.rabbitmq_url, origin="worker-extract")
|
|
await self._publisher.connect()
|
|
return self._publisher
|
|
|
|
async def _db_status(self, session: AsyncSession, document_id: uuid.UUID) -> str | None:
|
|
return await DocumentRepository(session).get_status_for_update(document_id)
|
|
|
|
async def handle(self, payload: DocumentUploaded) -> None:
|
|
async with self._session_factory() as session:
|
|
docs = DocumentRepository(session)
|
|
jobs = JobRepository(session)
|
|
current = await self._db_status(session, payload.document_id)
|
|
if current is None:
|
|
log.warning("document_not_found", document_id=str(payload.document_id))
|
|
return
|
|
if current in DOC_TERMINAL:
|
|
log.info("document_already_terminal", status=current)
|
|
return
|
|
|
|
await docs.update_status(payload.document_id, status="extracting", stage="downloading")
|
|
await jobs.claim_start(payload.document_id, "extract")
|
|
await session.commit()
|
|
|
|
document = await self._extract_document(payload)
|
|
extracted_text = document.markdown
|
|
fmt = document.metadata.get("format", "unknown")
|
|
ocr_used = document.metadata.get("is_scan", False)
|
|
is_structured = document.is_structured
|
|
has_tables = document.metadata.get("has_tables", False)
|
|
|
|
ext_key = extracted_key(str(payload.user_id), str(payload.document_id))
|
|
text_bytes = extracted_text.encode("utf-8")
|
|
|
|
await self._storage.put(ext_key, text_bytes, content_type="text/markdown; charset=utf-8")
|
|
|
|
publisher = await self._publisher_instance()
|
|
if self._settings.prescreen_enabled:
|
|
msg = PrescreenRequested(
|
|
correlation_id=payload.correlation_id,
|
|
document_id=payload.document_id,
|
|
user_id=payload.user_id,
|
|
text_s3_key=ext_key,
|
|
filename=payload.filename,
|
|
char_count=len(extracted_text),
|
|
is_structured=is_structured,
|
|
has_tables=has_tables,
|
|
attempt=payload.attempt,
|
|
)
|
|
publish_rk = RK_PRESCREEN
|
|
queue_metric = "prescreen"
|
|
doc_status = "prescreening"
|
|
doc_stage = "queued_prescreen"
|
|
else:
|
|
msg = DocumentExtracted(
|
|
correlation_id=payload.correlation_id,
|
|
document_id=payload.document_id,
|
|
user_id=payload.user_id,
|
|
extracted_s3_key=ext_key,
|
|
char_count=len(extracted_text),
|
|
ocr_used=ocr_used,
|
|
is_structured=is_structured,
|
|
has_tables=has_tables,
|
|
attempt=payload.attempt,
|
|
)
|
|
publish_rk = RK_ANALYZE
|
|
queue_metric = "analyze"
|
|
doc_status = "analyzing"
|
|
doc_stage = "queued_analyze"
|
|
|
|
await publisher.publish(msg, routing_key=publish_rk)
|
|
mq_published.labels(queue=queue_metric).inc()
|
|
|
|
async with self._session_factory() as session:
|
|
docs = DocumentRepository(session)
|
|
jobs = JobRepository(session)
|
|
await docs.update_status(
|
|
payload.document_id,
|
|
status=doc_status,
|
|
stage=doc_stage,
|
|
extracted_s3_key=ext_key,
|
|
)
|
|
await jobs.mark_done(payload.document_id, "extract")
|
|
await session.commit()
|
|
|
|
log.info(
|
|
"extract_success",
|
|
document_id=str(payload.document_id),
|
|
format=fmt,
|
|
structured=is_structured,
|
|
tables=has_tables,
|
|
char_count=len(extracted_text),
|
|
)
|
|
|
|
async def _extract_document(self, payload: DocumentUploaded) -> ExtractedDocument:
|
|
data = await self._storage.get(payload.s3_key)
|
|
fmt = detect_format(data, mime=payload.mime, filename=payload.filename)
|
|
|
|
def _extract() -> ExtractedDocument:
|
|
return extract_document(data, mime=payload.mime, filename=payload.filename)
|
|
|
|
with extract_duration.labels(format=fmt).time():
|
|
document = await asyncio.to_thread(_extract)
|
|
|
|
extraction_total.labels(format=fmt, structured=str(document.is_structured).lower()).inc()
|
|
return document
|
|
|
|
def classify(self, exc: BaseException) -> FailureClass:
|
|
from src.contract_check.core.analysis.ocr import OCRError
|
|
|
|
if isinstance(exc, (ExtractionFailedError, UnsupportedFormatError)):
|
|
return "extraction_failed"
|
|
if isinstance(exc, OCRError):
|
|
return "ocr_failed"
|
|
return "infra"
|
|
|
|
async def on_failure(
|
|
self, payload: DocumentUploaded, failure_class: FailureClass, attempt: int, error: str
|
|
) -> None:
|
|
mq_failed.labels(queue="extract", failure_class=failure_class).inc()
|
|
async with self._session_factory() as session:
|
|
await JobRepository(session).mark_retrying(
|
|
payload.document_id,
|
|
"extract",
|
|
attempt=attempt,
|
|
failure_class=failure_class,
|
|
error=error,
|
|
)
|
|
await session.commit()
|
|
|
|
async def on_terminal_failure(
|
|
self, payload: DocumentUploaded, failure_class: FailureClass, error: str
|
|
) -> None:
|
|
mq_failed.labels(queue="extract", failure_class=failure_class).inc()
|
|
async with self._session_factory() as session:
|
|
docs = DocumentRepository(session)
|
|
jobs = JobRepository(session)
|
|
await jobs.mark_dlq(
|
|
payload.document_id, "extract", failure_class=failure_class, error=error
|
|
)
|
|
await docs.mark_failed(payload.document_id, stage=failure_class)
|
|
await release_document_slot(session, payload.document_id)
|
|
await refund_credit(
|
|
session,
|
|
payload.document_id,
|
|
failure_class,
|
|
self._settings.refund_policy,
|
|
)
|
|
await session.commit()
|
|
|
|
log.warning(
|
|
"extract_terminal_failure",
|
|
document_id=str(payload.document_id),
|
|
failure_class=failure_class,
|
|
)
|