update.py

"""
This file has functions to update the meilisearch index with new comments.


Payload from HF webhooklooks like this:
    {
  "event": {
    "action": "update",
    "scope": "discussion.comment"
  },
  "repo": {
    "type": "dataset",
    "name": "allenai/objaverse",
    "id": "63977bb96bdef8095268ded0",
    "private": false,
    "url": {
      "web": "https://huggingface.co/datasets/allenai/objaverse",
      "api": "https://huggingface.co/api/datasets/allenai/objaverse"
    },
    "owner": {
      "id": "5e70f3648ce3c604d78fe132"
    }
  },
  "discussion": {
    "id": "66f1a1092eb1ea2422555d24",
    "title": "PullRequest",
    "url": {
      "web": "https://huggingface.co/datasets/allenai/objaverse/discussions/63",
      "api": "https://huggingface.co/api/datasets/allenai/objaverse/discussions/63"
    },
    "status": "draft",
    "author": {
      "id": "6673e848436907f83a815ab0"
    },
    "num": 63,
    "isPullRequest": true,
    "changes": {
      "base": "refs/heads/main"
    }
  },
  "comment": {
    "id": "66f1a1092eb1ea2422555d25",
    "author": {
      "id": "6673e848436907f83a815ab0"
    },
    "hidden": true,
    "url": {
      "web": "https://huggingface.co/datasets/allenai/objaverse/discussions/63#66f1a1092eb1ea2422555d25"
    }
  },
  "webhook": {
    "id": "66d7991f9b7da501cd100d95",
    "version": 3
  }
}
"""
import time
import json
import os
from datetime import datetime, timezone

import requests
from dotenv import load_dotenv
from huggingface_hub import HfApi
from meilisearch import Client
from huggingface_hub import HfApi

from constants import MeilisearchIndexFields

load_dotenv(".env", override=True)

WEBHOOK_SECRET = os.getenv("WEBHOOK_SECRET")
MEILISEARCH_URL = os.getenv("MS_URL")
MEILISEARCH_KEY = os.getenv("MS_ADMIN_KEY")

ms_client = Client(MEILISEARCH_URL, MEILISEARCH_KEY)

api = HfApi(token=os.environ["HF_WEBHOOK_TOKEN"])

async def process_webhook(request):

    payload = await request.body()
    payload = payload.decode("utf-8")
    print(payload)

    payload = json.loads(payload)

    secret = request.headers.get("X-Webhook-Secret")
    if secret != WEBHOOK_SECRET:
        print("Invalid secret")
        return {"error": "Invalid secret"}, 400

    if payload["repo"]["type"] == "model":

        if "discussion" not in payload or payload["discussion"]["isPullRequest"]:
            return {"status": "skipped"}, 200
        
        changing_status = "comment" not in payload and payload["event"]["action"] == "update"
        if changing_status:
            update_discussion_status(payload)
        else:
            add_new_comment(payload)


    return {"status": "success"}, 200


def user_id_to_username(user_id):

    api_url = f"https://huggingface.co/api/users/{user_id}/overview"

    try:
        response = requests.get(api_url)
        return response.json()["user"]
    except Exception as e:
        print(f"Couldn't get username for id {user_id}: {e}")
        return user_id
    
def add_new_comment(payload):
    
    
    comment = payload["comment"].get("content", "")
    comment_id = payload["comment"]["id"]

    repo_id = payload["repo"]["name"]
    title = payload["discussion"]["title"]

    author_id = payload["comment"]["author"]["id"]
    author = user_id_to_username(author_id)

    url = payload["discussion"]["url"]["web"]
    updatedAt = int(datetime.now(timezone.utc).timestamp())

    status = payload["discussion"]["status"]

    melisearch_payload = {
        MeilisearchIndexFields.ID.value: comment_id,
        MeilisearchIndexFields.TITLE.value: title,
        MeilisearchIndexFields.STATUS.value: status,
        MeilisearchIndexFields.AUTHOR.value: author,
        MeilisearchIndexFields.URL.value: url,
        MeilisearchIndexFields.REPO_ID.value: repo_id,
        MeilisearchIndexFields.CONTENT.value: comment,
        MeilisearchIndexFields.UPDATED_AT.value: updatedAt,
    }

    ms_client.index(MeilisearchIndexFields.INDEX_NAME.value).add_documents([melisearch_payload])

def update_discussion_status(payload):
    
    # If closing and commenting at the same time, 
    # the comment comes with status = open after the webhook that says the discussion is closed.
    # Adding the sleep ensures the update comes afterwards
    time.sleep(1)
    
    url = payload["discussion"]["url"]["web"]
    status = payload["discussion"]["status"]

    existing_results = ms_client.index(MeilisearchIndexFields.INDEX_NAME.value).search(
        query="",
        opt_params={"filter": f"url = '{url}'"}
    )

    if len(existing_results["hits"]) > 0:
        docs2update = [
            {MeilisearchIndexFields.ID.value: d[MeilisearchIndexFields.ID.value], MeilisearchIndexFields.STATUS.value: status}
            for d in existing_results["hits"]
        ]

        update_request = ms_client.index(MeilisearchIndexFields.INDEX_NAME.value).update_documents(docs2update)
        print("Update request:", update_request)


def is_user(user_or_org):
    api_url = f"https://huggingface.co/api/users/{user_or_org}/overview"
    response = requests.get(api_url)
    return response.status_code == 200


def update_webhooks():
    """
    Update the old webhook every so often with trending models.
    """

    print("Updating webhook")
    
    existing_webhooks = api.list_webhooks()

    webhook_url = os.environ["HF_WEBHOOK_URL"]

    webhook2update = [x for x in existing_webhooks if x.url == webhook_url]

    if len(webhook2update) > 1:
        print("More than one webhook found")
        print(webhook2update)
        print("updating the first one")

    id2update = webhook2update[0].id
    
    watch_dict = {}

    for ww in webhook2update[0].watched:
        watch_dict[ww.name] = ww.type

    # get trending models

    trending_models = api.list_models(sort="likes7d", direction=-1, limit=1000)

    to_add = []

    for m in trending_models:
      org_or_user = m.id.split("/")[0]
      if org_or_user in watch_dict:
          continue
      if is_user(org_or_user):
          to_add.append({"name": m.id, "type": "user"})
      else:
          to_add.append({"name": m.id, "type": "org"})
      
      new_watched = webhook2update[0].watched + to_add

      print("There are now", len(new_watched), "items in the watched list")

      api.update_webhook(
          id=id2update, 
          url=webhook_url,
          watched=new_watched,
          domains=["discussion"],
          secret=WEBHOOK_SECRET,
      )