src/backend/routes/hf.py

import io
from dataclasses import asdict
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, BinaryIO, Literal

from fastapi import APIRouter, Depends
from loguru import logger

from src.backend.config import settings
from src.backend.schemas import (
    CommitInfo,
    CommunitySubmit_Params,
    GetModelInfo_QueryParams,
    GetModelInfo_RespData,
    HfRepoUrl,
    ResponseData,
    UploadFileContent_Params,
    UploadFileContent_RespData,
)

if TYPE_CHECKING:
    from huggingface_hub import ModelInfo


router = APIRouter(tags=["huggingface"])


@router.get("/models/info/")
async def get_model_info(
    params: Annotated[GetModelInfo_QueryParams, Depends()],
) -> ResponseData[GetModelInfo_RespData]:
    """Get model info with commit history."""
    model: ModelInfo = settings.hf_api.model_info(params.model_id, revision=params.revision or None)
    # Get model commit history
    commit_infos = settings.hf_api.list_repo_commits(repo_id=params.model_id, repo_type="model")
    commits = [CommitInfo.model_validate(asdict(c)) for c in commit_infos]
    # Response data
    data = GetModelInfo_RespData.model_validate({
        **asdict(model),
        "commits": commits or None,
    })
    return ResponseData(data=data)


@router.post("/upload-file/")
async def upload_file_content(
    params: Annotated[UploadFileContent_Params, Depends()],
) -> ResponseData[UploadFileContent_RespData]:
    """Submit a new evaluation request to the Hugging Face repository."""
    file_obj = io.BytesIO(params.content.encode("utf-8"))
    data = await upload_file_content_handler(
        path_or_fileobj=file_obj,
        path_in_repo=params.path_in_repo,
        commit_message=params.commit_message,
        repo_id=settings.QUEUE_REPO_ID,
        repo_type="dataset",
    )
    return ResponseData(data=data)


async def upload_file_content_handler(
    path_or_fileobj: str | Path | bytes | BinaryIO,
    path_in_repo: str,
    repo_id: str,
    repo_type: Literal["model", "dataset", "space"] = "dataset",
    commit_message: str | None = None,
):
    """Community submit handler."""

    # FIXME: debug
    import rich

    rich.print({
        "path_or_fileobj": path_or_fileobj,
        "path_in_repo": path_in_repo,
        "repo_id": repo_id,
        "repo_type": repo_type,
        "commit_message": commit_message,
    })
    raise RuntimeError("debug")

    commit_info = settings.hf_api.upload_file(
        path_or_fileobj=path_or_fileobj,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type=repo_type,
        commit_message=commit_message,
    )
    data_dict = asdict(commit_info)
    try:
        repo_url = HfRepoUrl.model_validate({
            "endpoint": commit_info.repo_url.endpoint,
            "namespace": commit_info.repo_url.namespace,
            "repo_name": commit_info.repo_url.repo_name,
            "repo_id": commit_info.repo_url.repo_id,
            "repo_type": commit_info.repo_url.repo_type,
            "url": commit_info.repo_url.url,
        })
        return UploadFileContent_RespData.model_validate({
            **data_dict,
            "repo_url": repo_url,
        })
    except Exception as e:
        msg = f"Failed to validate repo url: {e}"
        logger.warning(msg)
        return UploadFileContent_RespData.model_validate(data_dict)


@router.post("/community/submit/")
async def community_submit(
    params: CommunitySubmit_Params,
) -> ResponseData[UploadFileContent_RespData]:
    """Submit a new evaluation request to the Hugging Face repository."""
    file_obj = io.BytesIO(params.content.encode("utf-8"))
    path_in_repo = f"leaderboard-submissions/{params.filename}"
    data = await upload_file_content_handler(
        path_or_fileobj=file_obj,
        path_in_repo=path_in_repo,
        commit_message=params.commit_message,
        repo_id=settings.QUEUE_REPO_ID,
        repo_type="dataset",
    )
    return ResponseData(data=data)