yellow-bank-soal/backend/app/schemas/ai.py

"""
Pydantic schemas for AI generation endpoints.

Request/response models for admin AI generation playground.
"""

from typing import Dict, Literal, Optional

from pydantic import BaseModel, Field, field_validator

OPTION_LABELS = tuple("ABCDEFGHIJKLMNOPQRSTUVWXYZ")


class AIGeneratePreviewRequest(BaseModel):
    basis_item_id: int = Field(
        ..., description="ID of the basis item (must be sedang level)"
    )
    target_level: Literal["mudah", "sulit"] = Field(
        ..., description="Target difficulty level for generated question"
    )
    ai_model: str = Field(
        default="qwen/qwen2.5-32b-instruct",
        description="AI model to use for generation",
    )


class AIModelPricing(BaseModel):
    prompt: Optional[float] = Field(
        default=None, description="Input token price in USD per token"
    )
    completion: Optional[float] = Field(
        default=None, description="Output token price in USD per token"
    )
    prompt_per_million: Optional[float] = Field(
        default=None, description="Input token price in USD per 1M tokens"
    )
    completion_per_million: Optional[float] = Field(
        default=None, description="Output token price in USD per 1M tokens"
    )
    currency: str = "USD"
    source: str = "openrouter"


class AIUsageInfo(BaseModel):
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    total_tokens: Optional[int] = None
    cost_usd: Optional[float] = None


class AIGeneratePreviewResponse(BaseModel):
    success: bool = Field(..., description="Whether generation was successful")
    stem: Optional[str] = None
    options: Optional[Dict[str, str]] = None
    correct: Optional[str] = None
    explanation: Optional[str] = None
    ai_model: Optional[str] = None
    basis_item_id: Optional[int] = None
    target_level: Optional[str] = None
    usage: Optional[AIUsageInfo] = None
    error: Optional[str] = None
    cached: bool = False


class AISaveRequest(BaseModel):
    stem: str = Field(..., description="Question stem")
    options: Dict[str, str] = Field(
        ..., description="Answer options. Labels must match the basis item exactly."
    )
    correct: str = Field(..., description="Correct answer option label")
    explanation: Optional[str] = None
    tryout_id: str = Field(..., description="Tryout identifier")
    website_id: int = Field(..., description="Website identifier")
    basis_item_id: int = Field(..., description="Basis item ID")
    slot: int = Field(..., description="Question slot position")
    level: Literal["mudah", "sedang", "sulit"] = Field(
        ..., description="Difficulty level"
    )
    variant_status: Literal["active", "draft"] = Field(
        default="active",
        description="Lifecycle status for the saved variant. Workspace approvals save active variants.",
    )
    ai_model: str = Field(
        default="qwen/qwen2.5-32b-instruct",
        description="AI model used for generation",
    )

    @field_validator("correct")
    @classmethod
    def validate_correct(cls, v: str) -> str:
        label = v.upper()
        if label not in OPTION_LABELS:
            raise ValueError("Correct answer must be an option label A-Z")
        return label

    @field_validator("options")
    @classmethod
    def validate_options(cls, v: Dict[str, str]) -> Dict[str, str]:
        normalized = {
            str(key).strip().upper(): str(value).strip()
            for key, value in v.items()
            if str(key).strip() and str(value).strip()
        }
        if len(normalized) < 2:
            raise ValueError("Options must contain at least two non-empty choices")
        invalid_keys = sorted(set(normalized) - set(OPTION_LABELS))
        if invalid_keys:
            raise ValueError(f"Options contain invalid labels: {', '.join(invalid_keys)}")
        return normalized


class AISaveResponse(BaseModel):
    success: bool = Field(..., description="Whether save was successful")
    item_id: Optional[int] = None
    run_id: Optional[int] = None
    error: Optional[str] = None


class AIGenerateBatchRequest(BaseModel):
    basis_item_id: int = Field(
        ..., description="ID of the basis item (must be sedang level)"
    )
    target_level: Literal["mudah", "sulit"] = Field(
        ..., description="Target difficulty level for generated questions"
    )
    ai_model: str = Field(
        default="qwen/qwen2.5-32b-instruct",
        description="AI model to use for generation",
    )
    count: int = Field(default=3, ge=1, le=10, description="Number of variants to generate")
    operator_notes: Optional[str] = None


class AIBatchGeneratedItem(BaseModel):
    item_id: int
    stem: str
    options: Dict[str, str]
    correct: str
    explanation: Optional[str] = None
    level: str
    variant_status: str
    usage: Optional[AIUsageInfo] = None


class AIGenerateBatchResponse(BaseModel):
    success: bool
    run_id: Optional[int] = None
    item_ids: list[int] = Field(default_factory=list)
    items: list[AIBatchGeneratedItem] = Field(default_factory=list)
    generated_count: int = 0
    usage: Optional[AIUsageInfo] = None
    error: Optional[str] = None


class AIStatsResponse(BaseModel):
    total_ai_items: int = Field(..., description="Total AI-generated items")
    items_by_model: Dict[str, int] = Field(
        default_factory=dict, description="Items count by AI model"
    )
    cache_hit_rate: float = Field(
        default=0.0, description="Cache hit rate (0.0 to 1.0)"
    )
    total_cache_hits: int = Field(default=0, description="Total cache hits")
    total_requests: int = Field(default=0, description="Total generation requests")


class GeneratedQuestion(BaseModel):
    stem: str
    options: Dict[str, str]
    correct: str
    explanation: Optional[str] = None
    usage: Optional[AIUsageInfo] = None

    @field_validator("correct")
    @classmethod
    def validate_correct(cls, v: str) -> str:
        label = v.upper()
        if label not in OPTION_LABELS:
            raise ValueError("Correct answer must be an option label A-Z")
        return label