"""Partitioning: decide which messages a compression pass summarizes and which it keeps. Pure functions over a message list — no LLM call, no storage. `compress_context` and `ContextCompressor.would_compress` both run `_plan_compression`, so the prediction and the real pass can never drift apart. """ import re from collections.abc import Callable from navi.config import settings from navi.llm.base import Message _CRITICAL_TOOL_NAMES = frozenset({ "filesystem", "code_exec", "terminal", "ssh_exec", }) # Content markers that make a turn worth preserving verbatim longer. _CRITICAL_PATTERNS = [ re.compile(r"\b(error|exception|traceback|failed|failure)\b", re.IGNORECASE), re.compile(r"\b(user\s+said|no,\s+|don't\s+|do\s+not\s+|never\s+|instead\s+|wrong\s+|incorrect\s+|fix\s+|correct\s+)\b", re.IGNORECASE), re.compile(r"\b(edited|modified|created|deleted|wrote|added)\s+(file|function|class|method)\b", re.IGNORECASE), ] def should_compress(context_tokens: int, max_context_tokens: int, threshold: float) -> bool: return context_tokens >= int(max_context_tokens * threshold) def _turn_importance(turn: list[Message]) -> int: """Score a turn for adaptive keep_recent. Higher = more important to keep.""" score = 0 text = "\n".join((m.content or "") for m in turn) lowered = text.lower() # Strong signals: user corrections and explicit negatives if any(w in lowered for w in ("wrong", "incorrect", "fix", "don't use", "do not use", "instead use", "change to")): score += 3 for pattern in _CRITICAL_PATTERNS: score += len(pattern.findall(text)) for m in turn: if getattr(m, "is_compression_critical", False): score += 3 if m.role == "tool" and m.name in _CRITICAL_TOOL_NAMES: score += 1 if m.role == "user" and len((m.content or "").strip()) <= 20: # Very short user messages are usually social/filler; deprioritize score -= 2 return max(0, score) def _order_like(reference: list[Message], subset: list[Message]) -> list[Message]: """Reorder `subset` into the order its members appear in `reference`.""" wanted = {id(m) for m in subset} return [m for m in reference if id(m) in wanted] def _split_turn_tail( turn: list[Message], keep_recent_messages: int, *, budget_tokens: int | None = None, estimate: Callable[[list[Message]], int] | None = None, ) -> tuple[list[Message], list[Message], bool] | None: """Split an in-flight turn into (dropped_middle, kept_tail, trimmed_by_budget). The kept shape mirrors the midturn partition: the user request head plus the newest ``keep_recent_messages`` messages verbatim. ``budget_tokens`` bounds the kept region (head + tail) in the caller's units. The count window alone is useless once messages are huge: twelve messages at a sixth of the window each are two windows. With a budget the oldest tail messages move into the summary until the kept region fits (the newest message and the user head are always kept — dropping the in-flight request or its latest tool result is worse than overrunning). Returns None when nothing has to be dropped, so callers keep their "nothing to compress" semantics. """ if not turn: return None head_len = 1 if turn[0].role == "user" else 0 count_start = max(head_len, len(turn) - keep_recent_messages) tail_start = count_start trimmed = False if budget_tokens is not None and estimate is not None: head = turn[:head_len] while ( tail_start < len(turn) - 1 and estimate(head + turn[tail_start:]) > budget_tokens ): tail_start += 1 trimmed = True dropped = turn[head_len:tail_start] if dropped: # Never start the verbatim tail on a tool result whose call moved into # the summary: keep the assistant tool_call beside its results (an # orphaned tool message is invalid for chat-completion APIs). while tail_start > head_len and turn[tail_start].role == "tool": tail_start -= 1 dropped = turn[head_len:tail_start] if not dropped: return None return dropped, turn[:head_len] + turn[tail_start:], trimmed def partition_messages( messages: list[Message], keep_recent: int, keep_recent_messages: int | None = None, adaptive: bool = True, *, budget_tokens: int | None = None, estimate: Callable[[list[Message]], int] | None = None, ) -> tuple[list[Message], list[Message]]: """ Returns (to_summarize, to_keep). Keeps the system message and the last `keep_recent` conversational turns verbatim. Everything older goes into to_summarize. Tool call groups (assistant + tool results) always stay together. When adaptive=True, important turns (user corrections, errors, critical tools) are kept longer and social/filler turns are compressed sooner. ``budget_tokens``/``estimate`` bound the verbatim in-flight turn by size; see ``_split_turn_tail``. Summary messages (is_summary=True) are not part of any turn: they carry already-compressed history, not a new user request. They always go into to_summarize so the fresh summary replaces them (the single-summary invariant) — an old summary can never act as a turn's user-request head (which the midturn split would keep verbatim beside the fresh summary) or be rescued by the adaptive swap. This is what used to make repeated midturn compressions pile stale summary copies into the context. """ non_system = [m for m in messages if m.role != "system"] summary_msgs = [m for m in non_system if m.is_summary] convo = [m for m in non_system if not m.is_summary] # Group into turns: each turn starts with a user message. A non-user # message with no open turn (can only follow a stripped summary) trails # the user request that preceded it, so it attaches to the previous turn # instead of forming an orphan turn without a head. turns: list[list[Message]] = [] for msg in convo: if msg.role == "user" or not turns: turns.append([msg]) else: turns[-1].append(msg) if len(turns) <= keep_recent: if keep_recent_messages is not None: intra_turn = partition_current_turn_messages( turns, keep_recent_messages, budget_tokens=budget_tokens, estimate=estimate, ) if intra_turn is not None: old, to_keep = intra_turn # A compression pass replaces every old summary, never keeps # one beside the fresh one (single-summary invariant). return _order_like(non_system, old + summary_msgs), to_keep if summary_msgs: # Even when nothing else has aged out, summaries are never kept: # two or more of them must consolidate (the meta-summary path in # compress_context) instead of sitting in the context together. return _order_like(non_system, summary_msgs), convo return [], non_system # nothing old enough to compress # Adaptive: pull important older turns into the kept region and push # unimportant recent/filler turns out for summarization. base_keep = keep_recent recent_turns = turns[-base_keep:] old_turns = turns[:-base_keep] # Adaptive swap rescues important old turns in multi-turn conversations. It is # disabled in midturn mode (keep_recent_messages set): there the priority is to # compress the in-flight turn predictably (see the split below), and the swap # could otherwise move a filler in-flight turn into old_turns and summarize it # whole — losing the current request. No test exercises the swap with # keep_recent_messages set, so this guard is safe. if adaptive and keep_recent_messages is None: # Identify important old turns that should not be lost. important_old = [t for t in old_turns if _turn_importance(t) > 0] # Identify filler turns in the recent window that can be swapped out. filler_recent = [t for t in recent_turns if _turn_importance(t) == 0] swaps = min(len(important_old), len(filler_recent)) for i in range(swaps): # Replace the oldest filler in recent with the most important old turn. recent_turns[recent_turns.index(filler_recent[i])] = important_old[-(i + 1)] # Re-sort kept turns by original position so context order stays chronological. kept_set = {id(t) for t in recent_turns} recent_turns = [t for t in turns if id(t) in kept_set] old_turns = [t for t in turns if id(t) not in kept_set] to_summarize = [m for turn in old_turns for m in turn] # Midturn mode (keep_recent_messages set): the in-flight (last) turn can be very # long (many tool iterations). When there are already > keep_recent turns, the # turn-based keep window holds the whole in-flight turn verbatim, so midturn # compression barely shrinks the context. Split the in-flight turn too: keep its # head (the user request) + the newest keep_recent_messages verbatim, summarize # the middle — on top of compressing older turns. Mirrors the head/tail shape of # partition_current_turn_messages. Without this, the in-flight turn stays verbatim # whenever turns > keep_recent, which is the 150->141 shallow-compression symptom. if keep_recent_messages is not None and recent_turns: split = _split_turn_tail( recent_turns[-1], keep_recent_messages, budget_tokens=budget_tokens, estimate=estimate, ) if split is not None: dropped, kept, _ = split to_summarize = to_summarize + dropped to_keep = [m for turn in recent_turns[:-1] for m in turn] + kept else: to_keep = [m for turn in recent_turns for m in turn] else: to_keep = [m for turn in recent_turns for m in turn] # Fold every old summary into this pass; to_summarize is emitted in # chronological order so the summarizer sees history in sequence. to_summarize = _order_like(non_system, to_summarize + summary_msgs) return to_summarize, to_keep def partition_current_turn_messages( turns: list[list[Message]], keep_recent_messages: int, *, budget_tokens: int | None = None, estimate: Callable[[list[Message]], int] | None = None, ) -> tuple[list[Message], list[Message]] | None: """ Mid-turn fallback for long autonomous tool loops. A long chain of assistant/tool iterations after one user message is one conversational turn, so turn-based compression may have nothing to compress. Keep the current user request and the newest messages verbatim, then summarize older messages from the same in-flight turn. ``budget_tokens`` also trims a turn made of a few huge messages, which the message-count window alone cannot split (see ``_split_turn_tail``). """ if not turns: return None split = _split_turn_tail( turns[-1], keep_recent_messages, budget_tokens=budget_tokens, estimate=estimate, ) if split is None: return None dropped, kept, trimmed = split to_summarize = [m for turn in turns[:-1] for m in turn] + dropped to_keep = kept # A budget-driven trim may leave a single message to summarize; that message # is exactly the oversized one the count window could not address. if len(to_summarize) < (1 if trimmed else 2): return None return to_summarize, to_keep def _plan_compression( context: list[Message], keep_recent: int, keep_recent_messages: int | None = None, *, window: int | None = None, estimate: Callable[[list[Message]], int] | None = None, ) -> tuple[list[Message], list[Message], list[Message]] | None: """Decide what to summarize and what to keep — the single decision point shared by ``compress_context`` (real compression) and ``ContextCompressor.would_compress`` (dry-run prediction), so the two can never drift apart. Applies, in order: 1. Turn-based partition (``partition_messages``). 2. Target hysteresis (turn-based mode only): shrink the verbatim keep-recent window until the kept region fits the target fraction of the context window. A fixed keep_recent leaves roughly the same context size after compression (just below the 90% trigger), so the next few messages re-trigger it — the "125 -> 104, then 104 again after a couple of messages" cycle. Dropping the kept size to the 65% target leaves real headroom and folds the extra turns into the summary (same single LLM call, no extra round-trip). Midturn mode (keep_recent_messages set) already splits the in-flight turn aggressively, so a post-check in compress_session truncates if it is still over target. 3. Aggressive intra-turn fallback (midturn mode only): when the turn-based partition found nothing, retry with keep_recent_messages=2. Returns (system_msgs, to_summarize, to_keep) or None when there is nothing substantial to compress. """ system_msgs = [m for m in context if m.role == "system"] if estimate is None: # Imported here, not at module level: the default estimator lives with # the class, and the class imports this module. Callers that care about # units pass their calibrated estimate anyway. from .compressor import ContextCompressor estimate = ContextCompressor.estimate_context_tokens est = estimate target_tokens = int( (window or settings.ollama_num_ctx) * settings.context_compression_target ) # What the verbatim kept region may take: the target minus the system # messages. Passed down so the intra-turn split also works when a few huge # messages put the kept region over the target despite a small message count. keep_budget = max(1, target_tokens - est(system_msgs)) to_summarize, to_keep = partition_messages( context, keep_recent, keep_recent_messages=keep_recent_messages, budget_tokens=keep_budget, estimate=est, ) if keep_recent_messages is None: kept_tokens = est(system_msgs + to_keep) _KEEP_RECENT_FLOOR = 2 while kept_tokens > target_tokens and keep_recent > _KEEP_RECENT_FLOOR: keep_recent -= 1 to_summarize, to_keep = partition_messages( context, keep_recent, keep_recent_messages=keep_recent_messages, budget_tokens=keep_budget, estimate=est, ) kept_tokens = est(system_msgs + to_keep) if len(to_summarize) < 2 and keep_recent_messages is not None and keep_recent_messages > 2: to_summarize, to_keep = partition_messages( context, keep_recent, keep_recent_messages=2, budget_tokens=keep_budget, estimate=est, ) if len(to_summarize) < 2: return None return system_msgs, to_summarize, to_keep