neon-chat/main.py

#!/usr/bin/env python3
"""
NeonChat v2.1 — Lokaler RP Chatbot mit Ollama
Features: Streaming, Session-Passcodes, Multi-Device, Langzeitgedächtnis
"""

import json, os, time, asyncio, sqlite3, re, secrets, hashlib
from pathlib import Path
from datetime import datetime

import httpx
from fastapi import FastAPI, Request, UploadFile, File, HTTPException, Response, Cookie, Query
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

from comfyui_integration import is_comfyui_running, generate_portrait

# === Konfiguration ===
BASE_DIR = Path(__file__).parent
CHARACTERS_DIR = BASE_DIR / "characters"
DATA_DIR = BASE_DIR / "data"
DB_PATH = DATA_DIR / "memory.db"
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
DEFAULT_MODEL = os.environ.get("OLLAMA_MODEL", "gemma4:12b")
SESSION_COOKIE_NAME = "neonchat_session"
SESSION_DURATION = 60 * 60 * 24 * 365  # 1 Jahr

# === Datenbank ===
def init_db():
    DATA_DIR.mkdir(exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH))
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            character TEXT NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            timestamp REAL NOT NULL,
            summary TEXT
        );
        CREATE TABLE IF NOT EXISTS summaries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            character TEXT NOT NULL,
            summary TEXT NOT NULL,
            timestamp REAL NOT NULL
        );
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            passcode TEXT NOT NULL,
            created REAL NOT NULL,
            name TEXT DEFAULT ''
        );
        CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, character, timestamp);
        CREATE INDEX IF NOT EXISTS idx_summaries_session ON summaries(session_id, character, timestamp);
    """)
    conn.commit()
    conn.close()

init_db()

# === Session Management ===
def generate_passcode() -> str:
    """6-stelliger alphanumerischer Code."""
    return secrets.token_hex(3).upper()  # z.B. A3F2B1

def get_or_create_session(session_id: str = None) -> tuple:
    """Gibt (session_id, passcode) zurück. Erstellt neue Session falls nötig."""
    conn = sqlite3.connect(str(DB_PATH))
    if session_id:
        cur = conn.execute("SELECT id, passcode FROM sessions WHERE id = ?", (session_id,))
        row = cur.fetchone()
        if row:
            conn.close()
            return row[0], row[1]
    # Create new session
    new_id = secrets.token_hex(16)
    passcode = generate_passcode()
    conn.execute("INSERT INTO sessions (id, passcode, created) VALUES (?, ?, ?)", (new_id, passcode, time.time()))
    conn.commit()
    conn.close()
    return new_id, passcode

def get_session_by_passcode(passcode: str) -> str:
    """Findet Session-ID anhand eines Passcodes."""
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute("SELECT id FROM sessions WHERE passcode = ?", (passcode.upper(),))
    row = cur.fetchone()
    conn.close()
    return row[0] if row else None

# === Character Loading ===
def load_character(name: str) -> dict:
    path = CHARACTERS_DIR / f"{name}.json"
    if not path.exists():
        path = CHARACTERS_DIR / f"{name.lower()}.json"
    if not path.exists():
        for f in CHARACTERS_DIR.glob("*.json"):
            if f.stem.lower() == name.lower():
                path = f
                break
    if not path.exists():
        raise HTTPException(status_code=404, detail=f"Character {name} not found")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def list_characters() -> list:
    chars = []
    for f in CHARACTERS_DIR.glob("*.json"):
        with open(f, "r", encoding="utf-8") as fh:
            data = json.load(fh)
            chars.append({
                "name": data.get("name", f.stem),
                "filename": f.stem,
                "avatar": data.get("avatar", "👤"),
                "description": data.get("description", ""),
                "greeting": data.get("greeting", ""),
            })
    return chars

# === Memory System (Session-basiert) ===
def get_short_term_memory(session_id: str, character: str, limit: int = 20) -> list:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT role, content FROM messages WHERE session_id = ? AND character = ? ORDER BY timestamp DESC LIMIT ?",
        (session_id, character, limit)
    )
    rows = cur.fetchall()
    conn.close()
    return [{"role": r[0], "content": r[1]} for r in reversed(rows)]

def get_full_history(session_id: str, character: str, limit: int = 200) -> list:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT role, content, timestamp FROM messages WHERE session_id = ? AND character = ? ORDER BY timestamp ASC LIMIT ?",
        (session_id, character, limit)
    )
    rows = cur.fetchall()
    conn.close()
    return [{"role": r[0], "content": r[1], "timestamp": r[2]} for r in rows]

def get_long_term_memory(session_id: str, character: str) -> str:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT summary FROM summaries WHERE session_id = ? AND character = ? ORDER BY timestamp DESC LIMIT 3",
        (session_id, character)
    )
    rows = cur.fetchall()
    conn.close()
    if not rows:
        return ""
    return "\n\n".join([r[0] for r in rows])

def save_message(session_id: str, character: str, role: str, content: str):
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute(
        "INSERT INTO messages (session_id, character, role, content, timestamp) VALUES (?, ?, ?, ?, ?)",
        (session_id, character, role, content, time.time())
    )
    conn.commit()
    conn.close()

def needs_summary(session_id: str, character: str) -> bool:
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT COUNT(*) FROM messages WHERE session_id = ? AND character = ? AND summary IS NULL",
        (session_id, character)
    )
    count = cur.fetchone()[0]
    conn.close()
    return count >= 20

async def generate_summary(session_id: str, character: str):
    conn = sqlite3.connect(str(DB_PATH))
    cur = conn.execute(
        "SELECT id, role, content FROM messages WHERE session_id = ? AND character = ? AND summary IS NULL ORDER BY timestamp ASC LIMIT 20",
        (session_id, character)
    )
    rows = cur.fetchall()
    if not rows:
        conn.close()
        return
    
    chat_text = "\n".join([f"{'Du' if r[1] == 'user' else character}: {r[2]}" for r in rows])
    
    prompt = f"""Fasse das folgende Gespräch kurz zusammen. Konzentriere dich auf:
