Update main.py

main.py

@@ -3,6 +3,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import os
+import requests
 
 app = FastAPI()
 
@@ -54,16 +55,13 @@ def generate(item: Item):
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    # reject if key not the same 
+    # Reject if not authenticated
     apiKey = os.environ.get("API_KEY")
-    # if apiKey != key:
-    #     return jsonify({
-    #         "image_url": image_url,
-    #         "model": model,
-    #         "result": "Invalid API Key",
-    #     }), 400
-    
+    auth_header = request.headers['Authorization']
+    if not auth_header.startswith('Bearer ' + apiKey):
+        return jsonify({'message': 'Invalid or missing Bearer token'}), 401
+
     return {
         "response": generate(item),
-        "key": apiKey
+        "token": auth_header
     }