context_engine.py

"""
RUBRA — context_engine.py
Context Window Management & Session Memory Layer.

THE PROBLEM THIS REPLACES:
main.py previously did `hist = raw_hist[-30:]` — a flat, fixed slice with
no idea how many tokens that actually costs, no regard for which model is
about to receive it, and total amnesia for anything older than 15
exchanges (it's just gone, not summarized, not searchable). Separately,
RUBRA already had THREE overlapping memory paths that didn't know about
each other: `database.py` messages table (raw log), `database.py`
user_memory + `memory_engine.py` (durable cross-session facts), and
`evolution_engine.py`'s InfiniteMemory JSON blob (episodic/semantic/
personal — reloaded and rewritten to disk in FULL on every single write,
which doesn't scale). This module doesn't replace those — it sits above
them as the thing that actually DECIDES what goes into the prompt, and
introduces the piece that was fully missing: a token budget.

DESIGN — four tiers, assembled in priority order until the budget runs out:
  Tier 0  SYSTEM   — persona/rules/persistent facts. Always included,
                     never trimmed (caller's responsibility to keep this
                     tier itself reasonably sized).
  Tier 1  PINNED   — RAG hits / tool context / uploaded-file context for
                     THIS turn. High relevance, bounded size.
  Tier 2  RECENT   — verbatim recent turns, walked newest-first until the
                     remaining budget runs out. This is what "the model
                     actually sees as conversation" — exact wording
                     preserved, since paraphrasing recent turns loses
                     exactly the nuance multi-turn coherence depends on.
  Tier 3  SUMMARY  — a single rolling compacted paragraph covering
                     everything OLDER than what fit in Tier 2. Built
                     INCREMENTALLY (only the newly-pushed-out slice gets
                     summarized each time, folded into the existing
                     summary) so a long-running session never pays an
                     O(n) summarization cost on every turn.

Token counts are estimated (chars/4, tuned per empirical average for
mixed English/Bengali text) rather than an exact tokenizer — RUBRA talks
to several providers (Groq/OpenRouter/Gemini/Cerebras) each with their own
tokenizer, so no single exact count would be right for all of them anyway;
a slightly conservative heuristic with a safety margin is the correct
choice here, not a shortcut.
"""

import re
import logging
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Callable, AsyncIterator

from database import summary_load, summary_save

log = logging.getLogger("rubra.context")

# ═══════════════════════════════════════════════════════
#  TOKEN BUDGETS
# ═══════════════════════════════════════════════════════
# Raised substantially (2026-07, explicit instruction) — the models behind
# these cascades now include Nemotron 3 Super (1M real context) and
# Nemotron 3 Nano Omni (256K-300K), and brain.py no longer caps output
# length at all (max_tokens removed from every cascade entry — see
# brain.py's stream_llm). These numbers are still well under each model's
# real ceiling, not "as large as physically possible": the input side
# still needs SOME budget discipline (a model can only use its real
# context window once, shared between what we send in and what it writes
# back out), but there's much more room now than when every model here
# topped out around 32K-128K.
MODEL_CONTEXT_BUDGETS = {
    "coding":     80_000,
    "diff_merge": 150_000,  # Nemotron 3 Super primary here — 1M real context, this is still a small slice of it
    "general":    32_000,
    "vision":     60_000,   # Nemotron 3 Nano Omni primary here — 256K-300K real context
}
DEFAULT_BUDGET = 32_000

# Split of the budget across tiers once Tier 0 (system) is subtracted.
PINNED_TIER_FRACTION  = 0.25   # RAG/tool/file context for THIS turn
SUMMARY_TIER_MAX_CHARS = 3_200  # ~800 tokens ceiling on the rolling summary itself

# How many NEW raw messages must have accumulated past the last summary
# checkpoint before we bother running another summarization call — avoids
# re-summarizing a delta of 1-2 messages every single turn.
SUMMARIZE_TRIGGER_DELTA = 6