- Wichtige Fakten über den Nutzer
- Beziehungen und Emotionen
- Wichtige Ereignisse
- Charakterzüge die gezeigt wurden

Gespräch:
{chat_text}

Zusammenfassung (auf Deutsch, kurz):"""

    try:
        async with httpx.AsyncClient(timeout=60) as client:
            resp = await client.post(f"{OLLAMA_URL}/api/generate", json={
                "model": DEFAULT_MODEL,
                "prompt": prompt,
                "stream": False,
                "think": False,
                "options": {"temperature": 0.3, "num_predict": 300}
            })
            summary = resp.json().get("response", "").strip()
    except:
        summary = "Zusammenfassung nicht verfügbar."
    
    if not summary:
        summary = "Zusammenfassung nicht verfügbar."
    
    conn.execute(
        "INSERT INTO summaries (session_id, character, summary, timestamp) VALUES (?, ?, ?, ?)",
        (session_id, character, summary, time.time())
    )
    for r in rows:
        conn.execute("UPDATE messages SET summary = ? WHERE id = ?", (summary, r[0]))
    conn.commit()
    conn.close()

def clear_memory(session_id: str, character: str):
    conn = sqlite3.connect(str(DB_PATH))
    conn.execute("DELETE FROM messages WHERE session_id = ? AND character = ?", (session_id, character))
    conn.execute("DELETE FROM summaries WHERE session_id = ? AND character = ?", (session_id, character))
    conn.commit()
    conn.close()

# === Prompt Builder ===
def build_prompt(character: dict, short_term: list, long_term: str, user_message: str) -> list:
    name = character.get("name", "AI")
    personality = character.get("personality", "")
    background = character.get("background", "")
    behavior = character.get("behavior", "")
    tone = character.get("tone", "")
    nsfw = character.get("nsfw", "")
    scenario = character.get("scenario", "")
    
    system = f"""Du bist {name}. Antworte IMMER auf Deutsch.

PERSÖNLICHKEIT: {personality}

HINTERGRUND: {background}

VERHALTEN: {behavior}

SPRACHSTIL: {tone}

{f'SZENARIO: {scenario}' if scenario else ''}

{f'NSFW/SOZIALE REGELN: {nsfw}' if nsfw else ''}

