schemas.py

from __future__ import annotations

from typing import Optional

from pydantic import BaseModel


DEFAULT_OCR_MODEL = "google/gemini-3-flash-preview"
DEFAULT_EXPERT_A = "google/gemini-3-flash-preview"
DEFAULT_EXPERT_B = "qwen/qwen3.5-flash-02-23"
DEFAULT_JUDGE = "mistralai/mistral-large-2512"

DEFAULT_LANGUAGE = ""
DEFAULT_GUIDELINES = (
    "Preserve the original orthography exactly as written; do not modernise spelling.\n"
    "Preserve line breaks: one transcribed line per physical line on the page.\n"
    "Keep abbreviation marks, ligatures, long-s (ſ), historical letters and special signs as written.\n"
    "Do not silently expand abbreviations.\n"
    "Preserve capitalisation and punctuation as on the source."
)
DEFAULT_OUTPUT_MODE = "lines"  # 'lines' | 'custom_json'
DEFAULT_JSON_TEMPLATE = ""  # only used when output_mode == 'custom_json'

CURATED_MODELS = [
    "google/gemini-3-flash-preview",
    "qwen/qwen3.5-flash-02-23",
    "mistralai/mistral-large-2512",
    "openai/gpt-5-mini",
    "anthropic/claude-opus-4-7",
    "anthropic/claude-sonnet-4-6",
]


class SettingsReq(BaseModel):
    ocr_model: str = DEFAULT_OCR_MODEL
    moe_expert_a: str = DEFAULT_EXPERT_A
    moe_expert_b: str = DEFAULT_EXPERT_B
    moe_judge: str = DEFAULT_JUDGE
    language: str = DEFAULT_LANGUAGE
    guidelines: str = DEFAULT_GUIDELINES
    output_mode: str = DEFAULT_OUTPUT_MODE
    output_json_template: str = DEFAULT_JSON_TEMPLATE
    temperature: float = 0.0
    n_icl: int = 2
    use_icl: bool = True
    custom_ocr_system: Optional[str] = None
    custom_ocr_user: Optional[str] = None
    custom_expert_system: Optional[str] = None
    custom_expert_user: Optional[str] = None
    custom_judge_system: Optional[str] = None
    custom_judge_user: Optional[str] = None


class OcrReq(BaseModel):
    page_idx: int


class PostCorrReq(BaseModel):
    image_id: Optional[str] = None
    ocr_text: str
    page_idx: Optional[int] = None


class CorrectionReq(BaseModel):
    page_idx: int
    corrected_text: str
    add_to_icl: bool = False


class PasteIclReq(BaseModel):
    image_id: str
    text: str
    language: str = ""


class TestKeyReq(BaseModel):
    model: Optional[str] = None


class SamPointReq(BaseModel):
    page_idx: int
    x: int                 # click position in natural image pixels
    y: int
    label: Optional[str] = None  # initial label (user can rename afterwards)


class AnnotationLabelReq(BaseModel):
    page_idx: int
    annotation_id: str
    label: str