app.py

import os
import time
import traceback
from datetime import datetime, timezone

import gradio as gr
from huggingface_hub import HfApi


# ============================================================
# 基础配置
# ============================================================

# Hugging Face Space 通常会自动提供 SPACE_ID，例如：
# wenjiao/low_bit_open_llm_leaderboard_5
SPACE_REPO_ID = os.getenv("SPACE_ID")

# 如果 SPACE_ID 没拿到，就用你手动填的这个。
# 这里改成你的 Space 名字。
if not SPACE_REPO_ID:
    SPACE_REPO_ID = "wenjiao/timer"

# 你的 Space Notifications discussion 编号。
# 如果链接是 /discussions/3，这里就是 3。
DISCUSSION_NUM = int(os.getenv("DISCUSSION_NUM", "1"))

# 在 Space Settings -> Variables and secrets -> New secret 里添加：
# HF_TOKEN=hf_xxx
HF_TOKEN = os.getenv("HF_TOKEN")


# ============================================================
# 工具函数
# ============================================================

def get_hf_api():
    if not HF_TOKEN:
        raise RuntimeError(
            "没有检测到 HF_TOKEN。请在 Space Settings -> Variables and secrets 里添加 Secret：HF_TOKEN"
        )

    return HfApi(token=HF_TOKEN)


def clean_hf_username(hf_username):
    hf_username = str(hf_username or "").strip()
    hf_username = hf_username.lstrip("@")
    return hf_username


def notify_hf_user(hf_username, task_name, status, detail, result_url=None):
    hf_username = clean_hf_username(hf_username)

    if not hf_username:
        raise ValueError("Hugging Face 用户名不能为空。")

    api = get_hf_api()

    now = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")

    if status == "success":
        status_text = "✅ Space 任务已完成"
    else:
        status_text = "❌ Space 任务失败"

    comment_lines = [
        "@" + hf_username,
        "",
        status_text,
        "",
        "任务名称：" + str(task_name),
        "Space：" + str(SPACE_REPO_ID),
        "时间：" + now,
        "",
        str(detail),
    ]

    if result_url:
        comment_lines.extend([
            "",
            "结果链接：" + str(result_url),
        ])

    comment = "\n".join(comment_lines)

    api.comment_discussion(
        repo_id=SPACE_REPO_ID,
        repo_type="space",
        discussion_num=DISCUSSION_NUM,
        comment=comment,
    )

    discussion_url = "https://huggingface.co/spaces/" + SPACE_REPO_ID + "/discussions/" + str(DISCUSSION_NUM)
    return discussion_url


# ============================================================
# 这里替换成你自己的真实任务
# ============================================================

def run_real_task(task_name, progress=gr.Progress()):
    """
    这里现在是模拟任务。
    你后面可以把 time.sleep 这几行换成你的模型推理、评测、排行榜更新等代码。
    """

    progress(0.1, desc="开始任务")
    time.sleep(1)

    progress(0.4, desc="正在运行任务")
    time.sleep(1)

    progress(0.7, desc="正在生成结果")
    time.sleep(1)

    progress(1.0, desc="任务完成")

    return "任务 " + str(task_name) + " 已经成功完成。"


# ============================================================
# Gradio 回调函数
# ============================================================

def run_and_notify(task_name, hf_username, result_url, progress=gr.Progress()):
    task_name = str(task_name or "").strip()
    hf_username = clean_hf_username(hf_username)
    result_url = str(result_url or "").strip()

    if not task_name:
        task_name = "未命名任务"

    if not hf_username:
        return "❌ 请输入 Hugging Face 用户名，例如 wenjiao，不要写邮箱。"

    if not result_url:
        result_url = "https://huggingface.co/spaces/" + SPACE_REPO_ID

    try:
        result_text = run_real_task(task_name, progress=progress)

        discussion_url = notify_hf_user(
            hf_username=hf_username,
            task_name=task_name,
            status="success",
            detail=result_text,
            result_url=result_url,
        )

        return (
            "✅ 任务完成，并已尝试通过 Hugging Face @ 通知 @" + hf_username + "\n\n"
            + "任务结果：\n"
            + result_text
            + "\n\n通知位置：\n"
            + discussion_url
        )

    except Exception as e:
        error_text = traceback.format_exc()

        try:
            discussion_url = notify_hf_user(
                hf_username=hf_username,
                task_name=task_name,
                status="failed",
                detail="错误信息：" + str(e),
                result_url=result_url,
            )

            return (
                "❌ 任务失败，但已尝试通过 Hugging Face @ 通知 @" + hf_username + "\n\n"
                + "错误信息：\n"
                + str(e)
                + "\n\n通知位置：\n"
                + discussion_url
                + "\n\n完整错误：\n"
                + error_text
            )

        except Exception as notify_error:
            return (
                "❌ 任务失败，并且发送 Hugging Face 通知也失败了。\n\n"
                + "任务错误：\n"
                + str(e)
                + "\n\n通知错误：\n"
                + str(notify_error)
                + "\n\n完整错误：\n"
                + error_text
            )


# ============================================================
# Gradio 页面
# ============================================================

description_text = (
    "# Hugging Face Space 原生通知测试\n\n"
    + "点击按钮后，Space 会模拟运行一个任务。任务结束后，会自动往固定 Discussion 里评论，并 @ 目标用户。\n\n"
    + "当前 Space Repo ID：`" + str(SPACE_REPO_ID) + "`\n\n"
    + "当前 Discussion Num：`" + str(DISCUSSION_NUM) + "`\n\n"
    + "请确认你已经在 Space Settings 里添加 Secret：`HF_TOKEN=你的 Hugging Face token`。\n"
)

with gr.Blocks(title="HF Space Notification Demo") as demo:
    gr.Markdown(description_text)

    with gr.Row():
        task_name_input = gr.Textbox(
            label="任务名称",
            value="test-eval-job",
            placeholder="例如：llm-eval-run-001",
        )

        hf_username_input = gr.Textbox(
            label="要通知的 Hugging Face 用户名",
            placeholder="例如：wenjiao，不要带 @，不要写邮箱",
        )

    result_url_input = gr.Textbox(
        label="结果链接，可不填",
        placeholder="https://huggingface.co/spaces/" + str(SPACE_REPO_ID),
    )

    run_button = gr.Button("运行任务并触发 Hugging Face 通知", variant="primary")

    output = gr.Textbox(
        label="运行状态",
        lines=12,
    )

    run_button.click(
        fn=run_and_notify,
        inputs=[task_name_input, hf_username_input, result_url_input],
        outputs=output,
    )


demo.queue()
demo.launch()