from __future__ import annotations

from pydantic import BaseModel


class RankRequest(BaseModel):
    """
    Request body for POST /ai/rank.
    """

    trace_id: str
    speedups: list[float]
    diffs: list[str]
    optimization_ids: list[str]
    python_version: str | None = None
    function_references: str | None = None


class RankResponse(BaseModel):
    """
    Successful response from POST /ai/rank.
    """

    ranking: list[int]
    explanation: str
    scores: dict[int, dict[str, float]] | None = None


class RankErrorResponse(BaseModel):
    """
    Error response from POST /ai/rank.
    """

    error: str
