from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from langchain_community.chat_models import ChatLlamaCpp

from langchain_core.messages import (
    HumanMessage,
    SystemMessage,
)

import uvicorn


# =========================================================
# FASTAPI APP
# =========================================================

app = FastAPI(
    title="Local GGUF Chat API",
    version="1.0.0",
)


# =========================================================
# LOAD MODEL
# =========================================================

llm = ChatLlamaCpp(
    model_path="Qwen2.5-0.5B-Instruct.Q4_K_M.gguf",

    # Context Window
    n_ctx=2048,

    # CPU Threads
    n_threads=8,

    # Batch Size
    n_batch=512,

    # GPU Layers
    # 0 = Full CPU inference
    n_gpu_layers=0,

    # Generation Settings
    temperature=0.7,
    max_tokens=512,

    # Streaming
    streaming=True,

    verbose=False,
)


# =========================================================
# REQUEST MODEL
# =========================================================

class ChatRequest(BaseModel):
    message: str


# =========================================================
# ROOT ENDPOINT
# =========================================================

@app.get("/")
def home():

    return {
        "message": "GGUF Chat API Running"
    }


# =========================================================
# STREAM GENERATOR
# =========================================================

def generate_response(user_message: str):

    messages = [

        SystemMessage(
            content=(
                "You are a helpful AI assistant."
            )
        ),

        HumanMessage(
            content=user_message
        )
    ]

    for chunk in llm.stream(messages):

        if chunk.content:
            yield chunk.content


# =========================================================
# CHAT ENDPOINT
# =========================================================

@app.post("/chat")
async def chat(request: ChatRequest):

    return StreamingResponse(
        generate_response(request.message),
        media_type="text/plain",
    )


# =========================================================
# MAIN
# =========================================================

if __name__ == "__main__":

    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=8000,
        reload=False,
    )