Update app/main.py

app/main.py

@@ -1,5 +1,6 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Depends, Security
 from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from gradio_client import Client
 import time
@@ -11,6 +12,16 @@ DEFAULT_API = "/chat"
 
 client = Client(SPACE_ID)
 
+# Security setup
+security = HTTPBearer()
+VALID_API_KEY = "sk-1234"  # Replace with your actual API key
+
+async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    if credentials.scheme != "Bearer":
+        raise HTTPException(status_code=403, detail="Invalid authentication scheme")
+    if credentials.credentials != VALID_API_KEY:
+        raise HTTPException(status_code=403, detail="Invalid API key")
+    return credentials.credentials
 
 def chat_with_gradio(message: str, api_name: str = DEFAULT_API):
     """
@@ -21,14 +32,13 @@ def chat_with_gradio(message: str, api_name: str = DEFAULT_API):
     except Exception as e:
         raise RuntimeError(f"Gradio API error: {e}")
 
-
 class ChatRequest(BaseModel):
     message: str
     api_name: str = DEFAULT_API
 
 app = FastAPI()
 
-@app.post("/chat")
+@app.post("/chat", dependencies=[Depends(get_api_key)])
 async def chat_endpoint(req: ChatRequest):
     """Forward chat requests to the Gradio API."""
     try:
@@ -37,7 +47,7 @@ async def chat_endpoint(req: ChatRequest):
     except RuntimeError as e:
         raise HTTPException(status_code=502, detail=str(e))
 
-@app.post("/v1/chat/completions")
+@app.post("/v1/chat/completions", dependencies=[Depends(get_api_key)])
 async def openai_chat_completions(request: Request):
     """
     OpenAI-compatible chat completions endpoint that forwards to Gradio.