update model_catalog watchlist logic

main.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, Request, Response
+from datasets import load_dataset
 import os
 import requests
 import json
@@ -12,8 +13,7 @@ logger = logging.getLogger(__name__)
 # 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")
+CATALOG_WEBHOOKS_FLEET_DATASET = os.environ.get("CATALOG_WEBHOOKS_FLEET_DATASET")
 
 SIMSHIP_WEBHOOK = os.environ.get("SIMSHIP_WEBHOOK")
 SIMSHIP_WEBHOOK_SECRET = os.environ.get("SIMSHIP_WEBHOOK_SECRET")
@@ -173,16 +173,23 @@ app = FastAPI()
 @app.post("/webhook/model-catalog")
 async def model_catalog_webhook(request: Request):
     logger.info(f"Received model catalog webhook request from {request.client.host}")
-    if request.headers.get("X-Webhook-Secret") != MODEL_CATALOG_WEBHOOK_SECRET:
+    catalog_webhooks_fleet_dataset = load_dataset(CATALOG_WEBHOOKS_FLEET_DATASET)
+    catalog_webhooks_fleet_dataset = catalog_webhooks_fleet_dataset['train']
+    secrets = catalog_webhooks_fleet_dataset['secret']
+    
+    if request.headers.get("X-Webhook-Secret") not in secrets:
         logger.warning("Invalid webhook secret received for model catalog")
         return Response("Invalid secret", status_code=401)
+    else:
+        secret = request.headers.get("X-Webhook-Secret")
+        idx = secrets.index(secret)
     
     data = await request.json()
     logger.info(f"Model catalog webhook payload: {json.dumps(data, indent=2)}")
 
     webhook_id = data.get("webhook", {}).get("id", None)
     
-    if webhook_id == MODEL_CATALOG_WEBHOOK:
+    if webhook_id == catalog_webhooks_fleet_dataset[idx]['webhook_id']:
         logger.info("Processing model catalog webhook")
         process_model_catalog_webhook(data)
     else: