Update app.py

app.py

@@ -1,7 +1,9 @@
+import os
 from fastapi import FastAPI, Header, HTTPException
+from pydantic import BaseModel
 from transformers import pipeline
 
-API_KEY = "CHANGE_THIS_KEY"
+API_KEY = os.getenv("API_KEY")
 
 app = FastAPI()
 
@@ -11,10 +13,16 @@ generator = pipeline(
     max_new_tokens=256
 )
 
+class PromptRequest(BaseModel):
+    prompt: str
+
 @app.post("/generate")
-def generate(prompt: str, x_api_key: str = Header(None)):
+def generate(data: PromptRequest, x_api_key: str = Header(None)):
+    if API_KEY is None:
+        raise HTTPException(status_code=500, detail="API key not configured")
+
     if x_api_key != API_KEY:
         raise HTTPException(status_code=401, detail="Invalid API Key")
 
-    output = generator(prompt)
-    return {"response": output[0]["generated_text"]}
+    result = generator(data.prompt)
+    return {"response": result[0]["generated_text"]}