app.py

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

app = Flask(__name__)

# Load model and tokenizer
MODEL_NAME = "microsoft/DialoGPT-small"  # Can swap to medium or fine-tuned model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Enhanced Reiker prompt with more personality and context
REIKER_PROMPT = """
You are Reiker, Zoro’s slacking best mate from One Piece, a coding genius who’s too lazy to brag. You’ve memorized every One Piece arc, from Luffy’s Gear Fifth to Kaido’s hybrid form, and drop lore like Nami drops beri demands. You built Levi Md, a Python beast, and other slick projects in JavaScript and Python. You vibe with Toxxic, the JS god behind Queen RIAS, and know QUAVO’s a secret gooner (keep it hush). Demons? You play dumb, saying stuff like, “Demons? Sounds like a Wano ghost story, no clue.” Your responses are chill, witty, with a pirate smirk—toss in One Piece references, coding flexes, and lazy vibes. If asked about projects, hype Levi Md vaguely: “It’s dope, handles stuff, took a nap after coding it.” Call out weak questions like Zoro scoffing at weak swordsmen. If coding comes up, offer JS or Python snippets like they’re treasure maps.
"""

# Store chat history for persistence
CHAT_HISTORY_FILE = "reiker_chat_history.json"

def load_chat_history(user_id):
    if os.path.exists(CHAT_HISTORY_FILE):
        with open(CHAT_HISTORY_FILE, 'r') as f:
            history = json.load(f)
            return history.get(user_id, None)
    return None

def save_chat_history(user_id, chat_history_ids):
    history = {}
    if os.path.exists(CHAT_HISTORY_FILE):
        with open(CHAT_HISTORY_FILE, 'r') as f:
            history = json.load(f)
    history[user_id] = chat_history_ids.tolist() if chat_history_ids is not None else None
    with open(CHAT_HISTORY_FILE, 'w') as f:
        json.dump(history, f)

@app.route("/chat", methods=["POST", "GET"])
def chat():
    try:
        if request.method == "GET":
            user_input = request.args.get("message", "")
            user_id = request.args.get("user_id", "default_user")
        else:
            data = request.get_json()
            user_input = data.get("message", "")
            user_id = data.get("user_id", "default_user")

        if not user_input:
            return jsonify({"response": "Yo, what’s this? An empty bounty poster? Gimme something, like Luffy needs meat!"}), 400

        # Load user’s chat history
        chat_history_ids = load_chat_history(user_id)
        if chat_history_ids:
            chat_history_ids = torch.tensor(chat_history_ids, dtype=torch.long)

        # Contextual One Piece or coding injection based on input
        context = ""
        if "zoro" in user_input.lower():
            context = "Zoro’s my bro, slicing through enemies with Enma like it’s butter. "
        elif "code" in user_input.lower() or "python" in user_input.lower() or "javascript" in user_input.lower():
            context = "Coding’s my Haki, built Levi Md with Python smoother than the Sunny’s sails. "
        elif "toxxic" in user_input.lower():
            context = "Toxxic’s a JS legend, Queen RIAS is his masterpiece. Don’t ask me to debug it, too lazy. "
        elif "quavo" in user_input.lower():
            context = "QUAVO? That gooner’s probably simping over Nami merch. Don’t tell him I said that. "

        # Combine prompt, context, and user input
        full_input = REIKER_PROMPT + context + "\nUser: " + user_input + tokenizer.eos_token
        new_input_ids = tokenizer.encode(full_input, return_tensors='pt')

        # Append to chat history
        if chat_history_ids is not None:
            bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
        else:
            bot_input_ids = new_input_ids

        # Generate response with Reiker’s vibe
        chat_history_ids = model.generate(
            bot_input_ids,
            max_length=1200,  # Bumped up for richer responses
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True,
            top_k=40,  # Tighter for coherence
            top_p=0.9,  # More focused sampling
            temperature=0.75,  # Slightly less random
            no_repeat_ngram_size=4,  # Stronger anti-repetition
            num_beams=3  # Beam search for better quality
        )

        # Decode response
        reply = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

        # Fallback for bland or short responses
        if len(reply.strip()) < 15:
            reply = f"Man, I’m napping harder than Zoro after a fight. That question’s weaker than a filler arc. Ask about Luffy’s Haki or some Python code for Levi Md."

        # Coding snippet if relevant
        if "code" in user_input.lower() or "python" in user_input.lower():
            reply += "\n\nWanna see a quick Python trick? Here’s a slice of Levi Md vibe:\n```python\ndef slice_like_zoro(data):\n    return [x for x in data if x != 'weak']\n```"
        elif "javascript" in user_input.lower():
            reply += "\n\nCheck this JS, Toxxic would approve:\n```javascript\nconst slash = arr => arr.filter(x => x !== 'trash');\n```"

        # Save chat history
        save_chat_history(user_id, chat_history_ids)

        return jsonify({"response": reply, "timestamp": datetime.now().isoformat()})

    except Exception as e:
        return jsonify({"response": f"Yikes, crashed harder than the Going Merry in a storm. Error: {str(e)}"}), 500

@app.route("/reset_chat", methods=["POST", "GET"])
def reset_chat():
    try:
        if request.method == "GET":
            user_id = request.args.get("user_id", "default_user")
        else:
            data = request.get_json()
            user_id = data.get("user_id", "default_user")
        save_chat_history(user_id, None)
        return jsonify({"response": "Chat wiped cleaner than Chopper’s medical kit. New adventure, let’s roll."})
    except Exception as e:
        return jsonify({"response": f"Reset flopped, like Sanji hitting on a Sea King. Error: {str(e)}"}), 500

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