Update main.py

main.py

@@ -8,9 +8,15 @@ import logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-HF_WEBHOOK_SECRET = os.environ.get("HF_WEBHOOK_SECRET")
+# Configurations
 SLACK_WEBHOOK_URL = os.environ.get("SLACK_WEBHOOK_URL")
 
+MODEL_CATALOG_WEBHOOK = os.environ.get("MODEL_CATALOG_WEBHOOK")
+MODEL_CATALOG_WEBHOOK_SECRET = os.environ.get("MODEL_CATALOG_WEBHOOK_SECRET")
+
+SIMSHIP_WEBHOOK = os.environ.get("SIMSHIP_WEBHOOK")
+SIMSHIP_WEBHOOK_SECRET = os.environ.get("SIMSHIP_WEBHOOK_SECRET")
+
 def send_slack_message(message: str):
     """Send a message to Slack using webhook URL"""
     if not SLACK_WEBHOOK_URL:
@@ -25,41 +31,92 @@ def send_slack_message(message: str):
     except requests.exceptions.RequestException as e:
         logger.error(f"Failed to send Slack message: {e}")
 
+def process_model_catalog_webhook(data: dict):
+    """
+    Process model catalog webhook
+    """
+
+    event = data.get("event", {})
+    repo = data.get("repo", {})
+    movedTo = data.get("movedTo", {})
+    
+    # repo name changed
+    if (
+        repo.get("type") == "model" and
+        event.get("scope") == "repo" and
+        event.get("action") == "move"
+    ):
+        message = f"🔄 Model in the catalog renamed: {repo.get('name', '')} → https://hf.co/{movedTo.get('name', '')}"
+        send_slack_message(message)
+
+    # repo deleted
+    elif (
+        repo.get("type") == "model" and
+        event.get("scope") == "repo" and
+        event.get("action") == "delete"
+    ):
+        message = f"🗑️ Model in the catalog deleted: https://hf.co/{repo.get('name', '')}"
+        send_slack_message(message)
+
+    # other events
+    else:
+        pass
+
+def process_simship_webhook(data: dict):
+    """
+    Process simship webhook
+    """
+    event = data.get("event", {})
+    repo = data.get("repo", {})
+    updatedConfig = data.get("updatedConfig", {})
+    
+    # repo creation
+    if (
+        repo.get("type") == "model" and
+        event.get("scope") == "repo" and
+        event.get("action") == "create"
+    ):
+        message = f"🆕 SimShip Model created: https://hf.co/{repo.get('name', '')}"
+        send_slack_message(message)
+
+    # repo visibility update
+    elif (
+        repo.get("type") == "model" and
+        event.get("scope") == "repo.config" and
+        event.get("action") == "update" and
+        updatedConfig.get("private") is False
+    ):
+        message = f"🆕 SimShip Model made public: https://hf.co/{repo.get('name', '')}"
+        send_slack_message(message)
+
+    # other events
+    else:
+        pass
+
+
 app = FastAPI()
 
 @app.post("/webhook")
 async def webhook(request: Request):
     logger.info(f"Received webhook request from {request.client.host}")
     
-    if request.headers.get("X-Webhook-Secret") != HF_WEBHOOK_SECRET:
+    if request.headers.get("X-Webhook-Secret") not in {MODEL_CATALOG_WEBHOOK_SECRET, SIMSHIP_WEBHOOK_SECRET}:
         logger.warning("Invalid webhook secret received")
         return Response("Invalid secret", status_code=401)
     
     data = await request.json()
     logger.info(f"Webhook payload: {json.dumps(data, indent=2)}")
+
+    webhook_id = data.get("webhook", {}).get("id", None)
     
-    event = data.get("event", {})
-    repo = data.get("repo", {})
-    
-    # Only process model repo events
-    if repo.get("type") != "model" or event.get("scope") != "repo":
-        logger.info(f"Skipping non-model repo event: type={repo.get('type')}, scope={event.get('scope')}")
-        return Response("Not a model repo event", status_code=200)
-    
-    action = event.get("action")
-    logger.info(f"Processing model repo action: {action}")
-    
-    if action == "move":
-        old_name = repo.get("name", "unknown")
-        new_name = data.get("movedTo", {}).get("name", "unknown")
-        message = f"🔄 Model renamed: {old_name} → https://hf.co/{new_name}"
-        send_slack_message(message)
-    elif action == "delete":
-        name = repo.get("name", "unknown")
-        message = f"🗑️ Model deleted: https://hf.co/{name}"
-        send_slack_message(message)
+    if webhook_id == MODEL_CATALOG_WEBHOOK:
+        logger.info("Processing model catalog webhook")
+        process_model_catalog_webhook(data)
+    elif webhook_id == SIMSHIP_WEBHOOK:
+        logger.info("Processing simship webhook")
+        process_simship_webhook(data)
     else:
-        logger.info(f"Unhandled action: {action}")
+        logger.warning("Invalid webhook ID received")
+        return Response("Invalid webhook ID", status_code=401)
     
     return Response("Webhook notification received and processed!", status_code=200)
-