ARIA-AGENT/aria-brain/local_compact.py

"""
Compacting fuer ARIAs lokalen Modus (Fast-Lane / localOnly).

Ein lokales Modell hat ein kleines Kontextfenster (z.B. 32k auf einer
12-GB-Karte). Frueher bekam es starr die letzten 8 Nachrichten — alles davor
war fuer das lokale Modell einfach weg. Jetzt:

  * Fenster nach TOKENS statt fester Anzahl: so viele der letzten Nachrichten,
    wie in ~75 % des Kontextfensters passen (Systemprompt + Antwortplatz
    abgezogen), hoechstens LOCAL_MAX_TURNS (Tempo).
  * Was aus dem Fenster faellt, fasst das lokale Modell im HINTERGRUND zu
    einer laufenden Zusammenfassung zusammen ("Notizzettel"). Die liegt pro
    Projekt in /data/local_summaries.json und geht bei jeder lokalen Anfrage
    als "Bisheriger Verlauf" in den Systemprompt.

Unabhaengig davon laeuft weiter das Destillat (agent.distill_old_turns):
dauerhaft Wichtiges wandert als Fakt ins Gedaechtnis (Qdrant). Compacting =
Notizzettel fuer DIESES Gespraech, Destillat = Karteikarten fuers Archiv.
"""

from __future__ import annotations

import json
import logging
import os
import threading
import time
import urllib.request
from pathlib import Path

logger = logging.getLogger(__name__)

BRIDGE_URL = os.environ.get("BRIDGE_URL", "http://aria-bridge:8090")
SUMMARY_FILE = Path(os.environ.get("LOCAL_SUMMARY_FILE", "/data/local_summaries.json"))
LOCAL_MAX_TURNS = int(os.environ.get("LOCAL_MAX_TURNS", "12"))
DEFAULT_CTX = 8192
CHARS_PER_TOKEN = 2.8          # Deutsch/Code: eher vorsichtig schaetzen
COMPACT_MIN_TURNS = 4          # erst ab so vielen neu herausgefallenen Nachrichten
SUMMARY_MAX_TOKENS = 450

_lock = threading.Lock()
_running: set = set()          # Projekte, fuer die gerade verdichtet wird
_ctx_cache: dict = {"ts": 0.0, "models": {}}


def est_tokens(text: str) -> int:
    return int(len(text or "") / CHARS_PER_TOKEN) + 4


def model_ctx(model: str) -> int:
    """Kontextfenster des lokalen Modells laut Cluster (60 s gecacht)."""
    now = time.time()
    if now - _ctx_cache["ts"] > 60:
        try:
            with urllib.request.urlopen(f"{BRIDGE_URL}/internal/llm-models", timeout=3) as r:
                data = json.loads(r.read().decode("utf-8", "ignore"))
            _ctx_cache["models"] = {m.get("model"): int(m.get("ctx") or 0)
                                    for m in data.get("models") or []}
            _ctx_cache["ts"] = now
        except Exception as exc:
            logger.debug("[compact] Kontextfenster nicht abfragbar: %s", exc)
            _ctx_cache["ts"] = now - 50      # in 10 s nochmal versuchen
    return _ctx_cache["models"].get(model) or DEFAULT_CTX


def fit_window(turns: list, budget_tokens: int) -> list:
    """Die neuesten Nachrichten, die ins Budget passen (mind. die letzte)."""
    out, used = [], 0
    for t in reversed(turns[-LOCAL_MAX_TURNS:]):
        need = est_tokens(t.content)
        if out and used + need > budget_tokens:
            break
        out.append(t)
        used += need
    return list(reversed(out))


# ── Laufende Zusammenfassung pro Projekt ─────────────────────────────

def _load() -> dict:
    try:
        return json.loads(SUMMARY_FILE.read_text(encoding="utf-8"))
    except Exception:
        return {}


def get_summary(project_id: str) -> dict:
    """{text, upto_ts} — upto_ts = Zeitstempel der juengsten verdichteten Nachricht."""
    s = _load().get(project_id or "", {})
    return {"text": s.get("text", ""), "upto_ts": s.get("upto_ts", "")}


