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

app = FastAPI()

# 👇 Request body structure
class PromptRequest(BaseModel):
    prompt: str
    max_tokens: int = 200

# 👇 Download model from your HF repo
MODEL_PATH = hf_hub_download(
    repo_id="your-username/qwen-gguf",   # CHANGE THIS
    filename="qwen2-1_5b-instruct-q4_0.gguf"
)

# 👇 Load model once (very important)
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=1024,        # keep low for HF free tier
    n_threads=2        # reduce CPU usage
)

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

@app.post("/chat")
def chat(req: PromptRequest):
    output = llm(
        req.prompt,
        max_tokens=req.max_tokens,
        stop=["</s>"]
    )

    return {
        "response": output["choices"][0]["text"]
    }