Update main.py

main.py

@@ -1,25 +1,16 @@
 from fastapi import FastAPI, Request, Header, HTTPException
 from pydantic import BaseModel
-from typing import Optional, Any, Dict, List
-import httpx
+from typing import Optional, Any, Dict
+import requests
 import os
-import socket
-from urllib.parse import urlparse
 
 app = FastAPI()
 
-# --- Configuration & Cleaning ---
-def get_clean_env(key):
-    val = os.environ.get(key)
-    if val:
-        # Remove whitespace, single quotes, and double quotes
-        return val.strip().strip("'").strip('"')
-    return None
+# --- Configuration ---
+DISCORD_WEBHOOK_URL = os.environ.get("DISCORD_WEBHOOK_URL")
+WEBHOOK_SECRET = os.environ.get("WEBHOOK_SECRET")
 
-DISCORD_WEBHOOK_URL = get_clean_env("DISCORD_WEBHOOK_URL")
-WEBHOOK_SECRET = get_clean_env("WEBHOOK_SECRET")
-
-# --- Pydantic Models ---
+# --- Models ---
 class HFRepo(BaseModel):
     type: str
     name: str
@@ -37,114 +28,80 @@ class HFPayload(BaseModel):
     comment: Optional[Dict[str, Any]] = None
     webhook: Optional[Dict[str, Any]] = None
 
-@app.on_event("startup")
-async def startup_event():
-    print("--- STARTUP DIAGNOSTICS ---")
-    # 1. Check if Variable Exists
-    if DISCORD_WEBHOOK_URL:
-        # Mask the URL for safety in logs
-        masked = DISCORD_WEBHOOK_URL[:8] + "..." + DISCORD_WEBHOOK_URL[-4:]
-        print(f"INFO: Webhook URL loaded: {masked}")
-        
-        # 2. Extract Hostname
-        try:
-            parsed = urlparse(DISCORD_WEBHOOK_URL)
-            hostname = parsed.hostname
-            if not hostname:
-                print("ERROR: Could not extract hostname from URL. Is the format 'https://...'?")
-            else:
-                print(f"DEBUG: Hostname extracted: '{hostname}'")
-                
-                # 3. Test DNS Resolution (The "Network Doctor")
-                try:
-                    # Note: socket.gethostbyname is blocking, but okay for a quick startup check
-                    ip = socket.gethostbyname(hostname)
-                    print(f"SUCCESS: DNS resolved {hostname} to {ip}")
-                except socket.gaierror as e:
-                    print(f"CRITICAL FAILURE: Could not resolve {hostname}. Error: {e}")
-                    print("-> This means the URL is wrong OR the container has no internet access.")
-        except Exception as e:
-            print(f"ERROR parsing URL: {e}")
-    else:
-        print("WARNING: No DISCORD_WEBHOOK_URL found. Discord messages will not be sent.")
-    print("---------------------------")
-
 @app.get("/")
 def home():
     return {"status": "Webhook listener is running"}
 
 @app.post("/webhook")
-async def handle_webhook(
+def handle_webhook(
     payload: HFPayload, 
     x_webhook_secret: Optional[str] = Header(None)
 ):
     # 1. Security Check
-    if WEBHOOK_SECRET and x_webhook_secret != WEBHOOK_SECRET:
-        print(f"Unauthorized access attempt. Header: {x_webhook_secret}")
-        raise HTTPException(status_code=401, detail="Invalid Secret")
+    if WEBHOOK_SECRET:
+        if not x_webhook_secret or x_webhook_secret.strip() != WEBHOOK_SECRET.strip():
+            print("Auth Failed")
+            raise HTTPException(status_code=401, detail="Invalid Secret")
 
-    # 2. Format for Discord
+    # 2. Build Embed
     embed = build_discord_embed(payload)
     
     # 3. Send to Discord
     if DISCORD_WEBHOOK_URL:
-        print(f"Sending notification for {payload.repo.name}...")
+        # Clean the URL just in case
+        url = DISCORD_WEBHOOK_URL.strip().strip("'").strip('"')
+        
         try:
-            # Added a timeout to prevent hanging
-            async with httpx.AsyncClient(timeout=10.0) as client:
-                resp = await client.post(
-                    DISCORD_WEBHOOK_URL, 
-                    json={"embeds": [embed]}
-                )
-                resp.raise_for_status()
-                print(f"Sent to Discord: {resp.status_code}")
-        except httpx.ConnectError as e:
-            print(f"CONNECTION ERROR: Could not connect to Discord. Check DNS/Network. Details: {e}")
-            return {"message": "Connection Failed", "error": str(e)}
+            print(f"Sending to Discord: {payload.repo.name} -> {payload.event.action}")
+            # requests is synchronous and handles DNS/Proxies much better in containers
+            resp = requests.post(url, json={"embeds": [embed]}, timeout=10)
+            resp.raise_for_status()
+            print("Success: Message sent.")
         except Exception as e:
-            print(f"ERROR connecting to Discord: {e}")
+            print(f"Error sending to Discord: {e}")
             return {"message": "Error sending to Discord", "error": str(e)}
     
     return {"message": "Processed"}
 
 def build_discord_embed(data: HFPayload):
+    # Basic data extraction
     repo_name = data.repo.name
     repo_url = data.repo.url.get('web', '')
     event_scope = data.event.scope
     event_action = data.event.action
     
-    # Default Color (Grey)
-    color = 0x95a5a6
+    # Colors
+    COLOR_BLUE = 0x3498db
+    COLOR_GREEN = 0x2ecc71
+    COLOR_GREY = 0x95a5a6
+    
+    color = COLOR_GREY
     title = f"Update on {repo_name}"
-    description = f"**Action:** {event_action}\n**Scope:** {event_scope}"
+    description = f"Action: {event_action}\nScope: {event_scope}"
 
-    # Customize based on event type
-    if event_scope.startswith("repo.content"):
-        # Code/Model update (Blue)
-        color = 0x3498db
+    # Smart Formatting
+    if "repo.content" in event_scope:
+        color = COLOR_BLUE
         title = f"New Commit to {repo_name}"
-        description = "Model/Dataset files have been updated."
+        description = f"The repository content has been updated.\n[View Changes]({repo_url})"
     
-    elif event_scope.startswith("discussion"):
-        # Discussion/PR (Green)
-        color = 0x2ecc71
-        disc_title = data.discussion.get('title', 'No Title') if data.discussion else "Unknown"
-        title = f"Discussion Update: {disc_title}"
+    elif "discussion" in event_scope:
+        color = COLOR_GREEN
+        disc_title = data.discussion.get('title', 'Discussion') if data.discussion else "Discussion"
+        title = f"Discussion: {disc_title}"
         
         if data.comment:
             content = data.comment.get('content', '')
             if len(content) > 200:
                 content = content[:200] + "..."
-            description = f"**New Comment** on {repo_name}:\n> {content}"
+            description = f"**New Comment:**\n{content}\n\n[View Discussion]({repo_url})"
         else:
-            description = f"Discussion status: **{event_action}**"
+            description = f"Discussion status updated: **{event_action}**"
 
     return {
         "title": title,
         "description": description,
         "url": repo_url,
         "color": color,
-        "footer": {
-            "text": "Hugging Face Webhook"
-        }
+        "footer": {"text": "Hugging Face Webhooks"}
     }