def _save_summary(project_id: str, text: str, upto_ts: str) -> None:
    with _lock:
        data = _load()
        data[project_id or ""] = {"text": text, "upto_ts": upto_ts,
                                  "updated": time.strftime("%Y-%m-%dT%H:%M:%S")}
        tmp = SUMMARY_FILE.with_suffix(".tmp")
        tmp.write_text(json.dumps(data, ensure_ascii=False, indent=1), encoding="utf-8")
        tmp.replace(SUMMARY_FILE)


def reset(project_id=None) -> None:
    """Zusammenfassung(en) verwerfen (z.B. beim Zuruecksetzen des Gespraechs)."""
    with _lock:
        if project_id is None:
            try:
                SUMMARY_FILE.unlink()
            except FileNotFoundError:
                pass
            return
        data = _load()
        data.pop(project_id or "", None)
        SUMMARY_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=1), encoding="utf-8")


def summary_block(project_id: str) -> str:
    s = get_summary(project_id).get("text", "").strip()
    if not s:
        return ""
    return ("## BISHERIGER VERLAUF (Zusammenfassung aelterer Nachrichten)\n"
            "Das ist eine Zusammenfassung dessen, was vor den letzten Nachrichten "
            "besprochen wurde. Nutze es als Hintergrund, zitiere es nicht.\n" + s)


def maybe_compact(project_id: str, all_turns: list, window: list, model: str,
                  chat_fn) -> None:
    """Startet im Hintergrund eine Verdichtung, wenn genug Nachrichten aus dem
    Fenster gefallen sind, die noch nicht im Notizzettel stehen.
    chat_fn = local_llm_chat (wird nur im Hintergrund-Thread aufgerufen)."""
    if not window:
        return
    first_kept = window[0].ts
    upto = get_summary(project_id)["upto_ts"]
    dropped = [t for t in all_turns if t.ts < first_kept and t.ts > upto]
    if len(dropped) < COMPACT_MIN_TURNS:
        return
    key = project_id or ""
    with _lock:
        if key in _running:
            return
        _running.add(key)
    threading.Thread(target=_compact, args=(key, dropped, model, chat_fn),
                     daemon=True, name=f"compact-{key or 'main'}").start()


def _compact(project_id: str, dropped: list, model: str, chat_fn) -> None:
    try:
        prev = get_summary(project_id)["text"]
        ctx = model_ctx(model)
        # Was in die Verdichtungs-Anfrage passt: ~60 % des Fensters
        room = int(ctx * 0.6 * CHARS_PER_TOKEN) - len(prev) - 2000
        lines, used, last_ts = [], 0, ""
        for t in dropped:                       # aelteste zuerst
            who = "Stefan" if t.role == "user" else "ARIA"
            line = f"{who}: {t.content.strip()}"
            if used + len(line) > room:
                line = line[: max(0, room - used)] + " …"
            lines.append(line)
            used += len(line)
            last_ts = t.ts
            if used >= room:
                break                           # Rest beim naechsten Durchlauf
        system = ("Du fuehrst einen laufenden Notizzettel ueber ein Gespraech zwischen "
                  "Stefan und seiner Assistentin ARIA. Schreibe auf Deutsch eine knappe "
                  "Zusammenfassung (hoechstens ~250 Woerter, Stichpunkte erlaubt): worum "
                  "es ging, was entschieden oder erledigt wurde, offene Punkte, wichtige "
                  "Namen/Zahlen. Kein Smalltalk, keine Einleitung, nur der Notizzettel.")
        user = ((f"Bisheriger Notizzettel:\n{prev}\n\n" if prev else "")
                + "Diese Nachrichten kommen neu dazu:\n" + "\n".join(lines)
                + "\n\nSchreibe den aktualisierten Notizzettel.")
        res = chat_fn([{"role": "system", "content": system}, {"role": "user", "content": user}],
                      max_tokens=SUMMARY_MAX_TOKENS, temperature=0.2, model=model, timeout=180)
        text = (res.get("content") or "").strip() if res.get("ok") else ""
        if not text:
            logger.warning("[compact] Verdichtung fehlgeschlagen (%s) — Notizzettel bleibt",
                           res.get("error"))
            return
        _save_summary(project_id, text, last_ts)
        logger.info("[compact] %d Nachrichten verdichtet (Projekt %r) → %d Zeichen Notizzettel",
                    len(dropped), project_id, len(text))
    except Exception as exc:
        logger.warning("[compact] Fehler: %s", exc)
    finally:
        with _lock:
            _running.discard(project_id or "")