app.py

import os
import sys
import json
import uuid
import hmac
import threading
import importlib.util

import torch
from flask import Flask, request, jsonify, Response, stream_with_context
from flask_cors import CORS
from transformers import AutoConfig, AutoModelForCausalLM
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download


# ============================================================
# Configuration
# ============================================================

MODEL_ID = "Aobangaming/luna-2-70m"
MAX_LEN = 1024

TOP_K = 40
TOP_P = 0.6
PENALTY = 1.2
TEMPERATURE = 1.0

# Poe sends this as "Authorization: Bearer <key>" on every request.
# Set it as a Secret in your Hugging Face Space (POE_ACCESS_KEY).
# If unset, Poe requests are accepted without auth (fine for testing only).
POE_ACCESS_KEY = os.environ.get("POE_ACCESS_KEY", "")

POE_INTRO_MESSAGE = "Hi! I'm Aoban Luna 2 Flash. Say something."

app = Flask(__name__)
CORS(app)

# Protect shared model inference and session history.
generation_lock = threading.Lock()
sessions = {}


# ============================================================
# Load Luna 2's custom architecture
# ============================================================

print("[Bridge] Waking up Aoban Luna 2 Flash...")
print(f"[Bridge] Loading {MODEL_ID}...")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[Bridge] Selected device: {device}")

architecture_path = hf_hub_download(
    repo_id=MODEL_ID,
    filename="modeling_luna.py"
)

spec = importlib.util.spec_from_file_location(
    "modeling_luna",
    architecture_path
)

if spec is None or spec.loader is None:
    raise RuntimeError("Could not load Luna's custom architecture.")

module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)

config = AutoConfig.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    config=config,
    trust_remote_code=True
)

model = model.to(device)
model.eval()

tokenizer_path = hf_hub_download(
    repo_id=MODEL_ID,
    filename="luna_tokenizer.json"
)

tokenizer = Tokenizer.from_file(tokenizer_path)

print("[Bridge] Luna 2 loaded successfully.")
print(f"[Bridge] Running on {device}.")
print("[Bridge] Ready to accept requests.")


# ============================================================
# Text generation
# ============================================================

def generate_response(prompt, history):
    """
    Use Luna's original custom generation function.
    Do not replace this with model.generate().
    """

    response = module.generate_text(
        model,
        tokenizer,
        prompt,
        max_len=MAX_LEN,
        device=device,
        top_k=TOP_K,
        top_p=TOP_P,
        penalty=PENALTY,
        temperature=TEMPERATURE,
        chat_history=history
    )

    if response is None:
        return ""

    return str(response).strip()


# ============================================================
# Poe helpers
# ============================================================

def sse(event, data):
    """Format one Server-Sent Event the way Poe expects."""
    return f"event: {event}\ndata: {json.dumps(data)}\n\n"


def poe_authorized():
    if not POE_ACCESS_KEY:
        return True

    header = request.headers.get("Authorization", "")
    token = header[7:] if header.lower().startswith("bearer ") else ""
    return hmac.compare_digest(token, POE_ACCESS_KEY)


def poe_to_luna_history(query):
    """
    Poe sends the FULL conversation on every request, so Poe bots are
    stateless. Convert it to (prompt, history) in Luna's format.

    Poe roles: "system" | "user" | "bot"
    Luna roles: "user" | "assistant"
    """
    messages = []

    for msg in query:
        if not isinstance(msg, dict):
            continue

        role = msg.get("role")
        content = msg.get("content")

        if not isinstance(content, str) or not content.strip():
            continue

        if role == "user":
            messages.append({"role": "user", "content": content.strip()})
        elif role == "bot":
            messages.append({"role": "assistant", "content": content.strip()})
        # "system" messages are ignored; Luna has no system prompt.

    # The prompt is the latest user message; everything before it is history.
    prompt = ""
    history = []

    for i in range(len(messages) - 1, -1, -1):
        if messages[i]["role"] == "user":
            prompt = messages[i]["content"]
            history = messages[:i]
            break

    return prompt, history