CHARS_PER_TOKEN = 4.0  # heuristic — see module docstring


def estimate_tokens(text: str) -> int:
    """Heuristic token estimate. Always rounds up — better to
    under-fill the budget slightly than to overflow it."""
    if not text:
        return 0
    return int(len(text) / CHARS_PER_TOKEN) + 1


def estimate_message_tokens(msg: Dict) -> int:
    content = msg.get("content", "") or ""
    # +4 for role/formatting overhead per message, matching the rough
    # per-message padding most chat-completion tokenizers add.
    return estimate_tokens(content) + 4


@dataclass
class ContextWindow:
    """What context_engine hands back — main.py/agent.py consume this
    directly instead of hand-rolling their own history slice."""
    recent_hist:      List[Dict] = field(default_factory=list)   # verbatim, chronological
    summary_text:      str = ""                                   # Tier 3, empty if nothing summarized yet
    dropped_count:     int = 0                                    # messages neither shown nor yet summarized (should be 0 in steady state)
    budget_tokens:     int = 0
    used_tokens:       int = 0
    tiers_used:        List[str] = field(default_factory=list)


def _default_summary_prompt(old_messages: List[Dict], prior_summary: str) -> str:
    convo = "\n".join(
        f"{m.get('role','user')}: {(m.get('content') or '')[:400]}"
        for m in old_messages
    )
    prior = f"EXISTING SUMMARY OF EVEN OLDER CONTEXT:\n{prior_summary}\n\n" if prior_summary else ""
    return (
        f"{prior}"
        f"Summarize the following older conversation turns into a SHORT "
        f"(under 150 words) running summary that preserves: who the user "
        f"is, what they're working on, decisions already made, and any "
        f"open threads. Merge with the existing summary above if present "
        f"— produce ONE updated summary, not two separate ones. Write it "
        f"as plain prose, no headers, no bullet points.\n\n"
        f"OLDER TURNS TO FOLD IN:\n{convo}\n\n"
        f"Respond with ONLY the updated summary text."
    )


async def _run_summarizer(llm_func: Callable, old_messages: List[Dict], prior_summary: str) -> str:
    """llm_func follows the same convention used everywhere else in RUBRA
    (memory_engine.extract_and_update_memory, code_executor's autofix loop):
    an async generator over (messages, mode=...) yielding text tokens."""
    prompt = _default_summary_prompt(old_messages, prior_summary)
    out = ""
    try:
        async for tok in llm_func([{"role": "user", "content": prompt}], mode="fast"):
            out += tok
    except Exception as e:
        log.warning(f"[CONTEXT] summarization call failed: {e}")
        return prior_summary  # fail soft — keep the old summary rather than lose it
    out = out.strip()
    return out[:SUMMARY_TIER_MAX_CHARS] if out else prior_summary


