docqa/docxextract/schema.py

"""Typed request/response contracts and the error hierarchy."""

from __future__ import annotations

import enum
from typing import Any

from pydantic import BaseModel, Field, field_validator


class FieldStatus(str, enum.Enum):
    """Outcome of a single key extraction.

    `EXTRACTED` is only assigned above the confidence threshold; anything else
    is surfaced to the caller rather than being silently returned as fact.
    This distinction is the whole point of the service: the underlying model
    will always emit *some* span, so absence must be modelled explicitly.
    """

    EXTRACTED = "extracted"
    LOW_CONFIDENCE = "low_confidence"
    NOT_FOUND = "not_found"
    ERROR = "error"


class SourceType(str, enum.Enum):
    PDF_TEXT_LAYER = "pdf_text_layer"
    PDF_OCR = "pdf_ocr"
    IMAGE_OCR = "image_ocr"


class ExtractedField(BaseModel):
    key: str
    value: str | None = None
    status: FieldStatus = FieldStatus.NOT_FOUND
    confidence: float = Field(default=0.0, ge=0.0, le=1.0)
    page: int | None = None
    source: SourceType | None = None
    error: str | None = None

    @property
    def is_usable(self) -> bool:
        return self.status is FieldStatus.EXTRACTED


class DocumentInfo(BaseModel):
    document_id: str
    page_count: int
    pages_processed: int
    word_count: int
    source_type: SourceType
    truncated: bool = False


class ExtractRequest(BaseModel):
    """Request body. `file` arrives separately as an UploadFile."""

    keys: list[str] = Field(..., min_length=1)
    confidence_threshold: float | None = Field(default=None, ge=0.0, le=1.0)
    max_pages: int | None = Field(default=None, ge=1, le=50)

    @field_validator("keys")
    @classmethod
    def _validate_keys(cls, v: list[str]) -> list[str]:
        cleaned = [k.strip() for k in v if k and k.strip()]
        if not cleaned:
            raise ValueError("at least one non-empty key is required")
        if any(len(k) > 200 for k in cleaned):
            raise ValueError("each key must be 200 characters or fewer")
        return cleaned

    @field_validator("keys")
    @classmethod
    def _dedupe(cls, v: list[str]) -> list[str]:
        """Preserve order but drop exact duplicates.

        Repeating a key costs a full forward pass and returns an identical
        answer, so duplicates are collapsed rather than rejected.
        """
        seen: set[str] = set()
        out: list[str] = []
        for k in v:
            if k not in seen:
                seen.add(k)
                out.append(k)
        return out


class ExtractResponse(BaseModel):
    document: DocumentInfo
    fields: list[ExtractedField]
    latency_ms: float
    confidence_threshold: float

    @property
    def values(self) -> dict[str, str | None]:
        return {f.key: f.value for f in self.fields}


# --------------------------------------------------------------------------
# errors
# --------------------------------------------------------------------------
class DocExtractError(Exception):
    """Base class. `status_code` is what the HTTP layer will return."""

    status_code = 500

    def __init__(self, message: str, detail: Any = None):
        super().__init__(message)
        self.message = message
        self.detail = detail


class ValidationError(DocExtractError):
    status_code = 422


class UnsupportedMediaError(DocExtractError):
    status_code = 415


class PayloadTooLargeError(DocExtractError):
    status_code = 413


class CorruptDocumentError(DocExtractError):
    status_code = 422


class ExtractionTimeoutError(DocExtractError):
    status_code = 504


class ModelUnavailableError(DocExtractError):
    status_code = 503