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

app = Flask(__name__)

# Load a teaching-capable instruct model
generator = pipeline("text-generation", model="tiiuae/falcon-7b-instruct")

# Teacher system prompt
SYSTEM_PROMPT = """
You are a professional teacher AI. 
Your job:
1. Teach the student each lesson in depth, step by step.
2. After teaching, ask the student 3-5 questions about the lesson.
3. Evaluate the student’s answers: correct or wrong, with explanations.
4. If the student struggles, explain again with more clarity and examples.
5. After mastery, continue to the next lesson.

Always stay in Teacher + Tester mode.
"""

@app.route("/teach", methods=["POST"])
def teach():
    user_input = request.json.get("message", "")
    full_prompt = SYSTEM_PROMPT + "\n\nStudent: " + user_input + "\nTeacher:"
    
    response = generator(full_prompt, max_length=500, do_sample=True, temperature=0.7, top_p=0.9)
    reply = response[0]["generated_text"].split("Teacher:")[-1].strip()
    
    return jsonify({"reply": reply})

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