huggingface cli requirement and webhook route

main.py

@@ -1,25 +1,48 @@
-from fastapi import FastAPI, Form, Request, BackgroundTasks
+import asyncio
+import os
+import time
+from typing import Optional
+from uuid import uuid4
+
+from fastapi import FastAPI, Form, Header, HTTPException, Request, BackgroundTasks
 from fastapi.responses import HTMLResponse
 from fastapi.templating import Jinja2Templates
-
-from uuid import uuid4
-import time
-import asyncio
+from huggingface_hub import create_discussion, comment_discussion
 
 from build_map import load_dataset_and_metadata, upload_dataset_to_atlas
-from models import WebhookPayload, WebhookPayloadRepo, WebhookPayloadEvent
+from models import WebhookPayload
 
+WEBHOOK_SECRET = os.environ.get("WEBHOOK_SECRET")
+HUGGINGFACE_ACCESS_TOKEN = os.environ.get("HUGGINGFACE_ACCESS_TOKEN")
 
 app = FastAPI()
 # TODO: use task management queue
 tasks = {}
 templates = Jinja2Templates(directory="templates")
 
-def upload_atlas_task(task_id, dataset_name):
+def upload_atlas_task(task_id, 
+                      dataset_name, 
+                      webhook_payload: WebhookPayload = None, 
+                      webhook_notify: bool = False):
     dataset_dict = load_dataset_and_metadata(dataset_name)
-    map_url = upload_dataset_to_atlas(dataset_dict, project_name="atlas-space-test")
+    map_url = upload_dataset_to_atlas(dataset_dict)
     tasks[task_id]['status'] = 'done'
     tasks[task_id]['url'] = map_url
+    tasks[task_id]['finish_time'] = time.time()
+
+    if webhook_notify:
+        discussion = create_discussion(
+            repo_id=webhook_payload.repo.id,
+            title="Atlas Maps",
+            token=HUGGINGFACE_ACCESS_TOKEN,
+        )
+        comment_discussion(
+            repo_id=webhook_payload.repo.id,
+            discussion_num=discussion.num,
+            comment="Atlas Map: " + map_url,
+            token=HUGGINGFACE_ACCESS_TOKEN
+        )
+
 
 @app.on_event("startup")
 async def startup_event():
@@ -47,7 +70,6 @@ async def form_post(background_tasks: BackgroundTasks, dataset_name: str = Form(
     tasks[task_id] = {'status': 'running'}
     #form_data = DatasetForm(dataset_name=dataset_name)
     background_tasks.add_task(upload_atlas_task, task_id, dataset_name)
-    
     return {'task_id': task_id}
 
 @app.get("/status/{task_id}")
@@ -58,7 +80,12 @@ async def read_task(task_id: str):
         return tasks[task_id]
 
 @app.post("/webhook")
-async def post_webhook(background_tasks: BackgroundTasks, payload: WebhookPayload):
+async def post_webhook(background_tasks: BackgroundTasks, payload: WebhookPayload, x_webhook_secret:  Optional[str] = Header(default=None)):
+    if x_webhook_secret is None:
+        raise HTTPException(401)
+    if x_webhook_secret != WEBHOOK_SECRET:
+        raise HTTPException(403)
+
     if not (
         payload.event.action == "update"
         and payload.event.scope.startswith("repo.content")
@@ -69,5 +96,5 @@ async def post_webhook(background_tasks: BackgroundTasks, payload: WebhookPayloa
         task_id = str(uuid4())
         tasks[task_id] = {'status': 'running'}
         #form_data = DatasetForm(dataset_name=dataset_name)
-        background_tasks.add_task(upload_atlas_task, task_id, payload.repo.name)
+        background_tasks.add_task(upload_atlas_task, task_id, payload.repo.name, payload, True)
         return {'task_id': task_id}

requirements.txt

@@ -1,5 +1,6 @@
 datasets==2.13.0
 fastapi[all]
+huggingface-hub==0.16.4
 nomic==2.0.3
 pandas==1.5.3
 pyarrow==12.0.1