""" 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]" )