from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

pipe = pipeline(
    "text-generation",
    model="microsoft/phi-3-mini-4k-instruct"
)

class Req(BaseModel):
    message: str

def build_prompt(user_msg: str) -> str:
    return (
        "<|system|>\n"
        "You are a helpful, clear, and concise assistant. "
        "Give accurate and well-structured answers."
        "<|end|>\n"
        "<|user|>\n"
        f"{user_msg}"
        "<|end|>\n"
        "<|assistant|>\n"
    )

@app.post("/chat")
def chat(req: Req):
    prompt = build_prompt(req.message)

    out = pipe(
        prompt,
        max_new_tokens=200,
        temperature=0.6,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.1
    )

    text = out[0]["generated_text"]
    reply = text.split("<|assistant|>")[-1].strip()

    return {"reply": reply}