Update app.py

app.py

@@ -1,17 +1,33 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException, Depends
 from llama_cpp import Llama
 
 app = FastAPI()
-# Load your 3.6GB model
-llm = Llama(model_path="./mpt-7b-chat.ggmlv0.q4_0.bin")
 
-@app.post("/v1/chat")
-async def chat(request: Request):
-    data = await request.json()
-    # Basic API logic:
-    response = llm(data["prompt"], max_tokens=100)
-    return {"choices": [{"text": response["choices"][0]["text"]}]}
+# Set your API Key here
+API_KEY = "my-secret-key-123"
+
+# Load the model we downloaded in the Dockerfile
+# MPT models need the 'mpt' type specified for older llama-cpp
+llm = Llama(model_path="./model.bin", model_type="mpt")
+
+# Security check for your API Key
+def verify_key(request: Request):
+    key = request.headers.get("Authorization")
+    if key != f"Bearer {API_KEY}":
+        raise HTTPException(status_code=403, detail="Invalid API Key")
 
 @app.get("/")
 def home():
-    return {"status": "Your AI is Always Up!"}
+    return {"message": "Your MPT-7B Website is Online!", "api_endpoint": "/v1/chat"}
+
+@app.post("/v1/chat", dependencies=[Depends(verify_key)])
+async def chat(request: Request):
+    data = await request.json()
+    prompt = data.get("prompt", "Hello!")
+    
+    # Generate the response
+    output = llm(f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n", 
+                 stop=["<|im_end|>"], 
+                 max_tokens=256)
+                 
+    return {"response": output["choices"][0]["text"]}