Update app.py

app.py

@@ -1,12 +1,13 @@
-
 import os
 import sys
+import json
 import uuid
+import hmac
 import threading
 import importlib.util
 
 import torch
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response, stream_with_context
 from flask_cors import CORS
 from transformers import AutoConfig, AutoModelForCausalLM
 from tokenizers import Tokenizer
@@ -25,6 +26,13 @@ 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)
 
@@ -43,8 +51,6 @@ print(f"[Bridge] Loading {MODEL_ID}...")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"[Bridge] Selected device: {device}")
 
-# Download and import Luna's actual custom architecture.
-# No separate Python file is required in your application folder.
 architecture_path = hf_hub_download(
     repo_id=MODEL_ID,
     filename="modeling_luna.py"
@@ -62,7 +68,6 @@ module = importlib.util.module_from_spec(spec)
 sys.modules[spec.name] = module
 spec.loader.exec_module(module)
 
-# Load Luna's custom Transformers configuration and model.
 config = AutoConfig.from_pretrained(
     MODEL_ID,
     trust_remote_code=True
@@ -77,7 +82,6 @@ model = AutoModelForCausalLM.from_pretrained(
 model = model.to(device)
 model.eval()
 
-# Load Luna's tokenizer.
 tokenizer_path = hf_hub_download(
     repo_id=MODEL_ID,
     filename="luna_tokenizer.json"
@@ -120,7 +124,154 @@ def generate_response(prompt, history):
 
 
 # ============================================================
-# Chat endpoint
+# 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"])
@@ -150,11 +301,8 @@ def ask():
                 "logs": "Empty input."
             }), 400
 
-        # The browser should reuse this ID for the same conversation.
         session_id = str(data.get("session_id") or uuid.uuid4())
 
-        # Keep history retrieval, inference, and updates in one
-        # critical section so concurrent requests cannot corrupt it.
         with generation_lock:
             history = sessions.setdefault(session_id, [])
 
@@ -183,7 +331,7 @@ def ask():
 
 
 # ============================================================
-# Clear conversation history
+# Clear conversation history - unchanged
 # ============================================================
 
 @app.route("/clear", methods=["POST"])
@@ -205,20 +353,6 @@ def clear_history():
     })
 
 
-# ============================================================
-# Health check
-# ============================================================
-
-@app.route("/", methods=["GET"])
-def health():
-    return jsonify({
-        "service": "Aoban Luna 2 Flash",
-        "model": MODEL_ID,
-        "device": str(device),
-        "status": "online"
-    })
-
-
 # ============================================================
 # Start server
 # ============================================================
@@ -232,4 +366,4 @@ if __name__ == "__main__":
         host="0.0.0.0",
         port=port,
         threaded=True
-    )
+    )