Update main.py

main.py

@@ -54,15 +54,14 @@ def generate(item: Item):
     return output
 
 @app.post("/generate/")
-async def generate_text(item: Item):
+async def generate_text(item: Item, authorization: str = Header(None)):
     # Reject if not authenticated
     apiKey = os.environ.get("API_KEY")
-    authHeader = requests.headers['Authorization']
-    if not auth_header.startswith('Bearer ' + apiKey):
-        return jsonify({
-            "message": "Invalid or missing Bearer token",
-            "token": authHeader
-        }), 401
+    if not authorization or not authorization.startswith('Bearer ' + apiKey):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid or missing Bearer token"
+        )
 
     return {
         "response": generate(item)