Update main.py

main.py

@@ -1,9 +1,11 @@
-from fastapi import FastAPI, HTTPException, Header
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import os
-import requests
+from fastapi.security import OAuth2PasswordBearer
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
 app = FastAPI()
 
@@ -54,14 +56,14 @@ def generate(item: Item):
     return output
 
 @app.post("/generate/")
-async def generate_text(item: Item, authorization: str = Header(None)):
+async def generate_text(item: Item, token: Annotated[str, Depends(oauth2_scheme)]):
     # Reject if not authenticated
     apiKey = os.environ.get("API_KEY")
-    if not authorization or not authorization.startswith('Bearer ' + apiKey):
-        raise HTTPException(
-            status_code=401,
-            detail="Invalid or missing Bearer token"
-        )
+    if apiKey != token:
+        return {
+            "message": "Not authenticated",
+            "token": token
+        }
 
     return {
         "response": generate(item)