project-lyra/lyra/summary.py

"""Session summarization: compact a session's raw exchanges into a stored gist.

This is the first consolidation stage. Raw exchanges stay for detail recall; the
summary is what surfaces when an *older* session is recalled, and it's the input
to the profile (semantic memory) and era-rollup tiers.

Long sessions are summarized in chunks, then the partial gists are merged, so a
big imported conversation doesn't blow the local model's context window.
"""
from __future__ import annotations

import sys
import threading
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed

from lyra import config, llm, logbus, memory
from lyra.llm import Backend, Message

# Consolidation LLM budget. A gist is short (a handful of sentences), so cap the
# generation hard — an uncapped local model will otherwise ramble for thousands
# of tokens and, on a slow GPU, blow the request timeout. 768 is ~3x the longest
# real gist we've stored.
SUMMARY_MAX_TOKENS = 768
# Attempts on the primary backend before falling back to cloud.
MI50_ATTEMPTS = 2
# Per-call timeout (seconds). A capped 768-token gist finishes in ~60-90s on the
# MI50; 150s is headroom but bails a hung call fast so fallback isn't slow.
SUMMARY_TIMEOUT = 150

# Degenerate-output guard. A wedged local model (e.g. an overheated GPU) returns
# a single character repeated ("?????") as a *successful* 200, which no timeout or
# exception catches — so validate the text and treat junk as a failure. Real gists
# are diverse prose; flag output whose most-common non-space char dominates. Short
# outputs are exempt (nothing meaningful to judge).
_DEGENERATE_MIN_CHARS = 24
_DEGENERATE_CHAR_RATIO = 0.5


class DegenerateOutput(RuntimeError):
    """A backend returned junk (e.g. one char repeated) as a successful response."""


def _looks_degenerate(text: str) -> bool:
    stripped = "".join(text.split())
    if len(stripped) < _DEGENERATE_MIN_CHARS:
        return False
    return max(Counter(stripped).values()) / len(stripped) > _DEGENERATE_CHAR_RATIO

# Re-summarize a session once it has accumulated this many new raw exchanges.
SUMMARIZE_AFTER = 20
# Transcript budget per LLM call; longer sessions are chunked + merged. Cloud has
# a large context window; the local llama.cpp/Ollama servers have small ones, so a
# 24k-char chunk overflows them ("Context size has been exceeded") — keep local small.
MAX_TRANSCRIPT_CHARS = 24000
LOCAL_TRANSCRIPT_CHARS = 8000


def _budget(backend: Backend) -> int:
    return MAX_TRANSCRIPT_CHARS if backend == "cloud" else LOCAL_TRANSCRIPT_CHARS

_PROMPT = """You are compacting a conversation into a long-term memory record \
(not replying to anyone). Write a concise gist of the session below: what was \
discussed, key decisions or outcomes, concrete specifics worth keeping (names, \
places, numbers, hands), and the user's apparent mood/state. Third person, \
referring to the user as "Brian". 4-8 sentences. No preamble."""


def _transcript(exchanges: list[memory.Exchange]) -> str:
    return "\n".join(f"{ex.role}: {ex.content}" for ex in exchanges)


def _chunk(text: str, budget: int) -> list[str]:
    """Split on line boundaries into pieces under `budget` chars."""
    chunks, buf, size = [], [], 0
    for line in text.splitlines(keepends=True):
        if size + len(line) > budget and buf:
            chunks.append("".join(buf))
            buf, size = [], 0
        buf.append(line)
        size += len(line)
    if buf:
        chunks.append("".join(buf))
    return chunks


def _summarize_text(text: str, backend: Backend) -> str:
    messages: list[Message] = [
        {"role": "system", "content": _PROMPT},
        {"role": "user", "content": text},
    ]

    def _call(be: Backend) -> str:
        out = llm.complete(messages, backend=be,
                           max_tokens=SUMMARY_MAX_TOKENS, timeout=SUMMARY_TIMEOUT)
        if _looks_degenerate(out):
            raise DegenerateOutput(f"{be} returned degenerate output ({len(out)} chars)")
        return out

    # Try the primary backend a bounded number of times (each call fast-fails via
    # SUMMARY_TIMEOUT), with a short backoff for a transient blip / restarting GPU.
    last_exc: Exception | None = None
    for attempt in range(MI50_ATTEMPTS):
        try:
            return _call(backend)
        except Exception as exc:
            last_exc = exc
            logbus.log("debug", "summary retry", attempt=attempt + 1,
                       backend=backend, error=str(exc)[:80])
            if attempt < MI50_ATTEMPTS - 1:
                time.sleep(5 * (attempt + 1))

    # Primary exhausted. If it wasn't already cloud and cloud is configured, fall
    # back once so a stuck/offline MI50 doesn't sink consolidation for the night.
    if backend != "cloud" and config.load().openai_api_key:
        logbus.log("info", "summary fell back to cloud", primary=backend,
                   error=str(last_exc)[:80] if last_exc else None)
        return _call("cloud")

    raise last_exc if last_exc else RuntimeError("summary failed")


