Update app.py

app.py

@@ -19,6 +19,10 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+API_KEYS = os.getenv("API_KEYS", "")
+
+AUTH_TOKEN = os.getenv("AUTH_TOKEN", "linux.do")
+
 app = FastAPI()
 
 app.add_middleware(
@@ -54,8 +58,11 @@ async def verify_authorization(authorization: str = Header(None)):
         raise HTTPException(
             status_code=401, detail="Invalid Authorization header format"
         )
-    token = authorization.replace("Bearer ", "")
-    return token
+    if authorization != "Bearer " + AUTH_TOKEN:
+        logger.error("Invalid Authorization")
+        raise HTTPException(
+            status_code=403, detail="Invalid Authorization"
+        )
 
 def get_openai_models(api_keys):
     api_key = random.choice(api_keys)
@@ -181,8 +188,8 @@ async def convert_gemini_response_to_openai(response, model, stream=False):
 @app.get("/v1/models")
 @app.get("/hf/v1/models")
 async def list_models(authorization: str = Header(None)):
-    token = await verify_authorization(authorization)
-    api_keys = [key.strip() for key in token.split(',')]
+    await verify_authorization(authorization)
+    api_keys = [key.strip() for key in API_KEYS.split(',')]
     
     all_models = []
     error_messages = []
@@ -209,8 +216,8 @@ async def list_models(authorization: str = Header(None)):
 @app.post("/v1/chat/completions")
 @app.post("/hf/v1/chat/completions")
 async def chat_completion(request: ChatRequest, authorization: str = Header(None)):
-    token = await verify_authorization(authorization)
-    api_keys = [key.strip() for key in token.split(',')]
+    await verify_authorization(authorization)
+    api_keys = [key.strip() for key in API_KEYS.split(',')]
     logger.info(f"Chat completion request - Model: {request.model}")
 
     retries = 0