captcha_solver/models.py

"""Pydantic request/response schemas."""

from __future__ import annotations

from typing import Literal, Optional

from pydantic import BaseModel, Field, field_validator


CaptchaType = Literal["auto", "math", "audio", "text_ocr", "image_grid"]


class SolveRequest(BaseModel):
    """Request to solve a captcha."""

    type: CaptchaType = "auto"
    image_base64: Optional[str] = Field(
        default=None,
        description="Base64-encoded image of the captcha (PNG/JPG/WebP)",
    )
    audio_base64: Optional[str] = Field(
        default=None,
        description="Base64-encoded audio of the captcha (MP3/WAV/OGG)",
    )
    hint: Optional[str] = Field(
        default=None,
        description="Optional human hint to bias the solver (e.g. language code, category)",
    )
    timeout: int = Field(default=30, ge=1, le=120, description="Timeout in seconds")
    use_cache: bool = Field(default=True, description="Whether to use result cache")

    @field_validator("image_base64", "audio_base64")
    @classmethod
    def strip_data_url(cls, v: Optional[str]) -> Optional[str]:
        """Accept 'data:image/png;base64,XXX' or just 'XXX'."""
        if v is None:
            return v
        if "," in v and v.startswith("data:"):
            return v.split(",", 1)[1]
        return v.strip()


class SolveResponse(BaseModel):
    """Response with the captcha answer."""

    success: bool
    answer: Optional[str] = None
    confidence: float = Field(ge=0.0, le=1.0, default=0.0)
    solver: str = Field(description="Which solver produced the answer (e.g. 'math.regex')")
    elapsed_ms: int = 0
    cache_hit: bool = False
    attempts: int = 1
    error: Optional[str] = None
    metadata: dict = Field(default_factory=dict)


class HealthResponse(BaseModel):
    status: Literal["ok", "degraded", "down"]
    version: str
    engines: dict[str, str]
    uptime_s: float


class StatsResponse(BaseModel):
    total_requests: int
    by_type: dict[str, int]
    by_solver: dict[str, int]
    success_rate: float
    avg_elapsed_ms: float
    cache_hits: int


class ModelsResponse(BaseModel):
    loaded: list[str]
    available: list[dict]
    ollama_enabled: bool