services/ocr/contracts.py

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field, model_validator

from .config import ProviderName, default_provider_name, enabled_provider_names


class ImageSource(BaseModel):
    path: str | None = None
    base64: str | None = None

    @model_validator(mode="after")
    def validate_single_source(self) -> "ImageSource":
        supplied = int(self.path is not None) + int(self.base64 is not None)
        if supplied != 1:
            raise ValueError("image.path 和 image.base64 必须且只能提供一个")
        return self


class RecognizePayload(BaseModel):
    image: ImageSource
    provider: ProviderName = Field(default_factory=default_provider_name)
    languages: list[str] | None = None
    min_confidence: float = Field(default=0.0, ge=0.0, le=1.0)
    options: dict[str, Any] = Field(default_factory=dict)


class ComparePayload(BaseModel):
    image: ImageSource
    providers: list[ProviderName] = Field(default_factory=enabled_provider_names)
    languages: list[str] | None = None
    min_confidence: float = Field(default=0.0, ge=0.0, le=1.0)
    options: dict[str, dict[str, Any]] = Field(default_factory=dict)


class WarmupPayload(BaseModel):
    providers: list[ProviderName] = Field(default_factory=enabled_provider_names)
    languages: dict[str, list[str]] = Field(default_factory=dict)


class InvokeRequest(BaseModel):
    operation: Literal[
        "ocr.recognize", "ocr.compare", "providers.status", "providers.warmup"
    ]
    payload: dict[str, Any] = Field(default_factory=dict)


class OCRLine(BaseModel):
    text: str
    confidence: float | None = None
    quad: list[list[float]]
    bbox: dict[str, float]


class OCRResult(BaseModel):
    provider: ProviderName
    text: str
    lines: list[OCRLine]
    image_size: dict[str, int]
    elapsed_ms: float
    metadata: dict[str, Any] = Field(default_factory=dict)