diff --git a/navi/core/compression/__init__.py b/navi/core/compression/__init__.py new file mode 100644 index 0000000..05b48cc --- /dev/null +++ b/navi/core/compression/__init__.py @@ -0,0 +1,27 @@ +"""Context compression — partition history, summarize it, keep the context in budget. + +The implementation is split by responsibility: + +* `partition` — pure decisions about what to summarize and what to keep. +* `summarize` — formatting the summarizer prompt, the LLM call, the KEYS line. +* `compressor` — `ContextCompressor`: retry/truncate safety nets and persistence. + +`navi.core.compressor` remains as a compatibility shim; new code should import +from here. +""" + +from .compressor import ContextCompressor +from .partition import ( + partition_current_turn_messages, + partition_messages, + should_compress, +) +from .summarize import compress_context + +__all__ = [ + "ContextCompressor", + "compress_context", + "partition_current_turn_messages", + "partition_messages", + "should_compress", +] diff --git a/navi/core/compression/compressor.py b/navi/core/compression/compressor.py new file mode 100644 index 0000000..444b045 --- /dev/null +++ b/navi/core/compression/compressor.py @@ -0,0 +1,694 @@ +"""`ContextCompressor` — retry, truncate fallback and persistence around `compress_context`.""" + +from collections.abc import Callable +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +import structlog + +from navi.config import settings +from navi.llm.base import LLMBackend, Message +from navi.memory._chunks import build_chunks +from navi.memory.breadcrumbs import MAX_BREADCRUMBS, extract_breadcrumbs + +from ..context_builder import truncate_oversized +from ..events import ContextCompressed +from ..image_tokens import context_image_tokens +from ..token_calibration import calibration_for +from .partition import _plan_compression, _split_turn_tail, should_compress +from .summarize import compress_context + +if TYPE_CHECKING: + from navi.profiles.base import AgentProfile + +log = structlog.get_logger() + +# A call that reports fewer real prompt tokens than this share of our own +# estimate of the prompt was sent a truncated prompt (Ollama keeps roughly +# half the window and evaluates that). The estimate errs high on images, so +# the band is deliberately wide — a false positive costs one extra +# compression pass, a missed one costs the session. +_TRUNCATION_RATIO = 0.7 + +# Hard-truncate fallbacks keep at most this fraction of the context window's +# tokens of recent messages (the rest is dropped without summarizing). +_HARD_TRUNCATE_TOKEN_FRAC = 0.5 + + +class ContextCompressor: + """High-level context compression with retry strategy and hard-truncate fallback. + + Thin wrapper around `compress_context` that adds: + 1. Retry with keep_recent + 4 on LLM failure. + 2. Hard-truncate fallback (drop oldest messages without summarizing). + """ + + @staticmethod + def context_chars(context: list[Message]) -> int: + """Characters of message text — the input to every token estimate.""" + return sum(len(m.content or "") for m in context) + + @staticmethod + def estimate_context_tokens(context: list[Message]) -> int: + """Raw chars // 3 heuristic, no calibration. + + Kept as the reference unit for tests and for callers that want the + old, deliberately under-counting number. Window decisions must use + ``estimate_tokens`` instead — on prod this heuristic landed at ~56k + where the real prompt was 101k, and targets expressed in it were met + while the real prompt sat above the window. + """ + chars = ContextCompressor.context_chars(context) + return chars // 3 + context_image_tokens(context) + + @staticmethod + def estimate_tokens(context: list[Message], model=None) -> int: + """Token estimate in real-prompt units, used by every window decision. + + ``estimate_context_tokens`` converted to real prompt tokens by the + calibration fitted for ``model`` (see ``navi/core/token_calibration.py`` + for why chars // 3 is not usable as-is). Without observations the + calibration falls back to ``DEFAULT_CHARS_PER_TOKEN``. + """ + chars = ContextCompressor.context_chars(context) + return calibration_for(model).tokens_for(chars) + context_image_tokens(context) + + @staticmethod + def estimate_tokens_delta(context: list[Message], model=None) -> int: + """Price messages *appended since* a real baseline, same units. + + Uses the calibration slope without its intercept: the fixed prompt + overhead is already part of the baseline count. Pricing a delta with + the raw ``chars // 3`` left the pre/mid-turn gates in one unit and the + window in another for exactly the tool-heavy turns they exist for. + """ + chars = ContextCompressor.context_chars(context) + return calibration_for(model).delta_for(chars) + context_image_tokens(context) + + def __init__(self) -> None: + self._profile: "AgentProfile | None" = None + # Real prompt-token baseline from the last LLM call: + # (len(session.context) at call time, real prompt_tokens returned). + # The next estimate reuses these real tokens for the bulk of the context + # and only heuristics the messages appended since, instead of the + # chars//3 estimate which undercounts code-heavy tool output and fires + # compression too late. Cleared after compression (context shrank). + self._real_baseline: tuple[int, int] | None = None + # Set when the last call reported far fewer tokens than we sent: the + # backend truncated the prompt. Consumed by the post-turn worker, which + # compresses unconditionally in that case (see record_real_baseline). + self._context_truncated = False + + def set_profile(self, profile: "AgentProfile | None") -> None: + """Tell the compressor which profile is active so it can use profile-specific settings.""" + self._profile = profile + + @property + def context_truncated(self) -> bool: + """True when the last recorded call saw its prompt cut by the backend.""" + return self._context_truncated + + def pop_context_truncated(self) -> bool: + """Read and clear the truncation flag (the post-turn worker consumes it).""" + value = self._context_truncated + self._context_truncated = False + return value + + def record_real_baseline( + self, + context_len: int, + real_tokens: int | None, + *, + chars: int | None = None, + model=None, + sent_tokens: int | None = None, + image_tokens: int = 0, + ) -> None: + """Cache the real prompt-token count returned by the last LLM call. + + ``context_len`` is ``len(session.context)`` at the time of that call + (before the new assistant/tool messages are appended). The next + estimate reuses these real tokens for the bulk of the context and only + heuristics the new messages appended since. + + ``chars`` is the character count of the exact prompt that was sent + (the built context, not just session.context). Together with + ``real_tokens`` it feeds the per-model calibration line, which is what + converts hypothetical contexts (the proposed kept region) into real + tokens. Prompts that carried images are kept out of that fit: image + bytes are not in ``chars``, so their tokens would land in the fitted + intercept/slope and skew every later text estimate for that model. + + ``sent_tokens`` is our estimate of the prompt that was sent. When the + backend reports far fewer tokens than that, the prompt was truncated — + Ollama silently keeps about half the window and reports only what it + evaluated. Caching that number as the baseline used to pin the session + at ~50%: every later estimate started from the truncated count, the + 90% trigger never fired, and the model kept receiving a prompt whose + head (system prompt included) had been dropped. A truncated count is + therefore not recorded; the previous — honest — baseline stays, the + estimate keeps its size, and ``context_truncated`` tells the post-turn + worker to compress regardless of the threshold. + """ + if not real_tokens or real_tokens <= 0: + return + if ( + sent_tokens + and sent_tokens > 0 + and real_tokens < sent_tokens * _TRUNCATION_RATIO + ): + self._context_truncated = True + log.warning( + "compressor.prompt_truncated", + sent_tokens=sent_tokens, + real_tokens=real_tokens, + model=model, + ) + return + self._real_baseline = (context_len, int(real_tokens)) + if chars and not image_tokens: + calibration_for(model).observe(chars, real_tokens) + + def real_baseline_estimate( + self, + session_context: list[Message], + built_ctx: list[Message], + model=None, + ) -> int: + """Estimate context tokens using the real baseline when available. + + Falls back to the calibrated ``estimate_tokens`` on ``built_ctx`` when + no baseline is recorded or it is stale (the session shrank, e.g. right + after compression — ``base_len`` exceeds the current length). Every + branch is in real prompt tokens, the unit the window is measured in. + + Both the delta and the fallback are priced on the **truncated view** + (``truncate_oversized``), i.e. on what ``build()`` actually sends. The + raw stored context may hold a megabyte-long tool result that never + reaches the model; pricing it here made the gates fire for work the + compressor intentionally refuses (a single message cannot be split) + and — worse — made ``check_context_size`` kill the run with "context + too large" over a prompt the model never received. + """ + if self._real_baseline is not None: + base_len, base_tokens = self._real_baseline + if base_len <= len(session_context): + delta = self.estimate_tokens_delta( + truncate_oversized(session_context[base_len:]), model + ) + return base_tokens + delta + return self.estimate_tokens(truncate_oversized(built_ctx), model) + + async def compress_session( + self, + context: list[Message], + llm: LLMBackend, + model: "list[str] | str | None", + temperature: float, + keep_recent: int, + max_tokens: int | None = None, + keep_recent_messages: int | None = None, + window: int | None = None, + ) -> tuple[list[Message], str] | None: + """Compress context with retry + hard-truncate fallback. + + Returns (new_context, summary_text) or None if nothing changed. + Does NOT mutate the session — the caller is responsible for updating + session.context, session.context_token_count, and persisting. + """ + effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent + effective_max_tokens = getattr(self._profile, "compression_max_tokens", None) or max_tokens + window_tokens = window or settings.ollama_num_ctx + # Every size decision in this pass — the keep-region target, the + # post-check, the truncate budgets — must be in real-token units, the + # same units the trigger fired in, AND on the view the model is sent + # (``truncate_oversized``): pricing raw stored messages made a kept + # megabyte-long tool result look like the whole window and drop turns + # the model could have had. Mixed units are what let a compression + # pass finish at "65% of the window" while the real prompt it produced + # was above the window. + est = lambda msgs: self.estimate_tokens(truncate_oversized(msgs), model) + + # Attempt 1: normal compression + try: + result = await compress_context( + context=context, + llm=llm, + model=model, + temperature=temperature, + keep_recent=effective_keep_recent, + max_tokens=effective_max_tokens, + keep_recent_messages=keep_recent_messages, + profile=self._profile, + window=window_tokens, + estimate=est, + ) + except Exception: + # Attempt 2: keep more recent turns verbatim + try: + result = await compress_context( + context=context, + llm=llm, + model=model, + temperature=temperature, + keep_recent=effective_keep_recent + 4, + max_tokens=effective_max_tokens, + keep_recent_messages=(keep_recent_messages + 4) + if keep_recent_messages is not None + else None, + profile=self._profile, + window=window_tokens, + estimate=est, + ) + except Exception: + # Attempt 3: hard-truncate fallback + return self._hard_truncate(context, window=window_tokens, estimate=est) + + if result is None: + # Partition found nothing to summarize (e.g. few very large messages + # in a single turn). If the context is nonetheless over the token + # threshold, fall back to a token-budget hard-truncate so the gate + # firing always results in shrinkage — otherwise the agent runs until + # the window overflows. (The per-message truncation in + # context_builder.build caps single huge messages for the LLM view; + # this is the stored-context safety net.) + if should_compress( + est(context), + window_tokens, + settings.context_compression_threshold, + ): + return self._token_budget_truncate( + context, estimate=est, window=window_tokens + ) + return None + + # Target hysteresis safety net: if the compressed context is still over + # the target fraction (midturn mode kept a huge in-flight turn, or even + # the keep_recent floor could not fit the window), drop oldest kept + # turns by token budget so a compression firing always leaves real + # headroom below the trigger — otherwise the next few messages re-fire + # compression and the context yo-yos at the trigger line. The dropped + # turns were kept verbatim (not summarized) and are lost; this only + # triggers when the verbatim kept region alone exceeds the target. + new_context, summary_text = result + target_tokens = int(window_tokens * settings.context_compression_target) + if est(new_context) > target_tokens: + truncated = self._token_budget_truncate( + new_context, budget_tokens=target_tokens, estimate=est, window=window_tokens + ) + if truncated is not None: + new_context = truncated[0] + return new_context, summary_text + + def would_compress( + self, + context: list[Message], + keep_recent: int, + keep_recent_messages: int | None = None, + threshold: float | None = None, + max_context_tokens: int | None = None, + current_tokens: int | None = None, + model=None, + ) -> bool: + """True if ``compress_session`` would actually shrink the context. + + Runs the same ``_plan_compression`` decision as ``compress_context`` + (no LLM call). Also returns True when the context is over the token + threshold even though partition found nothing — the token-budget + fallback in ``compress_session`` handles that. Used to emit + ``CompressionStarted`` only when compression will really happen + (no misleading "compression" status on a no-op). + """ + effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent + window_tokens = max_context_tokens or settings.ollama_num_ctx + # Same unit as the caller's real-token estimate: the truncated view the + # model is actually sent, not the raw stored messages. + est = lambda msgs: self.estimate_tokens(truncate_oversized(msgs), model) + if _plan_compression( + context, + effective_keep_recent, + keep_recent_messages=keep_recent_messages, + window=window_tokens, + estimate=est, + ) is not None: + return True + # Partition no-op: compression can still shrink via the token-budget + # fallback — but only if dropping oldest turns actually reduces the + # context (a single huge turn that alone exceeds the budget cannot be + # split by message-boundary truncation, and is handled instead by the + # per-message view truncation in context_builder.build). Mirror that + # decision so would_compress and compress_session agree exactly. + # The size compared against the threshold is the caller's real-token + # estimate (`current_tokens`), not a fresh heuristic: the trigger fired + # on real tokens, and judging it again in chars // 3 units made the two + # disagree by ~1.8x. + if threshold is not None and max_context_tokens is not None: + total = current_tokens if current_tokens is not None else est(context) + if should_compress(total, window_tokens, threshold): + return self._token_budget_truncate( + context, estimate=est, window=window_tokens + ) is not None + return False + + def _hard_truncate( + self, + context: list[Message], + window: int | None = None, + estimate: Callable[[list[Message]], int] | None = None, + ) -> tuple[list[Message], str] | None: + """Last-resort fallback: drop oldest non-system messages without summarizing. + + Keeps whole conversational turns so tool call groups are never split. + Bounded by a token budget (half the context window) rather than a fixed + message count, so it still shrinks a context made of a few very large + messages — the fixed-count floor used to no-op on <= 6 messages even + when they were huge. + """ + budget = max(1, int((window or settings.ollama_num_ctx) * _HARD_TRUNCATE_TOKEN_FRAC)) + return self._token_budget_truncate( + context, budget_tokens=budget, estimate=estimate, window=window + ) + + def _token_budget_truncate( + self, + context: list[Message], + budget_tokens: int | None = None, + estimate: Callable[[list[Message]], int] | None = None, + window: int | None = None, + ) -> tuple[list[Message], str] | None: + """Drop oldest non-system messages (by turn groups) until under budget. + + Keeps system messages and the newest conversational turns whole (tool + call groups stay together). Returns None if the context already fits the + budget or there is nothing to drop. ``budget_tokens`` defaults to half + the context window. + """ + est = estimate or ContextCompressor.estimate_context_tokens + + if budget_tokens is None: + budget_tokens = max( + 1, int((window or settings.ollama_num_ctx) * _HARD_TRUNCATE_TOKEN_FRAC) + ) + + system_msgs = [m for m in context if m.role == "system"] + non_system = [m for m in context if m.role != "system"] + if not non_system: + return None + + # Group non-system into turns (user message starts a turn) + turns: list[list[Message]] = [] + current: list[Message] = [] + for msg in non_system: + if msg.role == "user" and current: + turns.append(current) + current = [msg] + else: + current.append(msg) + if current: + turns.append(current) + + # Keep newest turns (whole) until the token budget is exhausted; always + # keep at least the most recent turn so the agent retains the in-flight + # request + its latest tool results. + kept_turns: list[list[Message]] = [] + used = 0 + for turn in reversed(turns): + turn_tokens = est(turn) + if kept_turns and used + turn_tokens > budget_tokens: + break + kept_turns.insert(0, turn) + used += turn_tokens + + if kept_turns and used > budget_tokens: + # The newest turn alone is over budget (a first turn of a few huge + # messages): the loop above always keeps it whole, which used to + # make this fallback a no-op exactly when it was needed. Trim inside + # it — the user head and the newest messages stay verbatim. + newest = kept_turns[-1] + split = _split_turn_tail( + newest, + len(newest), + budget_tokens=budget_tokens, + estimate=est, + ) + if split is not None: + kept_turns[-1] = split[1] + + to_keep = [m for turn in kept_turns for m in turn] + if len(to_keep) == len(non_system): + return None # already within budget + + new_context = system_msgs + to_keep + # dropped > 0 is guaranteed here: to_keep == non_system returns None above. + summary_text = ( + "[Context was too large to summarize. Old messages were truncated to prevent " + "the model from exceeding its context window. Some earlier details may have been lost.]" + ) + return new_context, summary_text + + def check_context_size( + self, + context: list[Message], + session_context: list[Message] | None = None, + model=None, + window: int | None = None, + ) -> None: + """Raise ContextTooLargeError before an LLM call if the context is dangerously large.""" + # Local import: the window and the reserve are read at call time, so a + # runtime settings swap (tests, operator reload) reaches this check. + from navi.config import settings as _settings + from navi.exceptions import ContextTooLargeError + + if not context: + return + + window_tokens = window or _settings.ollama_num_ctx + # A reserve bigger than the window turns `available` negative and makes + # every check fail, however small the context — a 500-token model still + # has to be able to answer. Cap the reserve at a quarter of the window. + output_reserve = min( + _settings.output_reserve_tokens, max(1, window_tokens // 4) + ) + est = lambda msgs: self.estimate_tokens(msgs, model) + # When the caller supplies the raw session context, prefer the real-token + # baseline (real bulk + calibrated delta) over the calibrated estimate; + # both are in the same units as the real counts the window is measured in. + total = ( + self.real_baseline_estimate(session_context, context, model) + if session_context is not None + else est(context) + ) + available = window_tokens - output_reserve + + if total > available: + existing = est(context[:-1]) + new = est(context[-1:]) + remaining = available - existing + raise ContextTooLargeError( + f"Context too large: new content is ~{new:,} estimated tokens, " + f"but only ~{max(0, remaining):,} tokens are available " + f"(window {window_tokens:,}, already used ~{existing:,}, " + f"output_reserve {output_reserve:,}). " + "Split the file into smaller parts or delegate to a subagent." + ) + + async def compress_and_save_session( + self, + session, + session_store, + llm: LLMBackend, + model: str, + temperature: float, + session_id: str, + reason: str, + keep_recent: int, + max_tokens: int | None = None, + keep_recent_messages: int | None = None, + window: int | None = None, + memory_store=None, + ) -> ContextCompressed | None: + """Compresses the session context and persists the changes to the session store.""" + count_before = len(session.context) + window_tokens = window or settings.ollama_num_ctx + result = await self.compress_session( + context=session.context, + llm=llm, + model=model, + temperature=temperature, + keep_recent=keep_recent, + max_tokens=max_tokens, + keep_recent_messages=keep_recent_messages, + window=window_tokens, + ) + if result is None: + return None + + new_context, summary_text = result + + # Mark messages in session.messages as not context if they are no longer in new_context + # and are not system messages. + # Iterate session.context too: after a session reload, context-only rows + # (old summaries — is_display=False) are present in session.context but + # NOT in session.messages, so demoting only the latter left stale + # summaries with is_context=True in the DB — they piled up in the + # context at every reload (the half-the-context-is-summaries symptom). + new_context_ids = {id(m) for m in new_context} + demoted: list[Message] = [] + seen: set[int] = set() + for msg in [*session.messages, *session.context]: + if msg.role == "system" or id(msg) in new_context_ids or id(msg) in seen: + continue + seen.add(id(msg)) + if not msg.is_context: + # Dropped by an earlier compression. Iterating the display + # history surfaces it again on every later pass; indexing it + # again writes the same turns under a new chunk boundary (new + # hash) instead of once, forever. The current summary covers + # it; if it was dropped while indexing was off, that is the + # one turn range this trade-off leaves unindexed. + continue + msg.is_context = False + demoted.append(msg) + + # Rows that live only in session.context — a summary is is_display=False, + # so after a reload it is absent from session.messages — are invisible to + # save(), which updates flags for display history only. Their demotion + # used to stay in memory: every reload brought the stale summary back + # with is_context=True, and one prod session ended up carrying 15 live + # summaries, 80% of its context. Persist those flips explicitly. + in_messages = {id(m) for m in session.messages} + context_only = [ + m.sequence_number + for m in demoted + if m.sequence_number >= 0 and id(m) not in in_messages + ] + if context_only: + await session_store.set_context_flag(session_id, context_only, False) + + # Add the summary message to session.messages if it's not already there + summary_msg = next((m for m in new_context if m.is_summary), None) + if summary_msg and summary_msg not in session.messages: + summary_msg.is_display = False + session.messages.append(summary_msg) + + # Keep the dropped turns searchable before they leave the context for + # good. This must happen while `demoted` is still in hand: after the + # save there is no object that says "these were the turns compression + # removed", only the display history. + await self._index_demoted(memory_store, session, session_id, demoted, summary_msg) + + # Add a system message with is_compression=True and content=summary_text + session.messages.append( + Message( + role="system", + content=summary_text, + is_compression=True, + is_context=False, + created_at=datetime.now(timezone.utc), + ) + ) + + session.context = new_context + # The size the model is actually sent. A megabyte-long tool result stays + # in stored history (the user still sees it) but build() truncates it per + # message before the call — counting the raw message reported 817k tokens + # against a 128k window the model never received. + session.context_token_count = self.estimate_tokens( + truncate_oversized(new_context), model + ) + # The context just shrank, so the recorded real-token baseline is stale. + # Drop it — the next LLM call re-establishes a fresh one. The truncation + # flag goes with it: the prompt that will be sent is small again. + self._real_baseline = None + self._context_truncated = False + await session_store.save(session) + + log.info( + "compressor.compressed", + session_id=session_id, + reason=reason, + messages_before=count_before, + messages_after=len(new_context), + context_tokens=session.context_token_count, + ) + + # Archive old messages if the hot table exceeds the configured window. + if settings.session_messages_window > 0 and session.db_next_sequence > settings.session_messages_window: + threshold = session.db_next_sequence - settings.session_messages_window + archived = await session_store.archive_old_messages(session_id, threshold) + if archived > 0: + session.messages = [m for m in session.messages if m.sequence_number >= threshold] + session.context = [m for m in session.context if m.sequence_number >= threshold] + session.archive_threshold = threshold + + return ContextCompressed( + messages_before=count_before, + messages_after=len(new_context), + summary=summary_text, + context_tokens=session.context_token_count, + max_context_tokens=window_tokens, + ) + + async def _index_demoted( + self, + memory_store, + session, + session_id: str, + demoted: list[Message], + summary_msg: Message | None, + ) -> None: + """Write the turns compression just dropped into the raw-chunk index. + + Append-only: re-running compression on the same history hits the + `(session_id, content_hash)` uniqueness and writes nothing new. + + Breadcrumbs come from the summary's `KEYS:` line and, always, from the + deterministic extractor over the chunk text. The hard-truncate fallback + never calls a summarizer, so its index would otherwise be empty; and the + extractor also catches names the summarizer chose not to list. Every + chunk of this pass carries the union — the keys describe the batch, and + a searcher does not know which chunk of the batch holds the name. + """ + if memory_store is None or not demoted or not settings.memory_search_enabled: + return + try: + chunks = build_chunks(demoted) + if not chunks: + return + summary_keys = [] + if summary_msg is not None: + summary_keys = list( + (summary_msg.metadata or {}).get("breadcrumbs") or [] + ) + keys = _merge_breadcrumbs( + summary_keys, + extract_breadcrumbs( + [c["content"] for c in chunks], + limit=max(0, MAX_BREADCRUMBS - len(summary_keys)), + ), + ) + for chunk in chunks: + chunk["breadcrumbs"] = keys + await memory_store.add_chunks( + chunks, user_id=getattr(session, "user_id", None), session_id=session_id + ) + except Exception: + # Indexing must never break compression — a missing chunk is a lost + # search result, a raised exception here is a lost turn. + log.warning("compressor.chunk_index_failed", session_id=session_id, exc_info=True) + + +def _merge_breadcrumbs(*groups: list[str]) -> list[str]: + merged: list[str] = [] + seen: set[str] = set() + for group in groups: + for item in group: + key = " ".join(str(item).lower().split()) + if not key or key in seen: + continue + seen.add(key) + merged.append(str(item)) + if len(merged) >= MAX_BREADCRUMBS: + return merged + return merged + diff --git a/navi/core/compression/partition.py b/navi/core/compression/partition.py new file mode 100644 index 0000000..6f1c41e --- /dev/null +++ b/navi/core/compression/partition.py @@ -0,0 +1,354 @@ +"""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 + + diff --git a/navi/core/compression/summarize.py b/navi/core/compression/summarize.py new file mode 100644 index 0000000..8a886d1 --- /dev/null +++ b/navi/core/compression/summarize.py @@ -0,0 +1,339 @@ +"""Summary production: format history for the summarizer, call it, parse the KEYS line.""" + +import json +from collections.abc import Callable +from datetime import datetime, timezone +from pathlib import Path +from typing import TYPE_CHECKING + +from navi.config import settings +from navi.llm.base import LLMBackend, Message +from navi.memory.breadcrumbs import split_keys_line + +from ..image_tokens import estimate_image_tokens +from .partition import _CRITICAL_TOOL_NAMES, _plan_compression + +if TYPE_CHECKING: + from navi.profiles.base import AgentProfile + +def _format_for_summary( + messages: list[Message], image_budget: int | None = None +) -> tuple[str, list[str]]: + """ + Render messages as plain text for the summarization prompt. + + Returns (text, images) where images is a flat list of base64 strings + collected from all user messages. Vision-capable models will receive + the images alongside the text; non-vision models silently ignore them. + + ``image_budget`` caps the total cost of the collected images in tokens. + Every image in the history used to be forwarded in one user message with + only the text capped, so a session with a dozen screenshots produced a + summarization prompt that could not fit its own window — the call was + silently truncated by the backend and the summary was written as if the + history had been read. Images over the budget stay out of the call; their + per-message "[+ N image(s)]" note still tells the summarizer they existed. + """ + lines: list[str] = [] + images: list[str] = [] + remaining = image_budget if image_budget is None else max(0, image_budget) + i = 0 + while i < len(messages): + m = messages[i] + + if m.is_summary: + # Existing summary — include as-is (already compressed) + lines.append(m.content or "") + i += 1 + + elif m.role == "user": + if m.images: + for img in m.images: + if remaining is None: + images.append(img) + continue + cost = estimate_image_tokens(img) + if cost <= remaining: + images.append(img) + remaining -= cost + img_note = f" [+ {len(m.images)} image(s)]" + else: + img_note = "" + if m.content: + lines.append(f"User: {m.content}{img_note}") + elif img_note: + lines.append(f"User:{img_note}") + i += 1 + + elif m.role == "assistant" and m.tool_calls: + # Render tool calls + their results as a compact block + if m.content: + lines.append(f"Assistant: {m.content}") + for tc in m.tool_calls: + args_preview = json.dumps(tc.arguments)[:120] + lines.append(f"[Tool call: {tc.name}; arguments preview: {args_preview}]") + i += 1 + while i < len(messages) and messages[i].role == "tool": + tool_msg = messages[i] + result = tool_msg.content or "" + # Critical tool results and explicit critical flag survive verbatim + # up to a larger budget, so exact errors/file contents are preserved. + critical = ( + getattr(tool_msg, "is_compression_critical", False) + or tool_msg.name in _CRITICAL_TOOL_NAMES + ) + if critical and len(result) <= 4000: + preview = result + elif critical: + # Over the critical budget: keep head + tail halves rather + # than collapsing to the non-critical 800-char preview — + # the all-or-nothing cliff used to throw away 98% of a + # large file read (the common navi_code shape) while a + # 3999-char result survived whole. + half = 2000 + preview = result[:half] + "\n…[middle elided]…\n" + result[-half:] + else: + # Give the summarizer enough of the result to preserve + # diagnostic detail (the exact error line, a grep match, + # a key value), not just the first 300 chars — the summary + # is only as detailed as what the summarizer can see. + preview = result[:800] + ("…" if len(result) > 800 else "") + lines.append(f"[Tool result: {tool_msg.name}; preview: {preview}]") + i += 1 + + elif m.role == "assistant" and m.content: + lines.append(f"Assistant: {m.content}") + i += 1 + + else: + i += 1 # skip orphan tool messages + + return "\n".join(lines), images + + +# Safety limit: truncate formatted input to this many characters before sending to LLM. +# Prevents the summarizer from receiving near-context-sized input it can't fit alongside output. +# Typical ollama_num_ctx is 65k–128k, so 32k of input leaves ample room for a 6k-token summary. +_MAX_SUMMARY_INPUT_CHARS = 32_000 + +# Share of the window the historical images may take in a summarization call +# (see _format_for_summary). The text side is capped separately, in chars. +_SUMMARY_IMAGE_BUDGET_FRAC = 0.25 + + +# When existing summaries in to_summarize exceed this many chars combined, +# run a quick meta-summary to consolidate them before the main compression pass. +_META_SUMMARY_THRESHOLD = _MAX_SUMMARY_INPUT_CHARS // 3 # 10_666 + +_META_SUMMARY_SYSTEM = ( + "You are condensing multiple conversation summaries into a single compact summary. " + "Preserve all key facts, decisions, file paths, config values, errors, and user preferences. " + "Eliminate redundancy between overlapping summaries. " + "Write tight bullet points. Do not include filler or meta-commentary." +) + + +async def _meta_summarize( + summaries: list[Message], + llm: LLMBackend, + model, + temperature: float, +) -> Message: + """Consolidate multiple existing summary messages into one compact meta-summary.""" + combined = "\n\n---\n\n".join(m.content for m in summaries) + prompt = [ + Message(role="system", content=_META_SUMMARY_SYSTEM), + Message(role="user", content=combined), + ] + response = await llm.complete( + prompt, + tools=None, + temperature=temperature, + model=model, + think=False, + max_tokens=1500, + ) + text = (response.content or "").strip() or "(consolidated summary unavailable)" + return Message( + role="user", + content=f"[Consolidated Context Summary]\n{text}", + is_summary=True, + is_display=False, + is_context=True, + created_at=datetime.now(timezone.utc), + ) + + +# The section list and the instructions built from it define the summary +# format every compression pass produces. `KEYS` is parsed out by +# `split_keys_line` and becomes the summary's breadcrumbs. +_SUMMARY_SECTIONS = [ + ("Goal", "One clear sentence describing what the user is trying to accomplish in this session. Include deadlines or constraints if stated."), + ("Active Files", "Every file or directory the assistant touched, with absolute or project-relative path, status (created / modified / read / deleted), the purpose of the change, and the key function/class/symbol affected."), + ("Decisions & User Preferences", "Explicit choices, architecture decisions, style preferences, or corrections stated by the user — each with a one-line rationale and what triggered it. Things the user said NOT to do."), + ("Completed Work", "Concrete finished steps — step + file/function name + verification outcome (pass/fail). Be specific and complete; prefer listing over generalizing."), + ("Intermediate Findings", "Key results from read/grep/log/terminal inspection that informed decisions: the fact discovered and the file/section it came from. Drop verbose bulk, keep the diagnostic line that changed understanding."), + ("Pending Work / Todo", "Open tasks, in-progress items, or follow-ups that still need action."), + ("Errors & Blockers", "Failures, exceptions, or unresolved issues — each with the exact short error snippet, what was tried, and the outcome. Preserve the diagnostic line verbatim."), + ("Key Values", "Exact constants the assistant should remember: ports, config keys and values, versions, dependency names, important paths, IDs, exact CLI flags."), +] + +_SUMMARY_TEMPLATE_INSTRUCTIONS = ( + "You are summarizing a conversation history to free up context space. " + "The assistant will continue working using ONLY this summary — it will have no access " + "to the original messages. Be thorough and precise; prefer specifics and complete lists " + "over generalities — the cost of losing a fact is higher than the cost of a longer summary. " + "This summary is historical context, not a new user request.\n\n" + "Use EXACTLY the Markdown structure below. Every section must be present. " + "If a section has no relevant information, write its header and the literal word NONE. " + "Keep bullets information-dense, but do not generalize away specifics to save space. " + "Do not include greetings, filler, transitions, or meta-commentary.\n\n" + + "\n\n".join(f"## {title}\n{desc}" for title, desc in _SUMMARY_SECTIONS) + + "\n\n" + "Output rules:\n" + "- Preserve exact file paths, function/class/symbol names, config keys and values, " + "exact CLI flags, and short error/diagnostic snippets verbatim.\n" + "- Preserve short critical code the agent must reproduce exactly: final function/class " + "signatures, one-line fixes the user approved, key config lines. Do NOT paste long blocks, " + "patches, or full command output — keep only the short line that matters.\n" + "- Do not paraphrase values that must stay precise.\n" + "- Use Markdown headers exactly as shown.\n" + "- End with ONE line `KEYS: item, item, …` — the breadcrumbs of this history: " + "names, commands, file paths, config keys, amounts, ticket numbers, decisions and " + "links that a later turn may ask about by name. Plain items, no explanations, at " + "most 30, most important first. This line is an index, not prose; write " + "`KEYS: NONE` when nothing is worth indexing." +) + + +def _build_summary_system_prompt(profile: "AgentProfile | None") -> str: + """Build the system prompt used by the summarization LLM. + + Uses the profile-specific compression prompt file if configured. + """ + base = _SUMMARY_TEMPLATE_INSTRUCTIONS + extra = "" + if profile is not None and getattr(profile, "compression_prompt_file", None): + profile_dir = Path("navi/profiles") / profile.id + prompt_path = profile_dir / profile.compression_prompt_file + if prompt_path.exists(): + try: + extra = "\n\n---\n\n[Profile-specific compression instructions]\n\n" + prompt_path.read_text(encoding="utf-8").strip() + except Exception: + pass + return base + extra + + + +async def compress_context( + context: list[Message], + llm: LLMBackend, + model: "list[str] | str | None", + temperature: float, + keep_recent: int, + max_tokens: int | None = None, + keep_recent_messages: int | None = None, + profile: "AgentProfile | None" = None, + window: int | None = None, + estimate: Callable[[list[Message]], int] | None = None, +) -> tuple[list[Message], str] | None: + """ + Summarize old messages in the LLM context and return a shorter context list. + Only operates on `context` — the full display history (session.messages) is never touched. + Returns None if there is nothing to compress. + + Images from old user messages are passed to the summarization model. + Vision-capable models will incorporate image descriptions into the summary; + non-vision models silently ignore the images field. + + Uses the same model already loaded in memory (profile.model passed via WorkerContext) — + no model swap, no extra loading overhead. + + Profile settings override global defaults when provided: + - compression_keep_recent -> keep_recent + - compression_max_tokens -> max_tokens + - compression_prompt_file -> appended to summary system prompt + + Exceptions propagate to the caller (CompressionWorker catches them). + """ + effective_keep_recent = getattr(profile, "compression_keep_recent", None) or keep_recent + effective_max_tokens = getattr(profile, "compression_max_tokens", None) or max_tokens + + planned = _plan_compression( + context, + effective_keep_recent, + keep_recent_messages=keep_recent_messages, + window=window, + estimate=estimate, + ) + if planned is None: + return None # nothing substantial to compress + system_msgs, to_summarize, to_keep = planned + + # Meta-summary: if to_summarize contains multiple existing summary messages + # that are long enough to crowd the summarizer input, consolidate them first. + summary_msgs = [m for m in to_summarize if m.is_summary] + if len(summary_msgs) > 1: + total_summary_len = sum(len(m.content or "") for m in summary_msgs) + if total_summary_len > _META_SUMMARY_THRESHOLD: + try: + meta = await _meta_summarize( + summary_msgs, llm, model, temperature + ) + to_summarize = [m for m in to_summarize if not m.is_summary] + to_summarize.insert(0, meta) + except Exception: + # If meta-summary fails, continue with raw summaries as-is + pass + + # Images forwarded to the summarizer share one user message with the + # formatted text. Give them a quarter of the window at most: the summary + # itself needs the rest, and a summarization call that overflows is + # silently truncated by Ollama (to half the window), which would produce a + # summary written from an arbitrary slice of the history. + image_budget = int((window or settings.ollama_num_ctx) * _SUMMARY_IMAGE_BUDGET_FRAC) + summary_text_input, images = _format_for_summary( + to_summarize, image_budget=image_budget + ) + + # Truncate oversized input so the summarizer LLM has room to generate output. + # Keep head + tail, not just the head: to_summarize is chronological, so a + # head-only cut silently dropped the NEWEST summarized messages — the ones + # closest to the kept window, which bridge old context and the recent turn. + if len(summary_text_input) > _MAX_SUMMARY_INPUT_CHARS: + elision = "\n…[middle elided]…\n" + head = _MAX_SUMMARY_INPUT_CHARS * 3 // 4 + tail = _MAX_SUMMARY_INPUT_CHARS - head - len(elision) + summary_text_input = summary_text_input[:head] + elision + summary_text_input[-tail:] + + system_prompt = _build_summary_system_prompt(profile) + prompt = [ + Message(role="system", content=system_prompt), + Message(role="user", content=summary_text_input, images=images or None), + ] + + # think=False: compression must be fast — extended reasoning wastes context and hangs + response = await llm.complete( + prompt, tools=None, temperature=temperature, model=model, think=False, max_tokens=effective_max_tokens + ) + summary_text = (response.content or "").strip() or "(summary unavailable)" + # The summarizer's trailing `KEYS:` line is the breadcrumb index for the + # messages this pass drops. It is stripped from the summary the agent reads + # (it is an index, not prose) and carried on the summary message so the + # caller can write it next to the raw chunks. + summary_text, summary_keys = split_keys_line(summary_text) + if not summary_text.strip(): + summary_text = "(summary unavailable)" + + summary_msg = Message( + role="user", + content=f"[Context Summary - historical context only, not a new user request]\n{summary_text}", + is_summary=True, + is_display=False, + created_at=datetime.now(timezone.utc), + metadata={"breadcrumbs": summary_keys} if summary_keys else {}, + ) + + return system_msgs + [summary_msg] + to_keep, summary_text + + diff --git a/navi/core/compressor.py b/navi/core/compressor.py index e7cc534..3b260c9 100644 --- a/navi/core/compressor.py +++ b/navi/core/compressor.py @@ -1,1355 +1,37 @@ -""" -Context compressor — summarizes old messages to stay within the token limit. +"""Compatibility shim — the compressor lives in `navi.core.compression`. -Flow: -1. Partition session messages into "to_summarize" (old turns) and "to_keep" (recent turns). -2. Call the LLM to produce a structured summary of the old messages. -3. Replace the old turns with a structured summary message (role=user, is_summary=True). +Existing imports (`from navi.core.compressor import ContextCompressor`) keep +working; new code should import from `navi.core.compression`. -A "turn" is one user message plus all following assistant/tool messages up to the -next user message. Tool call groups (assistant + tool results) are never split. -Existing summary messages are always folded into the next compression pass. - -Compression is profile-aware: AgentProfile can provide compression_keep_recent, -compression_max_tokens, and a compression_prompt_file to specialize summaries. +Note for tests: the moved modules each bind their own `settings` object, so a +monkeypatch belongs on the module that reads it +(`navi.core.compression.partition`, `.summarize`, `.compressor`) — patching +this shim would change nothing. """ -import json -import re -import structlog -from datetime import datetime, timezone -from pathlib import Path -from collections.abc import Callable -from typing import TYPE_CHECKING - -from navi.llm.base import LLMBackend, Message -from navi.config import settings -from navi.memory._chunks import build_chunks -from navi.memory.breadcrumbs import MAX_BREADCRUMBS, extract_breadcrumbs, split_keys_line -from .events import ContextCompressed -from .token_calibration import calibration_for -from .context_builder import truncate_oversized -from .image_tokens import context_image_tokens, estimate_image_tokens - -if TYPE_CHECKING: - from navi.profiles.base import AgentProfile - -log = structlog.get_logger() - - -_SUMMARY_SECTIONS = [ - ("Goal", "One clear sentence describing what the user is trying to accomplish in this session. Include deadlines or constraints if stated."), - ("Active Files", "Every file or directory the assistant touched, with absolute or project-relative path, status (created / modified / read / deleted), the purpose of the change, and the key function/class/symbol affected."), - ("Decisions & User Preferences", "Explicit choices, architecture decisions, style preferences, or corrections stated by the user — each with a one-line rationale and what triggered it. Things the user said NOT to do."), - ("Completed Work", "Concrete finished steps — step + file/function name + verification outcome (pass/fail). Be specific and complete; prefer listing over generalizing."), - ("Intermediate Findings", "Key results from read/grep/log/terminal inspection that informed decisions: the fact discovered and the file/section it came from. Drop verbose bulk, keep the diagnostic line that changed understanding."), - ("Pending Work / Todo", "Open tasks, in-progress items, or follow-ups that still need action."), - ("Errors & Blockers", "Failures, exceptions, or unresolved issues — each with the exact short error snippet, what was tried, and the outcome. Preserve the diagnostic line verbatim."), - ("Key Values", "Exact constants the assistant should remember: ports, config keys and values, versions, dependency names, important paths, IDs, exact CLI flags."), -] - -_SUMMARY_TEMPLATE_INSTRUCTIONS = ( - "You are summarizing a conversation history to free up context space. " - "The assistant will continue working using ONLY this summary — it will have no access " - "to the original messages. Be thorough and precise; prefer specifics and complete lists " - "over generalities — the cost of losing a fact is higher than the cost of a longer summary. " - "This summary is historical context, not a new user request.\n\n" - "Use EXACTLY the Markdown structure below. Every section must be present. " - "If a section has no relevant information, write its header and the literal word NONE. " - "Keep bullets information-dense, but do not generalize away specifics to save space. " - "Do not include greetings, filler, transitions, or meta-commentary.\n\n" - + "\n\n".join(f"## {title}\n{desc}" for title, desc in _SUMMARY_SECTIONS) - + "\n\n" - "Output rules:\n" - "- Preserve exact file paths, function/class/symbol names, config keys and values, " - "exact CLI flags, and short error/diagnostic snippets verbatim.\n" - "- Preserve short critical code the agent must reproduce exactly: final function/class " - "signatures, one-line fixes the user approved, key config lines. Do NOT paste long blocks, " - "patches, or full command output — keep only the short line that matters.\n" - "- Do not paraphrase values that must stay precise.\n" - "- Use Markdown headers exactly as shown.\n" - "- End with ONE line `KEYS: item, item, …` — the breadcrumbs of this history: " - "names, commands, file paths, config keys, amounts, ticket numbers, decisions and " - "links that a later turn may ask about by name. Plain items, no explanations, at " - "most 30, most important first. This line is an index, not prose; write " - "`KEYS: NONE` when nothing is worth indexing." +from .compression import ( + ContextCompressor, + compress_context, + partition_current_turn_messages, + partition_messages, + should_compress, +) +from .compression.partition import _plan_compression, _turn_importance +from .compression.summarize import ( + _MAX_SUMMARY_INPUT_CHARS, + _build_summary_system_prompt, + _format_for_summary, ) - -# Tools whose full output is often needed later and should not be aggressively truncated. -_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), +__all__ = [ + "_MAX_SUMMARY_INPUT_CHARS", + "ContextCompressor", + "_build_summary_system_prompt", + "_format_for_summary", + "_plan_compression", + "_turn_importance", + "compress_context", + "partition_current_turn_messages", + "partition_messages", + "should_compress", ] - - -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 _format_for_summary( - messages: list[Message], image_budget: int | None = None -) -> tuple[str, list[str]]: - """ - Render messages as plain text for the summarization prompt. - - Returns (text, images) where images is a flat list of base64 strings - collected from all user messages. Vision-capable models will receive - the images alongside the text; non-vision models silently ignore them. - - ``image_budget`` caps the total cost of the collected images in tokens. - Every image in the history used to be forwarded in one user message with - only the text capped, so a session with a dozen screenshots produced a - summarization prompt that could not fit its own window — the call was - silently truncated by the backend and the summary was written as if the - history had been read. Images over the budget stay out of the call; their - per-message "[+ N image(s)]" note still tells the summarizer they existed. - """ - lines: list[str] = [] - images: list[str] = [] - remaining = image_budget if image_budget is None else max(0, image_budget) - i = 0 - while i < len(messages): - m = messages[i] - - if m.is_summary: - # Existing summary — include as-is (already compressed) - lines.append(m.content or "") - i += 1 - - elif m.role == "user": - if m.images: - for img in m.images: - if remaining is None: - images.append(img) - continue - cost = estimate_image_tokens(img) - if cost <= remaining: - images.append(img) - remaining -= cost - img_note = f" [+ {len(m.images)} image(s)]" - else: - img_note = "" - if m.content: - lines.append(f"User: {m.content}{img_note}") - elif img_note: - lines.append(f"User:{img_note}") - i += 1 - - elif m.role == "assistant" and m.tool_calls: - # Render tool calls + their results as a compact block - if m.content: - lines.append(f"Assistant: {m.content}") - for tc in m.tool_calls: - args_preview = json.dumps(tc.arguments)[:120] - lines.append(f"[Tool call: {tc.name}; arguments preview: {args_preview}]") - i += 1 - while i < len(messages) and messages[i].role == "tool": - tool_msg = messages[i] - result = tool_msg.content or "" - # Critical tool results and explicit critical flag survive verbatim - # up to a larger budget, so exact errors/file contents are preserved. - critical = ( - getattr(tool_msg, "is_compression_critical", False) - or tool_msg.name in _CRITICAL_TOOL_NAMES - ) - if critical and len(result) <= 4000: - preview = result - elif critical: - # Over the critical budget: keep head + tail halves rather - # than collapsing to the non-critical 800-char preview — - # the all-or-nothing cliff used to throw away 98% of a - # large file read (the common navi_code shape) while a - # 3999-char result survived whole. - half = 2000 - preview = result[:half] + "\n…[middle elided]…\n" + result[-half:] - else: - # Give the summarizer enough of the result to preserve - # diagnostic detail (the exact error line, a grep match, - # a key value), not just the first 300 chars — the summary - # is only as detailed as what the summarizer can see. - preview = result[:800] + ("…" if len(result) > 800 else "") - lines.append(f"[Tool result: {tool_msg.name}; preview: {preview}]") - i += 1 - - elif m.role == "assistant" and m.content: - lines.append(f"Assistant: {m.content}") - i += 1 - - else: - i += 1 # skip orphan tool messages - - return "\n".join(lines), images - - -# Safety limit: truncate formatted input to this many characters before sending to LLM. -# Prevents the summarizer from receiving near-context-sized input it can't fit alongside output. -# Typical ollama_num_ctx is 65k–128k, so 32k of input leaves ample room for a 6k-token summary. -_MAX_SUMMARY_INPUT_CHARS = 32_000 - -# Share of the window the historical images may take in a summarization call -# (see _format_for_summary). The text side is capped separately, in chars. -_SUMMARY_IMAGE_BUDGET_FRAC = 0.25 - -# A call that reports fewer real prompt tokens than this share of our own -# estimate of the prompt was sent a truncated prompt (Ollama keeps roughly -# half the window and evaluates that). The estimate errs high on images, so -# the band is deliberately wide — a false positive costs one extra -# compression pass, a missed one costs the session. -_TRUNCATION_RATIO = 0.7 - -# Hard-truncate fallbacks keep at most this fraction of the context window's -# tokens of recent messages (the rest is dropped without summarizing). -_HARD_TRUNCATE_TOKEN_FRAC = 0.5 - -# When existing summaries in to_summarize exceed this many chars combined, -# run a quick meta-summary to consolidate them before the main compression pass. -_META_SUMMARY_THRESHOLD = _MAX_SUMMARY_INPUT_CHARS // 3 # 10_666 - -_META_SUMMARY_SYSTEM = ( - "You are condensing multiple conversation summaries into a single compact summary. " - "Preserve all key facts, decisions, file paths, config values, errors, and user preferences. " - "Eliminate redundancy between overlapping summaries. " - "Write tight bullet points. Do not include filler or meta-commentary." -) - - -async def _meta_summarize( - summaries: list[Message], - llm: LLMBackend, - model, - temperature: float, -) -> Message: - """Consolidate multiple existing summary messages into one compact meta-summary.""" - combined = "\n\n---\n\n".join(m.content for m in summaries) - prompt = [ - Message(role="system", content=_META_SUMMARY_SYSTEM), - Message(role="user", content=combined), - ] - response = await llm.complete( - prompt, - tools=None, - temperature=temperature, - model=model, - think=False, - max_tokens=1500, - ) - text = (response.content or "").strip() or "(consolidated summary unavailable)" - return Message( - role="user", - content=f"[Consolidated Context Summary]\n{text}", - is_summary=True, - is_display=False, - is_context=True, - created_at=datetime.now(timezone.utc), - ) - - -def _build_summary_system_prompt(profile: "AgentProfile | None") -> str: - """Build the system prompt used by the summarization LLM. - - Uses the profile-specific compression prompt file if configured. - """ - base = _SUMMARY_TEMPLATE_INSTRUCTIONS - extra = "" - if profile is not None and getattr(profile, "compression_prompt_file", None): - profile_dir = Path("navi/profiles") / profile.id - prompt_path = profile_dir / profile.compression_prompt_file - if prompt_path.exists(): - try: - extra = "\n\n---\n\n[Profile-specific compression instructions]\n\n" + prompt_path.read_text(encoding="utf-8").strip() - except Exception: - pass - return base + extra - - -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"] - est = estimate or ContextCompressor.estimate_context_tokens - 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 - - -async def compress_context( - context: list[Message], - llm: LLMBackend, - model: "list[str] | str | None", - temperature: float, - keep_recent: int, - max_tokens: int | None = None, - keep_recent_messages: int | None = None, - profile: "AgentProfile | None" = None, - window: int | None = None, - estimate: Callable[[list[Message]], int] | None = None, -) -> tuple[list[Message], str] | None: - """ - Summarize old messages in the LLM context and return a shorter context list. - Only operates on `context` — the full display history (session.messages) is never touched. - Returns None if there is nothing to compress. - - Images from old user messages are passed to the summarization model. - Vision-capable models will incorporate image descriptions into the summary; - non-vision models silently ignore the images field. - - Uses the same model already loaded in memory (profile.model passed via WorkerContext) — - no model swap, no extra loading overhead. - - Profile settings override global defaults when provided: - - compression_keep_recent -> keep_recent - - compression_max_tokens -> max_tokens - - compression_prompt_file -> appended to summary system prompt - - Exceptions propagate to the caller (CompressionWorker catches them). - """ - effective_keep_recent = getattr(profile, "compression_keep_recent", None) or keep_recent - effective_max_tokens = getattr(profile, "compression_max_tokens", None) or max_tokens - - planned = _plan_compression( - context, - effective_keep_recent, - keep_recent_messages=keep_recent_messages, - window=window, - estimate=estimate, - ) - if planned is None: - return None # nothing substantial to compress - system_msgs, to_summarize, to_keep = planned - - # Meta-summary: if to_summarize contains multiple existing summary messages - # that are long enough to crowd the summarizer input, consolidate them first. - summary_msgs = [m for m in to_summarize if m.is_summary] - if len(summary_msgs) > 1: - total_summary_len = sum(len(m.content or "") for m in summary_msgs) - if total_summary_len > _META_SUMMARY_THRESHOLD: - try: - meta = await _meta_summarize( - summary_msgs, llm, model, temperature - ) - to_summarize = [m for m in to_summarize if not m.is_summary] - to_summarize.insert(0, meta) - except Exception: - # If meta-summary fails, continue with raw summaries as-is - pass - - # Images forwarded to the summarizer share one user message with the - # formatted text. Give them a quarter of the window at most: the summary - # itself needs the rest, and a summarization call that overflows is - # silently truncated by Ollama (to half the window), which would produce a - # summary written from an arbitrary slice of the history. - image_budget = int((window or settings.ollama_num_ctx) * _SUMMARY_IMAGE_BUDGET_FRAC) - summary_text_input, images = _format_for_summary( - to_summarize, image_budget=image_budget - ) - - # Truncate oversized input so the summarizer LLM has room to generate output. - # Keep head + tail, not just the head: to_summarize is chronological, so a - # head-only cut silently dropped the NEWEST summarized messages — the ones - # closest to the kept window, which bridge old context and the recent turn. - if len(summary_text_input) > _MAX_SUMMARY_INPUT_CHARS: - elision = "\n…[middle elided]…\n" - head = _MAX_SUMMARY_INPUT_CHARS * 3 // 4 - tail = _MAX_SUMMARY_INPUT_CHARS - head - len(elision) - summary_text_input = summary_text_input[:head] + elision + summary_text_input[-tail:] - - system_prompt = _build_summary_system_prompt(profile) - prompt = [ - Message(role="system", content=system_prompt), - Message(role="user", content=summary_text_input, images=images or None), - ] - - # think=False: compression must be fast — extended reasoning wastes context and hangs - response = await llm.complete( - prompt, tools=None, temperature=temperature, model=model, think=False, max_tokens=effective_max_tokens - ) - summary_text = (response.content or "").strip() or "(summary unavailable)" - # The summarizer's trailing `KEYS:` line is the breadcrumb index for the - # messages this pass drops. It is stripped from the summary the agent reads - # (it is an index, not prose) and carried on the summary message so the - # caller can write it next to the raw chunks. - summary_text, summary_keys = split_keys_line(summary_text) - if not summary_text.strip(): - summary_text = "(summary unavailable)" - - summary_msg = Message( - role="user", - content=f"[Context Summary - historical context only, not a new user request]\n{summary_text}", - is_summary=True, - is_display=False, - created_at=datetime.now(timezone.utc), - metadata={"breadcrumbs": summary_keys} if summary_keys else {}, - ) - - return system_msgs + [summary_msg] + to_keep, summary_text - - -class ContextCompressor: - """High-level context compression with retry strategy and hard-truncate fallback. - - Thin wrapper around `compress_context` that adds: - 1. Retry with keep_recent + 4 on LLM failure. - 2. Hard-truncate fallback (drop oldest messages without summarizing). - """ - - @staticmethod - def context_chars(context: list[Message]) -> int: - """Characters of message text — the input to every token estimate.""" - return sum(len(m.content or "") for m in context) - - @staticmethod - def estimate_context_tokens(context: list[Message]) -> int: - """Raw chars // 3 heuristic, no calibration. - - Kept as the reference unit for tests and for callers that want the - old, deliberately under-counting number. Window decisions must use - ``estimate_tokens`` instead — on prod this heuristic landed at ~56k - where the real prompt was 101k, and targets expressed in it were met - while the real prompt sat above the window. - """ - chars = ContextCompressor.context_chars(context) - return chars // 3 + context_image_tokens(context) - - @staticmethod - def estimate_tokens(context: list[Message], model=None) -> int: - """Token estimate in real-prompt units, used by every window decision. - - ``estimate_context_tokens`` converted to real prompt tokens by the - calibration fitted for ``model`` (see ``navi/core/token_calibration.py`` - for why chars // 3 is not usable as-is). Without observations the - calibration falls back to ``DEFAULT_CHARS_PER_TOKEN``. - """ - chars = ContextCompressor.context_chars(context) - return calibration_for(model).tokens_for(chars) + context_image_tokens(context) - - @staticmethod - def estimate_tokens_delta(context: list[Message], model=None) -> int: - """Price messages *appended since* a real baseline, same units. - - Uses the calibration slope without its intercept: the fixed prompt - overhead is already part of the baseline count. Pricing a delta with - the raw ``chars // 3`` left the pre/mid-turn gates in one unit and the - window in another for exactly the tool-heavy turns they exist for. - """ - chars = ContextCompressor.context_chars(context) - return calibration_for(model).delta_for(chars) + context_image_tokens(context) - - def __init__(self) -> None: - self._profile: "AgentProfile | None" = None - # Real prompt-token baseline from the last LLM call: - # (len(session.context) at call time, real prompt_tokens returned). - # The next estimate reuses these real tokens for the bulk of the context - # and only heuristics the messages appended since, instead of the - # chars//3 estimate which undercounts code-heavy tool output and fires - # compression too late. Cleared after compression (context shrank). - self._real_baseline: tuple[int, int] | None = None - # Set when the last call reported far fewer tokens than we sent: the - # backend truncated the prompt. Consumed by the post-turn worker, which - # compresses unconditionally in that case (see record_real_baseline). - self._context_truncated = False - - def set_profile(self, profile: "AgentProfile | None") -> None: - """Tell the compressor which profile is active so it can use profile-specific settings.""" - self._profile = profile - - @property - def context_truncated(self) -> bool: - """True when the last recorded call saw its prompt cut by the backend.""" - return self._context_truncated - - def pop_context_truncated(self) -> bool: - """Read and clear the truncation flag (the post-turn worker consumes it).""" - value = self._context_truncated - self._context_truncated = False - return value - - def record_real_baseline( - self, - context_len: int, - real_tokens: int | None, - *, - chars: int | None = None, - model=None, - sent_tokens: int | None = None, - image_tokens: int = 0, - ) -> None: - """Cache the real prompt-token count returned by the last LLM call. - - ``context_len`` is ``len(session.context)`` at the time of that call - (before the new assistant/tool messages are appended). The next - estimate reuses these real tokens for the bulk of the context and only - heuristics the new messages appended since. - - ``chars`` is the character count of the exact prompt that was sent - (the built context, not just session.context). Together with - ``real_tokens`` it feeds the per-model calibration line, which is what - converts hypothetical contexts (the proposed kept region) into real - tokens. Prompts that carried images are kept out of that fit: image - bytes are not in ``chars``, so their tokens would land in the fitted - intercept/slope and skew every later text estimate for that model. - - ``sent_tokens`` is our estimate of the prompt that was sent. When the - backend reports far fewer tokens than that, the prompt was truncated — - Ollama silently keeps about half the window and reports only what it - evaluated. Caching that number as the baseline used to pin the session - at ~50%: every later estimate started from the truncated count, the - 90% trigger never fired, and the model kept receiving a prompt whose - head (system prompt included) had been dropped. A truncated count is - therefore not recorded; the previous — honest — baseline stays, the - estimate keeps its size, and ``context_truncated`` tells the post-turn - worker to compress regardless of the threshold. - """ - if not real_tokens or real_tokens <= 0: - return - if ( - sent_tokens - and sent_tokens > 0 - and real_tokens < sent_tokens * _TRUNCATION_RATIO - ): - self._context_truncated = True - log.warning( - "compressor.prompt_truncated", - sent_tokens=sent_tokens, - real_tokens=real_tokens, - model=model, - ) - return - self._real_baseline = (context_len, int(real_tokens)) - if chars and not image_tokens: - calibration_for(model).observe(chars, real_tokens) - - def real_baseline_estimate( - self, - session_context: list[Message], - built_ctx: list[Message], - model=None, - ) -> int: - """Estimate context tokens using the real baseline when available. - - Falls back to the calibrated ``estimate_tokens`` on ``built_ctx`` when - no baseline is recorded or it is stale (the session shrank, e.g. right - after compression — ``base_len`` exceeds the current length). Every - branch is in real prompt tokens, the unit the window is measured in. - - Both the delta and the fallback are priced on the **truncated view** - (``truncate_oversized``), i.e. on what ``build()`` actually sends. The - raw stored context may hold a megabyte-long tool result that never - reaches the model; pricing it here made the gates fire for work the - compressor intentionally refuses (a single message cannot be split) - and — worse — made ``check_context_size`` kill the run with "context - too large" over a prompt the model never received. - """ - if self._real_baseline is not None: - base_len, base_tokens = self._real_baseline - if base_len <= len(session_context): - delta = self.estimate_tokens_delta( - truncate_oversized(session_context[base_len:]), model - ) - return base_tokens + delta - return self.estimate_tokens(truncate_oversized(built_ctx), model) - - async def compress_session( - self, - context: list[Message], - llm: LLMBackend, - model: "list[str] | str | None", - temperature: float, - keep_recent: int, - max_tokens: int | None = None, - keep_recent_messages: int | None = None, - window: int | None = None, - ) -> tuple[list[Message], str] | None: - """Compress context with retry + hard-truncate fallback. - - Returns (new_context, summary_text) or None if nothing changed. - Does NOT mutate the session — the caller is responsible for updating - session.context, session.context_token_count, and persisting. - """ - effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent - effective_max_tokens = getattr(self._profile, "compression_max_tokens", None) or max_tokens - window_tokens = window or settings.ollama_num_ctx - # Every size decision in this pass — the keep-region target, the - # post-check, the truncate budgets — must be in real-token units, the - # same units the trigger fired in, AND on the view the model is sent - # (``truncate_oversized``): pricing raw stored messages made a kept - # megabyte-long tool result look like the whole window and drop turns - # the model could have had. Mixed units are what let a compression - # pass finish at "65% of the window" while the real prompt it produced - # was above the window. - est = lambda msgs: self.estimate_tokens(truncate_oversized(msgs), model) - - # Attempt 1: normal compression - try: - result = await compress_context( - context=context, - llm=llm, - model=model, - temperature=temperature, - keep_recent=effective_keep_recent, - max_tokens=effective_max_tokens, - keep_recent_messages=keep_recent_messages, - profile=self._profile, - window=window_tokens, - estimate=est, - ) - except Exception: - # Attempt 2: keep more recent turns verbatim - try: - result = await compress_context( - context=context, - llm=llm, - model=model, - temperature=temperature, - keep_recent=effective_keep_recent + 4, - max_tokens=effective_max_tokens, - keep_recent_messages=(keep_recent_messages + 4) - if keep_recent_messages is not None - else None, - profile=self._profile, - window=window_tokens, - estimate=est, - ) - except Exception: - # Attempt 3: hard-truncate fallback - return self._hard_truncate(context, window=window_tokens, estimate=est) - - if result is None: - # Partition found nothing to summarize (e.g. few very large messages - # in a single turn). If the context is nonetheless over the token - # threshold, fall back to a token-budget hard-truncate so the gate - # firing always results in shrinkage — otherwise the agent runs until - # the window overflows. (The per-message truncation in - # context_builder.build caps single huge messages for the LLM view; - # this is the stored-context safety net.) - if should_compress( - est(context), - window_tokens, - settings.context_compression_threshold, - ): - return self._token_budget_truncate( - context, estimate=est, window=window_tokens - ) - return None - - # Target hysteresis safety net: if the compressed context is still over - # the target fraction (midturn mode kept a huge in-flight turn, or even - # the keep_recent floor could not fit the window), drop oldest kept - # turns by token budget so a compression firing always leaves real - # headroom below the trigger — otherwise the next few messages re-fire - # compression and the context yo-yos at the trigger line. The dropped - # turns were kept verbatim (not summarized) and are lost; this only - # triggers when the verbatim kept region alone exceeds the target. - new_context, summary_text = result - target_tokens = int(window_tokens * settings.context_compression_target) - if est(new_context) > target_tokens: - truncated = self._token_budget_truncate( - new_context, budget_tokens=target_tokens, estimate=est, window=window_tokens - ) - if truncated is not None: - new_context = truncated[0] - return new_context, summary_text - - def would_compress( - self, - context: list[Message], - keep_recent: int, - keep_recent_messages: int | None = None, - threshold: float | None = None, - max_context_tokens: int | None = None, - current_tokens: int | None = None, - model=None, - ) -> bool: - """True if ``compress_session`` would actually shrink the context. - - Runs the same ``_plan_compression`` decision as ``compress_context`` - (no LLM call). Also returns True when the context is over the token - threshold even though partition found nothing — the token-budget - fallback in ``compress_session`` handles that. Used to emit - ``CompressionStarted`` only when compression will really happen - (no misleading "compression" status on a no-op). - """ - effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent - window_tokens = max_context_tokens or settings.ollama_num_ctx - # Same unit as the caller's real-token estimate: the truncated view the - # model is actually sent, not the raw stored messages. - est = lambda msgs: self.estimate_tokens(truncate_oversized(msgs), model) - if _plan_compression( - context, - effective_keep_recent, - keep_recent_messages=keep_recent_messages, - window=window_tokens, - estimate=est, - ) is not None: - return True - # Partition no-op: compression can still shrink via the token-budget - # fallback — but only if dropping oldest turns actually reduces the - # context (a single huge turn that alone exceeds the budget cannot be - # split by message-boundary truncation, and is handled instead by the - # per-message view truncation in context_builder.build). Mirror that - # decision so would_compress and compress_session agree exactly. - # The size compared against the threshold is the caller's real-token - # estimate (`current_tokens`), not a fresh heuristic: the trigger fired - # on real tokens, and judging it again in chars // 3 units made the two - # disagree by ~1.8x. - if threshold is not None and max_context_tokens is not None: - total = current_tokens if current_tokens is not None else est(context) - if should_compress(total, window_tokens, threshold): - return self._token_budget_truncate( - context, estimate=est, window=window_tokens - ) is not None - return False - - def _hard_truncate( - self, - context: list[Message], - window: int | None = None, - estimate: Callable[[list[Message]], int] | None = None, - ) -> tuple[list[Message], str] | None: - """Last-resort fallback: drop oldest non-system messages without summarizing. - - Keeps whole conversational turns so tool call groups are never split. - Bounded by a token budget (half the context window) rather than a fixed - message count, so it still shrinks a context made of a few very large - messages — the fixed-count floor used to no-op on <= 6 messages even - when they were huge. - """ - budget = max(1, int((window or settings.ollama_num_ctx) * _HARD_TRUNCATE_TOKEN_FRAC)) - return self._token_budget_truncate( - context, budget_tokens=budget, estimate=estimate, window=window - ) - - def _token_budget_truncate( - self, - context: list[Message], - budget_tokens: int | None = None, - estimate: Callable[[list[Message]], int] | None = None, - window: int | None = None, - ) -> tuple[list[Message], str] | None: - """Drop oldest non-system messages (by turn groups) until under budget. - - Keeps system messages and the newest conversational turns whole (tool - call groups stay together). Returns None if the context already fits the - budget or there is nothing to drop. ``budget_tokens`` defaults to half - the context window. - """ - est = estimate or ContextCompressor.estimate_context_tokens - - if budget_tokens is None: - budget_tokens = max( - 1, int((window or settings.ollama_num_ctx) * _HARD_TRUNCATE_TOKEN_FRAC) - ) - - system_msgs = [m for m in context if m.role == "system"] - non_system = [m for m in context if m.role != "system"] - if not non_system: - return None - - # Group non-system into turns (user message starts a turn) - turns: list[list[Message]] = [] - current: list[Message] = [] - for msg in non_system: - if msg.role == "user" and current: - turns.append(current) - current = [msg] - else: - current.append(msg) - if current: - turns.append(current) - - # Keep newest turns (whole) until the token budget is exhausted; always - # keep at least the most recent turn so the agent retains the in-flight - # request + its latest tool results. - kept_turns: list[list[Message]] = [] - used = 0 - for turn in reversed(turns): - turn_tokens = est(turn) - if kept_turns and used + turn_tokens > budget_tokens: - break - kept_turns.insert(0, turn) - used += turn_tokens - - if kept_turns and used > budget_tokens: - # The newest turn alone is over budget (a first turn of a few huge - # messages): the loop above always keeps it whole, which used to - # make this fallback a no-op exactly when it was needed. Trim inside - # it — the user head and the newest messages stay verbatim. - newest = kept_turns[-1] - split = _split_turn_tail( - newest, - len(newest), - budget_tokens=budget_tokens, - estimate=est, - ) - if split is not None: - kept_turns[-1] = split[1] - - to_keep = [m for turn in kept_turns for m in turn] - if len(to_keep) == len(non_system): - return None # already within budget - - new_context = system_msgs + to_keep - # dropped > 0 is guaranteed here: to_keep == non_system returns None above. - summary_text = ( - "[Context was too large to summarize. Old messages were truncated to prevent " - "the model from exceeding its context window. Some earlier details may have been lost.]" - ) - return new_context, summary_text - - def check_context_size( - self, - context: list[Message], - session_context: list[Message] | None = None, - model=None, - window: int | None = None, - ) -> None: - """Raise ContextTooLargeError before an LLM call if the context is dangerously large.""" - # Local import: the window and the reserve are read at call time, so a - # runtime settings swap (tests, operator reload) reaches this check. - from navi.config import settings as _settings - from navi.exceptions import ContextTooLargeError - - if not context: - return - - window_tokens = window or _settings.ollama_num_ctx - # A reserve bigger than the window turns `available` negative and makes - # every check fail, however small the context — a 500-token model still - # has to be able to answer. Cap the reserve at a quarter of the window. - output_reserve = min( - _settings.output_reserve_tokens, max(1, window_tokens // 4) - ) - est = lambda msgs: self.estimate_tokens(msgs, model) - # When the caller supplies the raw session context, prefer the real-token - # baseline (real bulk + calibrated delta) over the calibrated estimate; - # both are in the same units as the real counts the window is measured in. - total = ( - self.real_baseline_estimate(session_context, context, model) - if session_context is not None - else est(context) - ) - available = window_tokens - output_reserve - - if total > available: - existing = est(context[:-1]) - new = est(context[-1:]) - remaining = available - existing - raise ContextTooLargeError( - f"Context too large: new content is ~{new:,} estimated tokens, " - f"but only ~{max(0, remaining):,} tokens are available " - f"(window {window_tokens:,}, already used ~{existing:,}, " - f"output_reserve {output_reserve:,}). " - "Split the file into smaller parts or delegate to a subagent." - ) - - async def compress_and_save_session( - self, - session, - session_store, - llm: LLMBackend, - model: str, - temperature: float, - session_id: str, - reason: str, - keep_recent: int, - max_tokens: int | None = None, - keep_recent_messages: int | None = None, - window: int | None = None, - memory_store=None, - ) -> ContextCompressed | None: - """Compresses the session context and persists the changes to the session store.""" - count_before = len(session.context) - window_tokens = window or settings.ollama_num_ctx - result = await self.compress_session( - context=session.context, - llm=llm, - model=model, - temperature=temperature, - keep_recent=keep_recent, - max_tokens=max_tokens, - keep_recent_messages=keep_recent_messages, - window=window_tokens, - ) - if result is None: - return None - - new_context, summary_text = result - - # Mark messages in session.messages as not context if they are no longer in new_context - # and are not system messages. - # Iterate session.context too: after a session reload, context-only rows - # (old summaries — is_display=False) are present in session.context but - # NOT in session.messages, so demoting only the latter left stale - # summaries with is_context=True in the DB — they piled up in the - # context at every reload (the half-the-context-is-summaries symptom). - new_context_ids = {id(m) for m in new_context} - demoted: list[Message] = [] - seen: set[int] = set() - for msg in [*session.messages, *session.context]: - if msg.role == "system" or id(msg) in new_context_ids or id(msg) in seen: - continue - seen.add(id(msg)) - if not msg.is_context: - # Dropped by an earlier compression. Iterating the display - # history surfaces it again on every later pass; indexing it - # again writes the same turns under a new chunk boundary (new - # hash) instead of once, forever. The current summary covers - # it; if it was dropped while indexing was off, that is the - # one turn range this trade-off leaves unindexed. - continue - msg.is_context = False - demoted.append(msg) - - # Rows that live only in session.context — a summary is is_display=False, - # so after a reload it is absent from session.messages — are invisible to - # save(), which updates flags for display history only. Their demotion - # used to stay in memory: every reload brought the stale summary back - # with is_context=True, and one prod session ended up carrying 15 live - # summaries, 80% of its context. Persist those flips explicitly. - in_messages = {id(m) for m in session.messages} - context_only = [ - m.sequence_number - for m in demoted - if m.sequence_number >= 0 and id(m) not in in_messages - ] - if context_only: - await session_store.set_context_flag(session_id, context_only, False) - - # Add the summary message to session.messages if it's not already there - summary_msg = next((m for m in new_context if m.is_summary), None) - if summary_msg and summary_msg not in session.messages: - summary_msg.is_display = False - session.messages.append(summary_msg) - - # Keep the dropped turns searchable before they leave the context for - # good. This must happen while `demoted` is still in hand: after the - # save there is no object that says "these were the turns compression - # removed", only the display history. - await self._index_demoted(memory_store, session, session_id, demoted, summary_msg) - - # Add a system message with is_compression=True and content=summary_text - session.messages.append( - Message( - role="system", - content=summary_text, - is_compression=True, - is_context=False, - created_at=datetime.now(timezone.utc), - ) - ) - - session.context = new_context - # The size the model is actually sent. A megabyte-long tool result stays - # in stored history (the user still sees it) but build() truncates it per - # message before the call — counting the raw message reported 817k tokens - # against a 128k window the model never received. - session.context_token_count = self.estimate_tokens( - truncate_oversized(new_context), model - ) - # The context just shrank, so the recorded real-token baseline is stale. - # Drop it — the next LLM call re-establishes a fresh one. The truncation - # flag goes with it: the prompt that will be sent is small again. - self._real_baseline = None - self._context_truncated = False - await session_store.save(session) - - log.info( - "compressor.compressed", - session_id=session_id, - reason=reason, - messages_before=count_before, - messages_after=len(new_context), - context_tokens=session.context_token_count, - ) - - # Archive old messages if the hot table exceeds the configured window. - if settings.session_messages_window > 0 and session.db_next_sequence > settings.session_messages_window: - threshold = session.db_next_sequence - settings.session_messages_window - archived = await session_store.archive_old_messages(session_id, threshold) - if archived > 0: - session.messages = [m for m in session.messages if m.sequence_number >= threshold] - session.context = [m for m in session.context if m.sequence_number >= threshold] - session.archive_threshold = threshold - - return ContextCompressed( - messages_before=count_before, - messages_after=len(new_context), - summary=summary_text, - context_tokens=session.context_token_count, - max_context_tokens=window_tokens, - ) - - async def _index_demoted( - self, - memory_store, - session, - session_id: str, - demoted: list[Message], - summary_msg: Message | None, - ) -> None: - """Write the turns compression just dropped into the raw-chunk index. - - Append-only: re-running compression on the same history hits the - `(session_id, content_hash)` uniqueness and writes nothing new. - - Breadcrumbs come from the summary's `KEYS:` line and, always, from the - deterministic extractor over the chunk text. The hard-truncate fallback - never calls a summarizer, so its index would otherwise be empty; and the - extractor also catches names the summarizer chose not to list. Every - chunk of this pass carries the union — the keys describe the batch, and - a searcher does not know which chunk of the batch holds the name. - """ - if memory_store is None or not demoted or not settings.memory_search_enabled: - return - try: - chunks = build_chunks(demoted) - if not chunks: - return - summary_keys = [] - if summary_msg is not None: - summary_keys = list( - (summary_msg.metadata or {}).get("breadcrumbs") or [] - ) - keys = _merge_breadcrumbs( - summary_keys, - extract_breadcrumbs( - [c["content"] for c in chunks], - limit=max(0, MAX_BREADCRUMBS - len(summary_keys)), - ), - ) - for chunk in chunks: - chunk["breadcrumbs"] = keys - await memory_store.add_chunks( - chunks, user_id=getattr(session, "user_id", None), session_id=session_id - ) - except Exception: - # Indexing must never break compression — a missing chunk is a lost - # search result, a raised exception here is a lost turn. - log.warning("compressor.chunk_index_failed", session_id=session_id, exc_info=True) - - -def _merge_breadcrumbs(*groups: list[str]) -> list[str]: - merged: list[str] = [] - seen: set[str] = set() - for group in groups: - for item in group: - key = " ".join(str(item).lower().split()) - if not key or key in seen: - continue - seen.add(key) - merged.append(str(item)) - if len(merged) >= MAX_BREADCRUMBS: - return merged - return merged diff --git a/tests/unit/core/test_compression_breadcrumbs.py b/tests/unit/core/test_compression_breadcrumbs.py index cc081aa..b2d1048 100644 --- a/tests/unit/core/test_compression_breadcrumbs.py +++ b/tests/unit/core/test_compression_breadcrumbs.py @@ -3,7 +3,7 @@ from types import SimpleNamespace import navi.config as _config -import navi.core.compressor as _cm +import navi.core.compression.compressor as _cm from navi.config import Settings from navi.core.compressor import ContextCompressor, compress_context from navi.core.session import InMemorySessionStore diff --git a/tests/unit/core/test_compressor.py b/tests/unit/core/test_compressor.py index fbc91fa..3b7c698 100644 --- a/tests/unit/core/test_compressor.py +++ b/tests/unit/core/test_compressor.py @@ -449,7 +449,7 @@ async def test_meta_summary_graceful_on_failure(self): """If meta-summary fails, compression continues with raw summaries.""" - import navi.core.compressor as compressor_module + import navi.core.compression.summarize as compressor_module backend = FakeLLMBackend(responses=["Final summary"]) original_meta = compressor_module._meta_summarize @@ -673,7 +673,7 @@ @pytest.mark.asyncio async def test_compress_session_retry_on_failure(self): """First attempt fails, second succeeds with keep_recent + 4.""" - import navi.core.compressor as compressor_module + import navi.core.compression.compressor as compressor_module backend = FakeLLMBackend(responses=["Summary text"]) original_compress_context = compressor_module.compress_context @@ -708,7 +708,7 @@ @pytest.mark.asyncio async def test_compress_session_hard_truncate_on_double_failure(self): """Both attempts fail → token-budget hard-truncate drops oldest turns.""" - import navi.core.compressor as compressor_module + import navi.core.compression.compressor as compressor_module from unittest.mock import AsyncMock original_compress_context = compressor_module.compress_context @@ -748,6 +748,37 @@ finally: compressor_module.compress_context = original_compress_context + def test_token_budget_truncate_trims_a_single_oversized_turn(self): + """The last-resort truncate must trim inside the newest turn when that + turn alone exceeds the budget — keeping it whole would make the fallback + a no-op exactly when it is needed. Also pins the `_split_turn_tail` + import that the compression split moved into `partition`.""" + head = Message(role="user", content="analyse this") + rest = [ + Message( + role="assistant", + content="", + tool_calls=[ToolCallRequest(id=str(i), name="fs", arguments={})], + ) + if i % 2 == 0 + else Message(role="tool", name="fs", content="x" * 10_000, tool_call_id=str(i - 1)) + for i in range(6) + ] + context = [Message(role="system", content="sys"), head, *rest] + est = lambda msgs: sum(len(m.content or "") for m in msgs) // 4 + + result = ContextCompressor()._token_budget_truncate( + context, budget_tokens=1_500, estimate=est, + ) + + assert result is not None + kept, summary = result + assert "truncated" in summary.lower() + # The user head survives; the middle of the turn went away. + assert kept[0].role == "system" + assert kept[1] is head + assert len(kept) < len(context) + @pytest.mark.asyncio async def test_compress_session_returns_none_when_nothing_to_compress(self): backend = FakeLLMBackend() @@ -904,15 +935,17 @@ @staticmethod def _override_settings(ollama_num_ctx: int): - """Patch both navi.config.settings and navi.core.compressor.settings. + """Patch `settings` on every module of the split compressor. - compress_context reads the module-level ``settings`` binding of - compressor.py; compress_session's safety net reads ``navi.config.settings`` - via a local ``import navi.config as _cfg``. Both must be patched for the - override to reach every code path. + `navi.core.compression.{partition,summarize,compressor}` each bind their + own `settings`, and `compress_session`'s safety net reads + `navi.config.settings` via a local import. Patching only one of them + leaves the override half-applied. """ import navi.config as _config - import navi.core.compressor as _cm + import navi.core.compression.compressor as _cm + import navi.core.compression.partition as _partition + import navi.core.compression.summarize as _summarize from navi.config import Settings new = Settings( @@ -922,19 +955,23 @@ navi_persona_file="", context_compression_target=0.65, ) - prev_config = _config.settings - prev_cm = _cm.settings + prev = ( + _config.settings, _cm.settings, _partition.settings, _summarize.settings, + ) _config.settings = new _cm.settings = new - return prev_config, prev_cm + _partition.settings = new + _summarize.settings = new + return prev @staticmethod - def _restore(prev_config, prev_cm): + def _restore(prev): import navi.config as _config - import navi.core.compressor as _cm + import navi.core.compression.compressor as _cm + import navi.core.compression.partition as _partition + import navi.core.compression.summarize as _summarize - _config.settings = prev_config - _cm.settings = prev_cm + _config.settings, _cm.settings, _partition.settings, _summarize.settings = prev @pytest.mark.asyncio async def test_preturn_adaptive_shrink_fits_target(self): @@ -942,7 +979,7 @@ leave the kept region above the 65% target, so compress_context shrinks keep_recent until the verbatim kept region fits. The extra turns fold into the summary (still one LLM call).""" - prev_config, prev_cm = self._override_settings(ollama_num_ctx=10_000) + prev = self._override_settings(ollama_num_ctx=10_000) try: # target_tokens = 10_000 * 0.65 = 6_500 backend = FakeLLMBackend(responses=["Summary text"]) @@ -972,13 +1009,13 @@ # Single LLM call — shrink used partition only, no extra summarization. assert backend._call_idx == 1 finally: - self._restore(prev_config, prev_cm) + self._restore(prev) @pytest.mark.asyncio async def test_preturn_no_shrink_when_already_under_target(self): """Small context: the kept region already fits the target, so keep_recent is not shrunk — normal compression behavior is preserved.""" - prev_config, prev_cm = self._override_settings(ollama_num_ctx=100_000) + prev = self._override_settings(ollama_num_ctx=100_000) try: backend = FakeLLMBackend(responses=["Summary text"]) context = [ @@ -1001,7 +1038,7 @@ assert len(new_context) == 6 assert ContextCompressor.estimate_context_tokens(new_context) <= 65_000 finally: - self._restore(prev_config, prev_cm) + self._restore(prev) @pytest.mark.asyncio async def test_safety_net_truncates_when_floor_still_over_target(self): @@ -1013,7 +1050,7 @@ Sizes are chosen against the *view* estimate, the unit the window is measured in: each 12k-char message is view-capped below num_ctx // 6. """ - prev_config, prev_cm = self._override_settings(ollama_num_ctx=10_000) + prev = self._override_settings(ollama_num_ctx=10_000) try: # target_tokens = 6_500. Each turn: user + 2 assistant messages, # each assistant view-capped at ~1_670 tokens -> a turn is ~3_300, @@ -1052,7 +1089,7 @@ assert new_context[1].content == "u5" assert not any(m.is_summary for m in new_context) finally: - self._restore(prev_config, prev_cm) + self._restore(prev) class TestWouldCompress: @@ -1296,7 +1333,9 @@ @staticmethod def _override_settings(monkeypatch): import navi.config as _config - import navi.core.compressor as _cm + import navi.core.compression.compressor as _cm + import navi.core.compression.partition as _partition + import navi.core.compression.summarize as _summarize from navi.config import Settings new = Settings( @@ -1306,8 +1345,8 @@ navi_persona_file="", context_compression_target=0.65, ) - monkeypatch.setattr(_config, "settings", new) - monkeypatch.setattr(_cm, "settings", new) + for module in (_config, _cm, _partition, _summarize): + monkeypatch.setattr(module, "settings", new) return new @staticmethod @@ -1643,7 +1682,9 @@ @staticmethod def _override_settings(): import navi.config as _config - import navi.core.compressor as _cm + import navi.core.compression.compressor as _cm + import navi.core.compression.partition as _partition + import navi.core.compression.summarize as _summarize from navi.config import Settings new = Settings( @@ -1652,25 +1693,30 @@ context_message_token_budget=100, # ~300 chars per message context_compression_target=0.65, ) - prev_config, prev_cm = _config.settings, _cm.settings + prev = ( + _config.settings, _cm.settings, _partition.settings, _summarize.settings, + ) _config.settings = new _cm.settings = new - return prev_config, prev_cm + _partition.settings = new + _summarize.settings = new + return prev @staticmethod - def _restore(prev_config, prev_cm): + def _restore(prev): import navi.config as _config - import navi.core.compressor as _cm + import navi.core.compression.compressor as _cm + import navi.core.compression.partition as _partition + import navi.core.compression.summarize as _summarize - _config.settings = prev_config - _cm.settings = prev_cm + _config.settings, _cm.settings, _partition.settings, _summarize.settings = prev @pytest.mark.asyncio async def test_oversized_kept_message_is_counted_as_the_model_sees_it(self): from navi.core.context_builder import truncate_oversized from navi.core.session import Session - prev_config, prev_cm = self._override_settings() + prev = self._override_settings() try: store = _RecordingStore() c = ContextCompressor() @@ -1705,4 +1751,4 @@ # The raw estimate is the scary number; the stored one is not it. assert raw > effective * 10 finally: - self._restore(prev_config, prev_cm) + self._restore(prev) diff --git a/tests/unit/workers/test_compression_worker.py b/tests/unit/workers/test_compression_worker.py index e9346bc..ff7ad8f 100644 --- a/tests/unit/workers/test_compression_worker.py +++ b/tests/unit/workers/test_compression_worker.py @@ -105,7 +105,7 @@ shrinks and ContextCompressed is emitted. The old worker called compress_context directly and silently no-op'ed on any LLM failure, leaving the session over the threshold until the next turn's gates.""" - import navi.core.compressor as compressor_module + import navi.core.compression.compressor as compressor_module worker = CompressionWorker() backend = FakeLLMBackend(responses=["unused"])