from fastapi import FastAPI
from pydantic import BaseModel
from llama_cpp import Llama

app = FastAPI()

# Load model once on startup
llm = Llama(
    model_path="qwen2-1_5b-instruct-q4_0.gguf",
    n_ctx=2048,
    n_threads=2
)

# Request body structure
class ChatRequest(BaseModel):
    message: str

@app.get("/")
def home():
    return {"status": "AI server running"}

@app.post("/chat")
def chat(req: ChatRequest):
    output = llm.create_completion(
        prompt=req.message,
        max_tokens=300,
        temperature=0.7
    )

    return {
        "response": output["choices"][0]["text"].strip()
    }