REGELN:
- Bleibe immer in der Rolle von {name}
- Antworte natürlich und umgangssprachlich
- Verwende keine Formulierungen wie "als KI" oder "als Sprachmodell"
- Reagiere emotional angemessen
- Sei nicht übermäßig freundlich — bleibe authentisch
- Verwende deutsche Umgangssprache wo passend
- Wiederhole dich nicht — variiere deine Ausdrücke
- Wenn der Nutzer etwas sagt das du schon weißt (aus Erinnerungen), beziehe dich darauf"""

    if long_term:
        system += f"\n\nERINNERUNGEN (Langzeitgedächtnis):\n{long_term}"

    messages = [{"role": "system", "content": system}]
    for msg in short_term:
        if msg["role"] == "user":
            messages.append({"role": "user", "content": msg["content"]})
        else:
            messages.append({"role": "assistant", "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})
    
    return messages

# === FastAPI ===
app = FastAPI(title="NeonChat")
templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))
app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")

def set_session_cookie(response, sid: str):
    response.set_cookie(SESSION_COOKIE_NAME, sid, max_age=SESSION_DURATION, httponly=False, samesite="lax")

@app.get("/", response_class=HTMLResponse)
async def index(request: Request, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, passcode = get_or_create_session(session_id)
    chars = list_characters()
    response = templates.TemplateResponse(request, "index.html", {"characters": chars, "passcode": passcode})
    if not session_id or session_id != sid:
        set_session_cookie(response, sid)
    return response

@app.get("/create", response_class=HTMLResponse)
async def create_page(request: Request):
    return templates.TemplateResponse(request, "create.html", {})

@app.get("/chat/{character_name}", response_class=HTMLResponse)
async def chat_page(request: Request, character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, passcode = get_or_create_session(session_id)
    char = load_character(character_name)
    history = get_full_history(sid, character_name, limit=200)
    long_term = get_long_term_memory(sid, character_name)
    response = templates.TemplateResponse(request, "chat.html", {
        "character": char,
        "history": history,
        "long_term": long_term,
        "session_id": sid,
        "passcode": passcode,
    })
    if not session_id or session_id != sid:
        set_session_cookie(response, sid)
    return response

@app.get("/join", response_class=HTMLResponse)
async def join_page(request: Request):
    """Passcode-Eingabe-Seite für andere Geräte."""
    return templates.TemplateResponse(request, "join.html", {})

@app.post("/api/session/join")
async def api_join_session(request: Request):
    """Verbindet ein anderes Gerät mit einer bestehenden Session via Passcode."""
    body = await request.json()
    passcode = body.get("passcode", "").strip().upper()
    if not passcode:
        return JSONResponse({"error": "Passcode erforderlich"}, status_code=400)
    
    sid = get_session_by_passcode(passcode)
    if not sid:
        return JSONResponse({"error": "Ungültiger Passcode"}, status_code=404)
    
    response = JSONResponse({"status": "ok", "session_id": sid})
    set_session_cookie(response, sid)
    return response

@app.get("/api/session/info")
async def api_session_info(session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    """Gibt Session-Info und Passcode zurück."""
    sid, passcode = get_or_create_session(session_id)
    return JSONResponse({"session_id": sid, "passcode": passcode})

@app.post("/api/chat/{character_name}")
async def chat_api(character_name: str, request: Request, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    char = load_character(character_name)
    body = await request.json()
    user_message = body.get("message", "")
    stream = body.get("stream", True)
    
    if not user_message.strip():
        return JSONResponse({"error": "Leere Nachricht"}, status_code=400)
    
    save_message(sid, character_name, "user", user_message)
    
    short_term = get_short_term_memory(sid, character_name, limit=20)
    long_term = get_long_term_memory(sid, character_name)
    messages = build_prompt(char, short_term, long_term, user_message)
    
    model = char.get("model") or DEFAULT_MODEL
    
    if stream:
        async def stream_response():
            full_response = ""
            try:
                async with httpx.AsyncClient(timeout=300) as client:
                    async with client.stream("POST", f"{OLLAMA_URL}/api/chat", json={
                        "model": model,
                        "messages": messages,
                        "stream": True,
                        "think": False,
                        "options": {
                            "temperature": char.get("temperature", 0.8),
                            "top_p": 0.9,
                            "num_predict": 800,
                            "repeat_penalty": 1.1,
                        }
                    }) as resp:
                        async for line in resp.aiter_lines():
                            if not line.strip():
                                continue
                            try:
                                chunk = json.loads(line)
                                content = chunk.get("message", {}).get("content", "")
                                if content:
                                    full_response += content
                                    yield f"data: {json.dumps({'token': content})}\n\n"
                                if chunk.get("done"):
                                    break
                            except json.JSONDecodeError:
                                continue
            except Exception as e:
                full_response = f"*(Fehler: {e})*"
                yield f"data: {json.dumps({'token': full_response})}\n\n"
            
            if full_response.strip():
                save_message(sid, character_name, "assistant", full_response.strip())
            
            if needs_summary(sid, character_name):
                asyncio.create_task(generate_summary(sid, character_name))
            
            yield f"data: {json.dumps({'done': True})}\n\n"
        
        return StreamingResponse(stream_response(), media_type="text/event-stream")
    else:
        try:
            async with httpx.AsyncClient(timeout=300) as client:
                resp = await client.post(f"{OLLAMA_URL}/api/chat", json={
                    "model": model,
                    "messages": messages,
                    "stream": False,
                    "think": False,
                    "options": {
                        "temperature": char.get("temperature", 0.8),
                        "top_p": 0.9,
                        "num_predict": 800,
                        "repeat_penalty": 1.1,
                    }
                })
                data = resp.json()
                response_text = data.get("message", {}).get("content", "").strip()
        except Exception as e:
            response_text = f"*(Fehler: {e})*"
        
        save_message(sid, character_name, "assistant", response_text)
        
        if needs_summary(sid, character_name):
            asyncio.create_task(generate_summary(sid, character_name))
        
        return JSONResponse({"response": response_text, "character": character_name})

@app.get("/api/characters")
async def api_list_characters():
    return JSONResponse({"characters": list_characters()})

@app.post("/api/character/create")
async def api_create_character(request: Request):
    body = await request.json()
    name = body.get("name", "unnamed")
    safe_name = re.sub(r'[^a-zA-Z0-9_\-]', '_', name.lower())
    path = CHARACTERS_DIR / f"{safe_name}.json"
    
    char = {
        "name": name,
        "avatar": body.get("avatar", "👤"),
        "description": body.get("description", ""),
        "personality": body.get("personality", ""),
        "background": body.get("background", ""),
        "behavior": body.get("behavior", ""),
        "tone": body.get("tone", ""),
        "scenario": body.get("scenario", ""),
        "nsfw": body.get("nsfw", ""),
        "greeting": body.get("greeting", ""),
        "model": body.get("model", ""),
        "temperature": body.get("temperature", 0.8),
        "appearance": body.get("appearance", {}),
    }
    
    with open(path, "w", encoding="utf-8") as f:
        json.dump(char, f, indent=2, ensure_ascii=False)
    
    return JSONResponse({"status": "ok", "character": safe_name})

@app.get("/api/character/{name}")
async def api_get_character(name: str):
    char = load_character(name)
    return JSONResponse(char)

@app.post("/api/memory/clear/{character_name}")
async def api_clear_memory(character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    clear_memory(sid, character_name)
    return JSONResponse({"status": "ok"})

@app.get("/api/memory/{character_name}")
async def api_get_memory(character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    short = get_short_term_memory(sid, character_name, limit=100)
    long = get_long_term_memory(sid, character_name)
    return JSONResponse({"short_term": short, "long_term": long})

@app.get("/api/history/{character_name}")
async def api_get_history(character_name: str, session_id: str = Cookie(None, alias=SESSION_COOKIE_NAME)):
    sid, _ = get_or_create_session(session_id)
    history = get_full_history(sid, character_name, limit=200)
    return JSONResponse({"history": history, "session_id": sid})

@app.get("/api/models")
async def api_list_models():
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"{OLLAMA_URL}/api/tags")
            models = [m["name"] for m in resp.json().get("models", [])]
            return JSONResponse({"models": models})
    except:
        return JSONResponse({"models": [], "error": "Ollama nicht erreichbar"})

@app.get("/api/comfyui/status")
async def api_comfyui_status():
    return JSONResponse({"running": is_comfyui_running()})

# In-memory tracking for portrait generation
_portrait_jobs = {}

@app.post("/api/character/{name}/portrait")
async def api_generate_portrait(name: str, request: Request):
    """Startet Portrait-Generierung — gibt sofort eine Job-ID zurück."""
    char = load_character(name)
    
    if not is_comfyui_running():
        return JSONResponse({"error": "ComfyUI läuft nicht. Starte ComfyUI zuerst."}, status_code=503)
    
    body = await request.json() if request.headers.get("content-type") == "application/json" else {}
    custom_prompt = body.get("prompt", "")
    
    import uuid
    job_id = str(uuid.uuid4())[:8]
    _portrait_jobs[job_id] = {"status": "generating", "image": None, "error": None}
    
    async def run_generation():
        image_path = generate_portrait(name, char, custom_prompt)
        if image_path:
            _portrait_jobs[job_id] = {"status": "done", "image": image_path, "error": None}
        else:
            _portrait_jobs[job_id] = {"status": "error", "image": None, "error": "Generierung fehlgeschlagen"}
    
    asyncio.create_task(run_generation())
    
    return JSONResponse({"status": "started", "job_id": job_id})

@app.get("/api/portrait/status/{job_id}")
async def api_portrait_status(job_id: str):
    """Fragt den Status einer Portrait-Generierung ab."""
    if job_id not in _portrait_jobs:
        return JSONResponse({"status": "not_found"}, status_code=404)
    return JSONResponse(_portrait_jobs[job_id])

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5252)