Update app.py

app.py

@@ -1,10 +1,13 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
-from typing import List
+from typing import List, Optional
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-
+HF_TOKEN = os.getenv('HF_TOKEN')
+API_KEY = os.getenv('API_KEY')
 app = FastAPI()
+security = HTTPBearer()
 
 # Load model
 quant_config = BitsAndBytesConfig(
@@ -20,6 +23,14 @@ model = AutoModelForCausalLM.from_pretrained(
     quantization_config=quant_config,
     device_map="auto" # Use 'auto' to let it balance between GPU/CPU
 )
+def validate_api_key(auth: HTTPAuthorizationCredentials = Depends(security)):
+    if auth.credentials != API_KEY:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or missing API Key",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    return auth.credentials
 class Message(BaseModel):
     role: str
     content: str
@@ -30,7 +41,7 @@ class ChatRequest(BaseModel):
     temperature: float = 0.7
 
 @app.post("/v1/chat/completions")
-async def chat_endpoint(request: ChatRequest):
+async def chat_endpoint(request: ChatRequest, token: str = Depends(validate_api_key)):
     # Get the last user message
     prompt = request.messages[-1].content