src/backend/schemas.py

from datetime import datetime
from typing import Annotated, Any, Generic, Literal, TypeVar

from pydantic import AfterValidator, BaseModel, ConfigDict, Field, computed_field
from pydantic_core import PydanticCustomError

T = TypeVar("T", bound=BaseModel)


class ResponseData(BaseModel, Generic[T]):
    code: int = 0
    msg: str = "ok"
    data: T | None = None


# --- Get Model Info ---
class GetModelInfo_QueryParams(BaseModel):
    """Parameters for model info query."""

    model_id: str
    revision: str | None = None


class GetModelInfo_RespData(BaseModel):
    model_config = ConfigDict(extra="allow")

    id: str
    author: str | None = None
    downloads: int | None = None
    likes: int | None = None
    tags: list[str] | None = None
    pipeline_tag: str | None = None
    sha: str | None = None
    created_at: datetime | None = None
    last_modified: datetime | None = None
    private: bool | None = None
    disabled: bool | None = None
    gated: bool | None = None
    # ... others
    safetensors: "SafetensorsField | None" = None
    # commits
    commits: list["CommitInfo"] | None = None


class SafetensorsField(BaseModel):
    model_config = ConfigDict(extra="allow")

    parameters: dict[str, Any] | None = None
    total: int | None = None


class CommitInfo(BaseModel):
    model_config = ConfigDict(extra="allow")

    commit_id: str

    authors: list[str]
    created_at: datetime
    title: str
    message: str

    formatted_title: str | None = None
    formatted_message: str | None = None


# --- Upload File Content ---
class UploadFileContent_Params(BaseModel):
    path_in_repo: str
    commit_message: str | None = None

    model_id: str
    model_sha: str = "main"
    model_dtype: str

    content: str


class UploadFileContent_RespData(BaseModel):
    model_config = ConfigDict(extra="allow")

    commit_url: str
    commit_message: str
    commit_description: str
    oid: Annotated[str, Field(description="Commit hash id.")]
    repo_url: "HfRepoUrl | None" = None


class HfRepoUrl(BaseModel):
    model_config = ConfigDict(extra="allow")

    endpoint: str
    namespace: str
    repo_name: str
    repo_id: str
    repo_type: Literal["model", "dataset", "space"]
    url: str


# --- Community Submit ---
class CommunitySubmit_Params(BaseModel):
    username: Annotated[str, Field(description='The username of the user. e.g. "Qwen"')]

    # Model info
    model_id: Annotated[str, Field(description='The model id. e.g. "Qwen/Qwen2.5-3B"')]

    @computed_field
    @property
    def model_key(self) -> str:
        if not self.model_id:
            raise PydanticCustomError(
                "model_id_invalid", "Model id {model_id!r} is invalid.", {"model_id": self.model_id}
            )
        return self.model_id.lower().replace("-", "_").replace("/", "-")

    base_model: Annotated[
        str | None, Field(description='The base model name (for delta or adapter weights). e.g. "Qwen/Qwen2.5-3B"')
    ] = None
    model_sha: Annotated[
        str | None, Field(description='The model sha or "main". e.g. "3aab1f1954e9cc14eb9509a215f9e5ca08227a9b"')
    ] = None
    model_dtype: Literal[
        # float types
        "bfloat16",
        "float32",
        "float16",
        "float64",
        # int types
        "int8",
        "uint8",
        "int16",
        "uint15",
        "int32",
        "uint32",
        "int64",
        "uint64",
        "bool",
        # rare types
        "complex32",
        "complex64",
        "complex128",
        "float8_e4m3fn",
        "float8_e5m2",
        "float8_e4m3fnuz",
        "float8_e5m2fnuz",
        "float8_e8m0fnu",
        "float4_e2m1fn_x2",
    ]
    weight_type: Literal["Original", "Delta", "Adapter"] = "Original"
    model_type: Annotated[
        str,
        Field(description='The model type. e.g. "pretrained", "fine-tuned", "instruction-tuned", "RL-tuned"'),
    ]

    # Submission info
    content: Annotated[str, Field(description='The content of the file in JSON format to upload.')]
    submit_time: Annotated[
        str, Field(description='The submission time in UTC ISO format. e.g. "2025-01-15T10:30:00Z"')
    ]
    commit_message: str | None = None

    @computed_field
    @property
    def filename(self) -> str:
        """Filename of the file to upload."""
        if not self.model_key:
            raise PydanticCustomError(
                "model_key_invalid", "Model key {model_key!r} is invalid.", {"model_key": self.model_key}
            )
        if not self.username:
            raise PydanticCustomError(
                "username_invalid", "Username {username!r} is invalid.", {"username": self.username}
            )
        # "2025-01-15T10:30:00Z" -> "20250115T103000"
        submit_time = self.submit_time.replace(":", "").replace("-", "").rstrip("Z")
        return f"{submit_time}_{self.model_key}_{self.model_dtype}_{self.model_sha}_{self.username}.json"