async def build_context_window(
    session_id: str,
    raw_hist: List[Dict],
    pinned_context: str = "",
    system_tokens_used: int = 0,
    budget_key: str = "general",
    llm_func: Optional[Callable] = None,
) -> ContextWindow:
    """
    The main entry point. Replaces `hist[-30:]`.

    Args:
        session_id:          for summary checkpoint lookup/storage
        raw_hist:             FULL message history for this session, chronological
        pinned_context:       already-built Tier 1 string (RAG hits, tool
                               output, file context) — sized by the caller,
                               just used here for budget accounting
        system_tokens_used:   estimated tokens already spent on Tier 0
                               (system prompt + persona + persistent facts)
        budget_key:           which MODEL_CONTEXT_BUDGETS entry to use
        llm_func:             optional async-generator LLM call used to
                               produce/update the rolling summary. If not
                               provided, older turns beyond the verbatim
                               window are simply dropped (old behavior) —
                               summarization is a strict improvement, not
                               a hard requirement, so this degrades
                               gracefully rather than raising.

    Returns a ContextWindow with recent_hist (verbatim, put straight into
    the messages list) and summary_text (inject as a single system_note,
    same pattern main.py already uses for memory_ctx/file_ctx).
    """
    total_budget = MODEL_CONTEXT_BUDGETS.get(budget_key, DEFAULT_BUDGET)
    remaining = max(total_budget - system_tokens_used, 1000)  # floor so a huge system prompt can't zero this out

    pinned_tokens = estimate_tokens(pinned_context)
    pinned_cap = int(total_budget * PINNED_TIER_FRACTION)
    if pinned_tokens > pinned_cap:
        # Caller-supplied pinned context is itself too big — this function
        # doesn't truncate it (that'd cut mid-fact), it just accounts for
        # the real size so Tier 2 doesn't over-allocate against a budget
        # that pinned context already blew past.
        remaining -= pinned_tokens
    else:
        remaining -= pinned_tokens
    remaining = max(remaining, 500)

    # ── Tier 2: walk newest-first, keep everything that fits ──────────
    kept: List[Dict] = []
    used = 0
    cutoff_index = len(raw_hist)  # index into raw_hist where verbatim coverage starts
    for i in range(len(raw_hist) - 1, -1, -1):
        m = raw_hist[i]
        cost = estimate_message_tokens(m)
        if used + cost > remaining and kept:
            # Stop BEFORE this message, but only once we already have at
            # least one message kept — always keep the most recent turn
            # even if it alone blows the budget (better to slightly
            # overflow than send zero context).
            cutoff_index = i + 1
            break
        kept.append(m)
        used += cost
        cutoff_index = i
    kept.reverse()

    tiers_used = ["recent"] if kept else []
    if pinned_context:
        tiers_used.insert(0, "pinned")

    # ── Tier 3: incremental summary of everything before cutoff_index ──
    summary_text = ""
    dropped_count = 0
    older = raw_hist[:cutoff_index]
    if older:
        checkpoint = summary_load(session_id)
        prior_summary = checkpoint["summary"]
        through_index = checkpoint["through_index"]

        # Only the slice we haven't summarized yet.
        new_slice = older[through_index:] if through_index < len(older) else []

        if new_slice and len(new_slice) >= SUMMARIZE_TRIGGER_DELTA and llm_func is not None:
            summary_text = await _run_summarizer(llm_func, new_slice, prior_summary)
            try:
                summary_save(session_id, summary_text, len(older))
            except Exception as e:
                log.warning(f"[CONTEXT] summary_save failed for {session_id}: {e}")
        elif prior_summary:
            # Nothing new worth summarizing yet, but we still have a
            # perfectly good summary of everything older than that —
            # keep surfacing it rather than going silent on that history.
            summary_text = prior_summary
            dropped_count = len(new_slice)  # not yet folded in, not shown verbatim either
        else:
            # No summarizer available and nothing summarized yet — this is
            # the graceful-degradation path; older turns are simply not
            # in this prompt, exactly like the old `hist[-30:]` behavior,
            # but at least tracked/visible via dropped_count for logging.
            dropped_count = len(new_slice)

        if summary_text:
            tiers_used.append("summary")

    return ContextWindow(
        recent_hist=kept,
        summary_text=summary_text,
        dropped_count=dropped_count,
        budget_tokens=total_budget,
        used_tokens=system_tokens_used + pinned_tokens + used,
        tiers_used=tiers_used,
    )


def format_summary_for_prompt(summary_text: str) -> str:
    """Same system_note convention main.py already uses for memory_ctx —
    caller does: `hist = [{"role": "system_note", "_summary_ctx": ...}] + hist`."""
    if not summary_text:
        return ""
    return (
        "[EARLIER IN THIS CONVERSATION — summarized, not verbatim]\n"
        f"{summary_text}\n"
        "[END SUMMARY — the messages below this are the exact recent turns]"
    )