import os
from llama_cpp import Llama
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

print("Loading Qwen 2.5 Coder 7B (4 GB)...", flush=True)
llm = Llama(
    model_path="/app/models/qwen2.5-coder-7b-instruct-q4_k_m.gguf",
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=0,
    verbose=False,
)
print("Qwen 7B loaded! Ready for inference.", flush=True)

class ChatRequest(BaseModel):
    messages: list
    max_tokens: int = 200
    temperature: float = 0.7

@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
    prompt = ""
    for msg in req.messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if role == "user":
            prompt += f"<|im_start|>user\n{content}<|im_end|>\n"
        elif role == "system":
            prompt += f"<|im_start|>system\n{content}<|im_end|>\n"
        elif role == "assistant":
            prompt += f"<|im_start|>assistant\n{content}<|im_end|>\n"
    prompt += "<|im_start|>assistant\n"
    
    output = llm(prompt, max_tokens=req.max_tokens, temperature=req.temperature, stop=["<|im_end|>"])
    text = output["choices"][0]["text"]
    
    return {
        "choices": [{"message": {"role": "assistant", "content": text}}],
        "model": "qwen2.5-coder-7b",
    }

@app.get("/")
async def root():
    return {"status": "ready", "model": "Qwen 2.5 Coder 7B", "humaneval": "65%"}

@app.get("/health")
async def health():
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)