app.py

import json
import os
import subprocess
import tempfile

from fastapi import FastAPI, HTTPException, Request
from huggingface_hub import HfApi, login
import glob

app = FastAPI()
KEY = os.environ.get("WEBHOOK_SECRET")
HF_READ = os.environ.get("HF_READ")

api = HfApi(token=HF_READ)
login(HF_READ)

LATEST_PAYLOAD = "latest_payload.json"

def look_up_tflite(repo_id):
    with tempfile.TemporaryDirectory() as tmpdir:
        repo_url = f"https://huggingface.co/{repo_id}"
        subprocess.run(["git", "clone", repo_url, tmpdir], check=True)
        result = subprocess.run(["ls", "-l", tmpdir], capture_output=True, text=True)
        tflite_files = glob.glob(os.path.join(tmpdir, "*.tflite"))
        tflite_message = ""
        for tflite_path in tflite_files:
            tflite_file = os.path.basename(tflite_path)
            print(f"Found tflite file: {tflite_file}")
            tflite_message += f" ⚙️ Found tflite file **{tflite_file}** - Running cycle estimation\n"
        return f"```\n{result.stdout}\n```\n{tflite_message}\n(This is an automated message)"
    
def post_discussion_comment(repo_id, discussion_num, message):
    api.comment_discussion(
        repo_id=repo_id,
        discussion_num=discussion_num,
        comment=message
    )

@app.get("/")
def greet_json():
    try:
        with open(LATEST_PAYLOAD, "r") as file:
            payload = json.load(file)
    except FileNotFoundError:
        payload = {"message": "No payload received yet"}
    except json.JSONDecodeError:
        payload = {"message": "Error reading payload file"}
    return payload

@app.post("/benchmark")
async def handle_webhook(request: Request):
    secret = request.headers.get("X-Webhook-Secret")
    if secret != KEY:
        raise HTTPException(status_code=403, detail="Invalid webhook secret")

    try:
        payload = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON payload")

    with open(LATEST_PAYLOAD, "w") as file:
        json.dump(payload, file)

    # Check for "automatic model conversion succeeded" in the message
    message = payload.get("comment", {}).get("content", "")
    if "an automatic model conversion and evaluation has been performed and suceeded" in message:
        # Extract repo and discussion info from payload
        repo_id = payload.get("repo", {}).get("name", "gbahlnxp/mobilenetv1")
        if not repo_id:
            raise HTTPException(status_code=400, detail="Repository name missing in payload")
        discussion_num = payload.get("discussion", {}).get("num")
        if not discussion_num:
            raise HTTPException(status_code=400, detail="Discussion num missing in payload")

        try:
            ls_result = look_up_tflite(repo_id)
            post_discussion_comment(repo_id, discussion_num, f"❗A conversion was successful for Repo **{repo_id}** \n Starting cycle estimation \n Cloning the repo \n {ls_result}")
        except Exception as e:
            post_discussion_comment(repo_id, discussion_num, f"Error running ls: {e}")

    return {"status": "success", "received": payload}