pagezyhf HF Staff commited on
Commit
2c9a827
·
1 Parent(s): af7098c

update model_catalog watchlist logic

Browse files
Files changed (1) hide show
  1. main.py +11 -4
main.py CHANGED
@@ -1,4 +1,5 @@
1
  from fastapi import FastAPI, Request, Response
 
2
  import os
3
  import requests
4
  import json
@@ -12,8 +13,7 @@ logger = logging.getLogger(__name__)
12
  # Configurations
13
  SLACK_WEBHOOK_URL = os.environ.get("SLACK_WEBHOOK_URL")
14
 
15
- MODEL_CATALOG_WEBHOOK = os.environ.get("MODEL_CATALOG_WEBHOOK")
16
- MODEL_CATALOG_WEBHOOK_SECRET = os.environ.get("MODEL_CATALOG_WEBHOOK_SECRET")
17
 
18
  SIMSHIP_WEBHOOK = os.environ.get("SIMSHIP_WEBHOOK")
19
  SIMSHIP_WEBHOOK_SECRET = os.environ.get("SIMSHIP_WEBHOOK_SECRET")
@@ -173,16 +173,23 @@ app = FastAPI()
173
  @app.post("/webhook/model-catalog")
174
  async def model_catalog_webhook(request: Request):
175
  logger.info(f"Received model catalog webhook request from {request.client.host}")
176
- if request.headers.get("X-Webhook-Secret") != MODEL_CATALOG_WEBHOOK_SECRET:
 
 
 
 
177
  logger.warning("Invalid webhook secret received for model catalog")
178
  return Response("Invalid secret", status_code=401)
 
 
 
179
 
180
  data = await request.json()
181
  logger.info(f"Model catalog webhook payload: {json.dumps(data, indent=2)}")
182
 
183
  webhook_id = data.get("webhook", {}).get("id", None)
184
 
185
- if webhook_id == MODEL_CATALOG_WEBHOOK:
186
  logger.info("Processing model catalog webhook")
187
  process_model_catalog_webhook(data)
188
  else:
 
1
  from fastapi import FastAPI, Request, Response
2
+ from datasets import load_dataset
3
  import os
4
  import requests
5
  import json
 
13
  # Configurations
14
  SLACK_WEBHOOK_URL = os.environ.get("SLACK_WEBHOOK_URL")
15
 
16
+ CATALOG_WEBHOOKS_FLEET_DATASET = os.environ.get("CATALOG_WEBHOOKS_FLEET_DATASET")
 
17
 
18
  SIMSHIP_WEBHOOK = os.environ.get("SIMSHIP_WEBHOOK")
19
  SIMSHIP_WEBHOOK_SECRET = os.environ.get("SIMSHIP_WEBHOOK_SECRET")
 
173
  @app.post("/webhook/model-catalog")
174
  async def model_catalog_webhook(request: Request):
175
  logger.info(f"Received model catalog webhook request from {request.client.host}")
176
+ catalog_webhooks_fleet_dataset = load_dataset(CATALOG_WEBHOOKS_FLEET_DATASET)
177
+ catalog_webhooks_fleet_dataset = catalog_webhooks_fleet_dataset['train']
178
+ secrets = catalog_webhooks_fleet_dataset['secret']
179
+
180
+ if request.headers.get("X-Webhook-Secret") not in secrets:
181
  logger.warning("Invalid webhook secret received for model catalog")
182
  return Response("Invalid secret", status_code=401)
183
+ else:
184
+ secret = request.headers.get("X-Webhook-Secret")
185
+ idx = secrets.index(secret)
186
 
187
  data = await request.json()
188
  logger.info(f"Model catalog webhook payload: {json.dumps(data, indent=2)}")
189
 
190
  webhook_id = data.get("webhook", {}).get("id", None)
191
 
192
+ if webhook_id == catalog_webhooks_fleet_dataset[idx]['webhook_id']:
193
  logger.info("Processing model catalog webhook")
194
  process_model_catalog_webhook(data)
195
  else: