from fastapi import FastAPI, Request, HTTPException
import uvicorn
from llama_cpp import Llama

app = FastAPI()

print("Loading Llama model...")
llm = Llama(
    model_path="model.gguf",
    n_ctx=4096,
    n_gpu_layers=0,
)
print("Model loaded successfully!")

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        data = await request.json()
    except:
        raise HTTPException(status_code=400, detail="Invalid JSON body")
        
    messages = data.get("messages", [])
    max_tokens = data.get("max_tokens", 1024)
    temperature = data.get("temperature", 0.7)
    
    # Construct Dicta-LM Prompt
    prompt = "<s>"
    for m in messages:
        if m["role"] == "user":
            prompt += f"[INST] {m['content']} [/INST]\n"
        elif m["role"] == "assistant":
            prompt += f"{m['content']}</s> "
            
    import asyncio
    try:
        response = await asyncio.to_thread(
            llm,
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            stop=["</s>", "[INST]", "[/INST]"],
            echo=False
        )
        
        return {
            "id": response["id"],
            "object": "chat.completion",
            "created": response["created"],
            "model": "dictalm2.0-instruct",
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response["choices"][0]["text"].strip()
                    },
                    "finish_reason": response["choices"][0]["finish_reason"]
                }
            ],
            "usage": response["usage"]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/v1/models")
def get_models():
    return {
        "data": [
            {"id": "dictalm2.0-instruct", "object": "model", "owned_by": "dicta"}
        ]
    }

@app.get("/")
def health_check():
    return {"status": "running"}