app.py

import os
from pathlib import Path
from typing import Literal

from fastapi import BackgroundTasks, HTTPException, Response, status
from huggingface_hub import (
    CommitOperationAdd,
    CommitOperationDelete,
    comment_discussion,
    create_commit,
    create_repo,
    delete_repo,
    get_repo_discussions,
    snapshot_download,
    space_info,
)
from huggingface_hub.repocard import RepoCard
from requests import HTTPError
from gradio_webhooks import GradioWebhookApp, WebhookPayload
from huggingface_hub import login
from ui import generate_ui
from database import is_space_registered

login(token=os.getenv("HF_TOKEN"))

CI_BOT_NAME = "spaces-ci-bot"

app = GradioWebhookApp(ui=generate_ui())


@app.add_webhook("/webhook")
async def post_webhook(payload: WebhookPayload, task_queue: BackgroundTasks):
    if payload.repo.type != "space":
        raise HTTPException(400, f"Must be a Space, not {payload.repo.type}")

    space_id = payload.repo.name
    if not is_space_registered(space_id):
        return "Space not in the watchlist."

    has_task = False
    if (
        # Means "a new PR has been opened"
        payload.event.scope.startswith("discussion")
        and payload.event.action == "create"
        and payload.discussion is not None
        and payload.discussion.isPullRequest
        and payload.discussion.status == "open"
    ):
        if not is_pr_synced(space_id=space_id, pr_num=payload.discussion.num):
            # New PR! Sync task scheduled
            task_queue.add_task(
                sync_ci_space,
                space_id=space_id,
                pr_num=payload.discussion.num,
                private=payload.repo.private,
            )
            has_task = True
    elif (
        # Means "a PR has been merged or closed"
        payload.event.scope.startswith("discussion")
        and payload.event.action == "update"
        and payload.discussion is not None
        and payload.discussion.isPullRequest
        and (
            payload.discussion.status == "merged"
            or payload.discussion.status == "closed"
        )
    ):
        task_queue.add_task(
            delete_ci_space,
            space_id=space_id,
            pr_num=payload.discussion.num,
        )
        has_task = True
    elif (
        # Means "some content has been pushed to the Space" (any branch)
        payload.event.scope.startswith("repo.content")
        and payload.event.action == "update"
    ):
        # New repo change. Is it a commit on a PR?
        # => loop through all PRs and check if new changes happened
        for discussion in get_repo_discussions(repo_id=space_id, repo_type="space"):
            if discussion.is_pull_request and discussion.status == "open":
                if not is_pr_synced(space_id=space_id, pr_num=discussion.num):
                    # Found a PR that is not yet synced
                    task_queue.add_task(
                        sync_ci_space,
                        space_id=space_id,
                        pr_num=discussion.num,
                        private=payload.repo.private,
                    )
                    has_task = True

    if has_task:
        return Response(
            "Task scheduled to sync/delete Space", status_code=status.HTTP_202_ACCEPTED
        )
    else:
        return Response("No task scheduled", status_code=status.HTTP_202_ACCEPTED)


def is_pr_synced(space_id: str, pr_num: int) -> bool:
    # What is the last synced commit for this PR?
    ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
    try:
        card = RepoCard.load(repo_id_or_path=ci_space_id, repo_type="space")
        last_synced_sha = getattr(card.data, "synced_sha", None)
    except HTTPError:
        last_synced_sha = None

    # What is the last commit id for this PR?
    info = space_info(repo_id=space_id, revision=f"refs/pr/{pr_num}")
    last_pr_sha = info.sha

    # Is it up to date ?
    return last_synced_sha == last_pr_sha


def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
    # Create a temporary space for CI if didn't exist
    ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)

    try:
        create_repo(
            ci_space_id,
            repo_type="space",
            space_sdk="docker",
            private=private,
        )
        is_new = True
    except HTTPError as err:
        if err.response.status_code == 409:  # already exists
            is_new = False
        else:
            raise

    # Download space codebase from PR revision
    snapshot_path = Path(
        snapshot_download(
            repo_id=space_id, revision=f"refs/pr/{pr_num}", repo_type="space"
        )
    )

    # Sync space codebase with PR revision
    operations = [  # little aggressive but works
        CommitOperationDelete(".", is_folder=True)
    ]
    for filepath in snapshot_path.glob("**/*"):
        if filepath.is_file():
            path_in_repo = str(filepath.relative_to(snapshot_path))

            # Upload all files without changes except for the README file
            if path_in_repo == "README.md":
                card = RepoCard.load(filepath)
                setattr(card.data, "synced_sha", snapshot_path.name)  # latest sha
                path_or_fileobj = str(card).encode()
            else:
                path_or_fileobj = filepath

            operations.append(
                CommitOperationAdd(
                    path_in_repo=path_in_repo, path_or_fileobj=path_or_fileobj
                )
            )

    create_commit(
        repo_id=ci_space_id,
        repo_type="space",
        operations=operations,
        commit_message=f"Sync CI Space with PR {pr_num}.",
    )

    # Post a comment on the PR
    notify_pr(space_id=space_id, pr_num=pr_num, action="create" if is_new else "update")


def delete_ci_space(space_id: str, pr_num: int) -> None:
    # Delete
    ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
    delete_repo(repo_id=ci_space_id, repo_type="space")

    # Notify about deletion
    notify_pr(space_id=space_id, pr_num=pr_num, action="delete")


def notify_pr(
    space_id: str, pr_num: int, action: Literal["create", "update", "delete"]
) -> None:
    ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)
    if action == "create":
        comment = NOTIFICATION_TEMPLATE_CREATE.format(ci_space_id=ci_space_id)
    elif action == "update":
        comment = NOTIFICATION_TEMPLATE_UPDATE.format(ci_space_id=ci_space_id)
    elif action == "delete":
        comment = NOTIFICATION_TEMPLATE_DELETE
    else:
        raise ValueError(f"Status {action} not handled.")

    comment_discussion(
        repo_id=space_id, repo_type="space", discussion_num=pr_num, comment=comment
    )


def _get_ci_space_id(space_id: str, pr_num: int) -> str:
    return f"{CI_BOT_NAME}/{space_id.replace('/', '-')}-ci-pr-{pr_num}"


NOTIFICATION_TEMPLATE_CREATE = """\
Hey there!
Following the creation of this PR, an ephemeral Space [{ci_space_id}](https://huggingface.co/spaces/{ci_space_id}) has been launched.
Any changes pushed to this PR will be synced with the test Space.

If your Space needs configuration (secrets or upgraded hardware), you must duplicate this ephemeral Space to your account and configure the settings by yourself.
You are responsible of making sure that the changes introduced in the PR are not harmful (leak secret, run malicious scripts,...)

(This is an automated message. To disable the Spaces CI Bot, please unregister using [this form](https://huggingface.co/spaces/spaces-ci-bot/webhook))
"""

NOTIFICATION_TEMPLATE_UPDATE = """\
Hey there!
Following new commits that happened in this PR, the ephemeral Space [{ci_space_id}](https://huggingface.co/spaces/{ci_space_id}) has been updated.

(This is an automated message. To disable the Spaces CI Bot, please unregister using [this form](https://huggingface.co/spaces/spaces-ci-bot/webhook))
"""

NOTIFICATION_TEMPLATE_DELETE = """\
Hey there!
PR is now merged/closed. The ephemeral Space has been deleted.

(This is an automated message. To disable the Spaces CI Bot, please unregister using [this form](https://huggingface.co/spaces/spaces-ci-bot/webhook))
"""


app.ready()