Update main.py

main.py

@@ -1,14 +1,17 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Header
 from pydantic import BaseModel
 import os, requests
 
 app = FastAPI()
 
-# Pull your secret from the HF Spaces Secrets
+# Get both secrets from HF settings
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+ACCESS_KEY = os.getenv("API_key")  # Your custom VirtuComm API key
+
 if not GROQ_API_KEY:
-    # If this triggers, check Settings → Secrets → GROQ_API_KEY
-    raise RuntimeError("GROQ_API_KEY environment variable not set")
+    raise RuntimeError("GROQ_API_KEY not set in Hugging Face secrets.")
+if not ACCESS_KEY:
+    raise RuntimeError("API_key not set in Hugging Face secrets.")
 
 GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
 
@@ -17,7 +20,11 @@ class ChatRequest(BaseModel):
     message: str
 
 @app.post("/chat")
-def chat(req: ChatRequest):
+def chat(req: ChatRequest, x_api_key: str = Header(None)):
+    # Check that client included the correct API key
+    if x_api_key != ACCESS_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key.")
+
     headers = {
         "Authorization": f"Bearer {GROQ_API_KEY}",
         "Content-Type": "application/json"