from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
import json

app = Flask(__name__)

HISTORY_FILE = "history.json"
MAX_HISTORY = 10  # Keep only the last 10 lines per user

# Load history from file (or create empty)
if os.path.exists(HISTORY_FILE):
    with open(HISTORY_FILE, "r", encoding="utf-8") as f:
        user_histories = json.load(f)
else:
    user_histories = {}

# Choose a Transformers-compatible DeepSeek Distill model
MODEL_NAME = "Qwen/Qwen3-1.7B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto"
)
def save_history():
    """Save all histories to a JSON file."""
    with open(HISTORY_FILE, "w", encoding="utf-8") as f:
        json.dump(user_histories, f, ensure_ascii=False, indent=2)

@app.route("/message", methods=["GET"])
def handle_message():
    user_message = request.args.get("message")
    user_id = request.args.get("userid")

    if not user_message or not user_id:
        return jsonify({"error": "Both 'message' and 'userid' are required"}), 400

    # Retrieve or initialize history
    history = user_histories.get(user_id, [])


    history.append(f"User: {user_message}")


    history = history[-MAX_HISTORY:]


    conversation_text = "\n".join(history) + "\nAI:"


    inputs = tokenizer(conversation_text, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        temperature=0.6,
        do_sample=True
    )

    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)

 
    if "AI:" in reply:
        reply_text = reply.split("AI:")[-1].strip()
    else:
        reply_text = reply.strip()


    history.append(f"AI: {reply_text}")
    user_histories[user_id] = history

    
    save_history()

    return jsonify({"response": reply_text})

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