def _summarize_transcript(transcript: str, backend: Backend) -> str:
    """Transcript -> gist (LLM only, no DB). Chunks + merges if oversized, and
    recurses so even the merged partials never exceed the backend's window."""
    budget = _budget(backend)
    if len(transcript) <= budget:
        return _summarize_text(transcript, backend)
    partials = [_summarize_text(c, backend) for c in _chunk(transcript, budget)]
    merged = "Partial summaries to merge:\n\n" + "\n\n".join(partials)
    return _summarize_transcript(merged, backend)


def summarize_session(session_id: str, backend: Backend | None = None) -> str | None:
    """(Re)generate and store the gist for a session. Returns the summary text."""
    exchanges = memory.history(session_id)
    if not exchanges:
        return None
    backend = backend or config.load().summary_backend
    gist = _summarize_transcript(_transcript(exchanges), backend)
    memory.store_summary(session_id, gist, exchanges[-1].id)
    logbus.log("info", "summarized session", session=session_id, exchanges=len(exchanges))
    return gist


def maybe_summarize(session_id: str, backend: Backend | None = None) -> None:
    """Summarize the session if enough new turns have accumulated since last time."""
    if memory.unsummarized_count(session_id) >= SUMMARIZE_AFTER:
        summarize_session(session_id, backend=backend)


_inflight: set[str] = set()
_inflight_lock = threading.Lock()


def maybe_summarize_async(session_id: str, backend: Backend | None = None) -> None:
    """Run maybe_summarize off the chat turn's critical path. Consolidation is
    background maintenance — it must never stall the reply or surface an error to
    the user (a slow/oversized local model would otherwise block the turn). At most
    one summary per session runs at a time."""
    with _inflight_lock:
        if session_id in _inflight:
            return
        _inflight.add(session_id)

    def _run() -> None:
        try:
            maybe_summarize(session_id, backend=backend)
        except Exception as exc:
            logbus.log("error", "summary skipped", session=session_id, error=str(exc)[:120])
        finally:
            with _inflight_lock:
                _inflight.discard(session_id)

    threading.Thread(target=_run, daemon=True, name="summarize").start()


def summarize_all(
    backend: Backend | None = None, limit: int | None = None, workers: int | None = None
) -> dict:
    """Summarize every session that needs it. Idempotent and resumable.

    Concurrency is backend-aware: the cloud API parallelizes happily, but the
    local/MI50 GPU servers run a single slot (llama.cpp --parallel 1) — firing N
    requests at them just queues, blows the client timeout, and thrashes the KV
    cache (wasted compute + heat). So GPU backends run serially unless overridden.
    DB reads/writes (store_summary embeds) stay on the main thread, so the single
    SQLite connection is never touched from multiple threads.
    """
    backend = backend or config.load().summary_backend
    if workers is None:
        workers = 8 if backend == "cloud" else 1

    # Main thread: collect the work (transcripts) for sessions needing a summary.
    todo: list[tuple[str, str, int]] = []
    for s in memory.list_sessions():
        sid = s["id"]
        if memory.get_summary(sid) and memory.unsummarized_count(sid) == 0:
            continue
        exchanges = memory.history(sid)
        if not exchanges:
            continue
        todo.append((sid, _transcript(exchanges), exchanges[-1].id))
        if limit is not None and len(todo) >= limit:
            break

    done, failed = 0, 0
    logbus.log("info", "summarize-all starting", todo=len(todo), backend=backend, workers=workers)

    def work(item: tuple[str, str, int]) -> tuple[str, str, int]:
        sid, transcript, last_id = item
        return sid, _summarize_transcript(transcript, backend), last_id

    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(work, item): item for item in todo}
        for fut in as_completed(futures):
            sid = futures[fut][0]
            try:
                _, gist, last_id = fut.result()
                memory.store_summary(sid, gist, last_id)  # main thread: embed + write
                done += 1
            except Exception as exc:
                failed += 1
                logbus.log("error", "summarize failed", session=sid, error=str(exc)[:120])
            if (done + failed) % 25 == 0:
                logbus.log("info", "summarize-all progress", done=done, failed=failed, total=len(todo))

    report = {"summarized": done, "failed": failed, "total": len(todo)}
    logbus.log("info", "summarize-all complete", **report)
    return report


def main() -> int:
    limit = int(sys.argv[1]) if len(sys.argv) > 1 else None
    print(summarize_all(limit=limit))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())