refacto + add example

app.py

@@ -1,9 +1,8 @@
-# Taken from https://huggingface.co/spaces/huggingface-projects/auto-retrain
 import os
 from pathlib import Path
-from typing import Literal, Optional
+from typing import Literal
 
-from fastapi import BackgroundTasks, Header, HTTPException
+from fastapi import BackgroundTasks, HTTPException
 from huggingface_hub import (
     CommitOperationAdd,
     CommitOperationDelete,
@@ -16,55 +15,18 @@ from huggingface_hub import (
     space_info,
 )
 from huggingface_hub.repocard import RepoCard
-from pydantic import BaseModel
 from requests import HTTPError
 
-from gradio_webhooks import WebhookGradioApp
+from gradio_webhooks import GradioWebhookApp, WebhookPayload
 
-WEBHOOK_SECRET = os.getenv("WEBHOOK_SECRET")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 
-class WebhookPayloadEvent(BaseModel):
-    action: Literal["create", "update", "delete"]
-    scope: str
-
-
-class WebhookPayloadRepo(BaseModel):
-    type: Literal["dataset", "model", "space"]
-    name: str
-    private: bool
-
-
-class WebhookPayloadDiscussion(BaseModel):
-    num: int
-    isPullRequest: bool
-    status: Literal["open", "closed", "merged"]
-
-
-class WebhookPayload(BaseModel):
-    event: WebhookPayloadEvent
-    repo: WebhookPayloadRepo
-    discussion: Optional[WebhookPayloadDiscussion]
-
-
-app = WebhookGradioApp()
+app = GradioWebhookApp()
 
 
 @app.add_webhook("/webhook")
-async def post_webhook(
-    payload: WebhookPayload,
-    task_queue: BackgroundTasks,
-    x_webhook_secret: Optional[str] = Header(default=None),
-):
-    print("Received new hook!")
-    if x_webhook_secret is None:
-        print("HTTP 401: No webhook secret")
-        raise HTTPException(401)
-    if x_webhook_secret != WEBHOOK_SECRET:
-        print("HTTP 403: wrong webhook secret")
-        raise HTTPException(403)
-
+async def post_webhook(payload: WebhookPayload, task_queue: BackgroundTasks):
     if payload.repo.type != "space":
         print("HTTP 400: not a space")
         raise HTTPException(400, f"Must be a Space, not {payload.repo.type}")
@@ -273,4 +235,4 @@ PR is now merged/closed. The temporary test Space has been deleted.
 """
 
 
-app.block_thread()
+app.ready()

gradio_webhooks.py

@@ -1,15 +1,18 @@
+import os
 from pathlib import Path
-from typing import Set, Union
+from typing import Literal, Optional, Set, Union
 
 import gradio as gr
+from fastapi import HTTPException, Request
+from pydantic import BaseModel
 
 
-class WebhookGradioApp:
+class GradioWebhookApp:
     """
     ```py
-    from gradio_webhooks import WebhookGradioApp
+    from gradio_webhooks import GradioWebhookApp
 
-    app = WebhookGradioApp()
+    app = GradioWebhookApp()
 
 
     @app.add_webhook("/test_webhook")
@@ -17,11 +20,15 @@ class WebhookGradioApp:
         return {"in_gradio": True}
 
 
-    app.block_thread()
+    app.ready()
     ```
     """
 
-    def __init__(self, landing_path: Union[str, Path] = "README.md") -> None:
+    def __init__(
+        self,
+        landing_path: Union[str, Path] = "README.md",
+        webhook_secret: Optional[str] = None,
+    ) -> None:
         # Use README.md as landing page or provide any markdown file
         landing_path = Path(landing_path)
         landing_content = landing_path.read_text()
@@ -41,11 +48,28 @@ class WebhookGradioApp:
         self.fastapi_app = app
         self.webhook_paths: Set[str] = set()
 
+        # Add auth middleware to check the "X-Webhook-Secret" header
+        self._webhook_secret = webhook_secret or os.getenv("WEBHOOK_SECRET")
+        if self._webhook_secret is None:
+            print(
+                "\nWebhook secret is not defined. This means your webhook endpoints will be open to everyone."
+            )
+            print(
+                "To add a secret, set `WEBHOOK_SECRET` as environment variable or pass it at initialization: "
+                "\n\t`app = GradioWebhookApp(webhook_secret='my_secret', ...)`"
+            )
+            print(
+                "For more details about Webhook secrets, please refer to https://huggingface.co/docs/hub/webhooks#webhook-secret."
+            )
+        app.middleware("http")(self._webhook_secret_middleware)
+
     def add_webhook(self, path: str):
+        """Decorator to add a webhook to the server app."""
         self.webhook_paths.add(path)
-        return self.fastapi_app.post(path)
+        return self.fastapi_app.get(path)
 
-    def block_thread(self) -> None:
+    def ready(self) -> None:
+        """Set the app as "ready" and block main thread to keep it running."""
         url = (
             self.gradio_app.share_url
             if self.gradio_app.share_url is not None
@@ -55,3 +79,37 @@ class WebhookGradioApp:
         print("\n".join(f"  - POST {url}{webhook}" for webhook in self.webhook_paths))
         print(f"Checkout {url}/docs for more details.")
         self.gradio_app.block_thread()
+
+    async def _webhook_secret_middleware(self, request: Request, call_next) -> None:
+        """Middleware to check "X-Webhook-Secret" header on every webhook request."""
+        if request.url.path in self.webhook_paths:
+            if self._webhook_secret is not None:
+                request_secret = request.headers.get("x-webhook-secret")
+                if request_secret is None:
+                    raise HTTPException(401)
+                if request_secret != self._webhook_secret:
+                    raise HTTPException(403)
+        return await call_next(request)
+
+
+class WebhookPayloadEvent(BaseModel):
+    action: Literal["create", "update", "delete"]
+    scope: str
+
+
+class WebhookPayloadRepo(BaseModel):
+    type: Literal["dataset", "model", "space"]
+    name: str
+    private: bool
+
+
+class WebhookPayloadDiscussion(BaseModel):
+    num: int
+    isPullRequest: bool
+    status: Literal["open", "closed", "merged"]
+
+
+class WebhookPayload(BaseModel):
+    event: WebhookPayloadEvent
+    repo: WebhookPayloadRepo
+    discussion: Optional[WebhookPayloadDiscussion]

simple.py

@@ -0,0 +1,12 @@
+from gradio_webhooks import GradioWebhookApp, WebhookPayload
+
+app = GradioWebhookApp()
+
+
+@app.add_webhook("/my_webhook")
+async def hello(payload: WebhookPayload):
+    print(f"Received webhook for repo {payload.repo.name}")
+    return {"processed": True}
+
+
+app.ready()