def handle_poe_query(body):
    prompt, history = poe_to_luna_history(body.get("query", []))

    def stream():
        yield sse("meta", {
            "content_type": "text/markdown",
            "linkify": True,
            "suggested_replies": False
        })

        if not prompt:
            yield sse("error", {
                "text": "I didn't get a message to respond to.",
                "allow_retry": False
            })
            yield sse("done", {})
            return

        try:
            with generation_lock:
                response = generate_response(prompt, history)

            if not response:
                response = "..."

            yield sse("text", {"text": response})

        except Exception as e:
            app.logger.exception("Luna generation failed (Poe)")
            yield sse("error", {
                "text": f"Brain Error: {e}",
                "allow_retry": True
            })

        yield sse("done", {})

    return Response(
        stream_with_context(stream()),
        mimetype="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no"
        }
    )


def handle_poe_request(body):
    if not poe_authorized():
        return jsonify({"error": "Unauthorized."}), 401

    request_type = body.get("type")

    if request_type == "query":
        return handle_poe_query(body)

    if request_type == "settings":
        return jsonify({
            "server_bot_dependencies": {},
            "allow_attachments": False,
            "introduction_message": POE_INTRO_MESSAGE
        })

    # "report_feedback" and "report_error" just need a 200 JSON reply.
    return jsonify({})


# ============================================================
# Poe endpoint + health check (same URL: "/")
# ============================================================
# Poe POSTs to the bot URL you register (the Space's root URL).
# The browser/health check keeps using GET on the same route.

@app.route("/", methods=["GET", "POST"])
def root():
    if request.method == "POST":
        body = request.get_json(silent=True)

        if not isinstance(body, dict) or "type" not in body:
            return jsonify({"error": "Invalid Poe request."}), 400

        return handle_poe_request(body)

    return jsonify({
        "service": "Aoban Luna 2 Flash",
        "model": MODEL_ID,
        "device": str(device),
        "status": "online"
    })


# ============================================================
# Chat endpoint (website / your JS script) - unchanged
# ============================================================

@app.route("/ask", methods=["POST"])
def ask():
    try:
        data = request.get_json(silent=True)

        if not isinstance(data, dict):
            return jsonify({
                "response": "Invalid request body.",
                "logs": "Expected a JSON object."
            }), 400

        prompt = data.get("input", "")

        if not isinstance(prompt, str):
            return jsonify({
                "response": "Input must be text.",
                "logs": "Invalid input type."
            }), 400

        prompt = prompt.strip()

        if not prompt:
            return jsonify({
                "response": "Please enter a message.",
                "logs": "Empty input."
            }), 400

        session_id = str(data.get("session_id") or uuid.uuid4())

        with generation_lock:
            history = sessions.setdefault(session_id, [])

            response = generate_response(prompt, list(history))

            history.extend([
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": response}
            ])

        print(f"[Bridge] Generated response for session {session_id[:8]}")

        return jsonify({
            "response": response,
            "session_id": session_id,
            "logs": "Generation successful."
        })

    except Exception as e:
        app.logger.exception("Luna generation failed")

        return jsonify({
            "response": f"Brain Error: {str(e)}",
            "logs": str(e)
        }), 500


# ============================================================
# Clear conversation history - unchanged
# ============================================================

@app.route("/clear", methods=["POST"])
def clear_history():
    data = request.get_json(silent=True) or {}
    session_id = data.get("session_id")

    if not isinstance(session_id, str) or not session_id.strip():
        return jsonify({
            "error": "Missing or invalid session_id."
        }), 400

    with generation_lock:
        sessions.pop(session_id, None)

    return jsonify({
        "status": "History cleared.",
        "session_id": session_id
    })


# ============================================================
# Start server
# ============================================================

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))

    print(f"[Bridge] Starting Flask server on port {port}...")

    app.run(
        host="0.0.0.0",
        port=port,
        threaded=True
    )