from flask import Flask, request, jsonify
from transformers import pipeline

app = Flask(__name__)

generator = pipeline(
    "text-generation",
    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0"
)

@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json()
    user_message = data.get("message", "")

    if not user_message:
        return jsonify({"reply": "No message received"})

    prompt = f"""
<|system|>
You are Stairo AI, a smart helpful chatbot. Answer clearly and simply.
</s>
<|user|>
{user_message}
</s>
<|assistant|>
"""

    result = generator(
        prompt,
        max_new_tokens=150,
        do_sample=True,
        temperature=0.6,
        repetition_penalty=1.2
    )

    output = result[0]["generated_text"]

    if "<|assistant|>" in output:
        reply = output.split("<|assistant|>")[-1].strip()
    else:
        reply = output

    return jsonify({"reply": reply})

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