Coverage for src/lilbee/retrieval/query/compaction.py: 100%
94 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Condense the turns that no longer fit the prompt into a rolling summary.
3Windowing alone drops the oldest turns outright, so a resumed conversation the
4user can still scroll is one the model cannot see: it answers as if the earlier
5turns never happened. Compaction keeps a summary of what was dropped and carries
6it at the head of the prompt, so old context degrades to a gist instead of
7vanishing. The transcript on disk and on screen is never touched.
8"""
10from __future__ import annotations
12from dataclasses import dataclass
13from typing import TYPE_CHECKING
15from lilbee.retrieval.query.history_window import (
16 chars_for_tokens,
17 estimate_tokens,
18 windowed_history,
19)
21if TYPE_CHECKING:
22 from lilbee.retrieval.query.searcher import ChatMessage
24# Kept simple: 0.6B models botch elaborate structured-summary instructions.
25# {words} is filled from summary_word_budget so the ask agrees with the cap.
26COMPACT_PROMPT = (
27 "Condense the conversation below into brief factual notes that let an "
28 "assistant carry it on. Keep names, numbers, decisions, and anything left "
29 "unresolved; drop pleasantries. Under {words} words. Return ONLY the notes.\n\n"
30 "Conversation:\n{transcript}"
31)
33# Ceiling on a summary's tokens; ctx/8 governs below it. A tighter ceiling
34# flattens long conversations to a few hundred words even on a 32k window.
35COMPACT_MAX_TOKENS = 1024
38def summary_word_budget(ctx_target: int) -> int:
39 """The word count the prompt asks for, derived from the token cap.
41 Roughly three words per four tokens, so the instruction and ``num_predict``
42 agree: asking for more words than the cap can hold guarantees a truncated
43 final sentence, and asking for far fewer wastes the window.
44 """
45 return summary_cap(ctx_target) * 3 // 4
48# A summary must never eat the window it exists to protect: at a 2048 target the
49# flat 320-token cap would be a third of the history budget. Scale it down with
50# the model, with a floor that still fits a useful note.
51_SUMMARY_CTX_FRACTION = 8
52_SUMMARY_MIN_TOKENS = 64
54# Rough cost of the instruction wrapper around a batch transcript.
55_PROMPT_OVERHEAD_TOKENS = 64
56# Fraction of the window a batch may claim: chars/4 under-counts terse text by
57# up to ~1.8x (measured), and an overflowing batch strands its turns.
58_ESTIMATE_SAFETY_FRACTION = 0.6
59# Never build a batch smaller than this, however tight the window.
60_MIN_BATCH_TOKENS = 128
62# Most model calls one compaction may spend. Switching a 100k-token conversation
63# onto a 2k model produces ~100 batches; folding them all would stall the user's
64# next turn for minutes to produce a 256-token note, i.e. ~99% of the content is
65# discarded either way. Condense the most recent slice well and say plainly that
66# the rest was dropped, rather than stalling to produce mush.
67MAX_COMPACT_CALLS = 4
69# Compaction fires at a fraction of the budget rather than at the limit, and
70# clears down to a summary plus the newest exchanges.
71#
72# Firing AT the limit and folding only the overflow leaves the history still at
73# the limit, so the next turn overflows again and every later turn pays a model
74# call. Triggering early and clearing deep buys many turns of headroom for the
75# same one call: driving these functions over a 40-turn conversation of ~350
76# token turns costs 2 calls at the 8192 context floor and 18 at 2048 (0 at 32k,
77# which never fills). The fire-at-the-limit shape this replaces modelled at ~34
78# and ~70 for the same runs -- modelled, not measured, since that code is gone.
79# test_a_long_chat_does_not_compact_on_every_turn pins the property.
80#
81# This is the shape Anthropic's own compaction uses (a token threshold, then the
82# conversation replaced by a summary). The difference here is that the newest
83# exchanges stay verbatim: a chat has to answer the question just asked, not a
84# paraphrase of it.
85COMPACT_TRIGGER_FRACTION = 0.8
86# Messages (not exchanges) kept verbatim when compaction clears the history.
87COMPACT_KEEP_RECENT = 4
89# Fraction of ``chat_n_ctx_target`` a conversation may spend on its history;
90# the rest is for the system prompt, RAG context, question, and reasoning.
91HISTORY_TOKEN_BUDGET_FRACTION = 0.5
94def history_budget(ctx_target: int) -> int:
95 """Token budget for everything a conversation carries into the prompt."""
96 return int(ctx_target * HISTORY_TOKEN_BUDGET_FRACTION)
99@dataclass(frozen=True)
100class CompactionResult:
101 """Notes produced by one compaction, and what they do and do not cover."""
103 summary: str
104 condensed: int
105 """Messages folded into the notes (not exchanges: a user+assistant pair is two)."""
106 stranded: int
107 """Messages dropped with no notes. Non-zero means the conversation lost detail
108 outright, which the UI must say plainly rather than let the model appear to
109 have forgotten for no reason."""
112@dataclass(frozen=True)
113class CompactionPlan:
114 """What one compaction will attempt, and what it gives up on before it starts."""
116 batches: list[list[ChatMessage]]
117 stranded: int
118 """Messages dropped with no notes because the backlog exceeded MAX_COMPACT_CALLS.
120 Deliberately no ``condensed`` counterpart: a plan cannot know what will be
121 condensed, only what it will try. Whether a batch lands depends on the model
122 answering, so only the caller that watched it can count. A planned-not-actual
123 count is what made the UI claim turns were summarized when they were lost.
124 """
127# The summary rides in as a user/assistant pair rather than a second system
128# message: the prompt is assembled as [system] + history + [user], and most chat
129# templates accept only the leading system message, silently dropping or
130# rejecting a later one. A pair keeps user/assistant alternation intact for
131# every template, and windowed_history drops in pairs for the same reason.
132SUMMARY_REQUEST = "Before we go on, remind me what we have covered so far."
135def summary_messages(summary: str) -> list[ChatMessage]:
136 """The synthetic pair that carries *summary* into the prompt, or nothing."""
137 if not summary.strip():
138 return []
139 return [
140 {"role": "user", "content": SUMMARY_REQUEST},
141 {"role": "assistant", "content": summary},
142 ]
145def prompt_history(
146 history: list[ChatMessage], summary: str, *, max_tokens: int
147) -> list[ChatMessage]:
148 """Assemble the history a prompt should carry: the summary, then recent turns.
150 The summary is charged against the same budget and reserved first, so adding
151 it can never push the prompt over the limit it exists to respect.
153 When even that does not fit, the summary is dropped rather than stacked on
154 top. windowed_history deliberately keeps the newest pair whatever it costs
155 (an empty prompt is useless), so adding notes to an already-oversized window
156 would push the prompt further past the budget than carrying no summary at
157 all -- and overflow is an engine failure, not a worse answer. Faced with a
158 turn too big to share, the live question beats notes about old ones.
159 """
160 pair = summary_messages(summary)
161 reserved = sum(estimate_tokens(m) for m in pair)
162 recent = windowed_history(history, max_tokens=max(1, max_tokens - reserved))
163 if reserved + sum(estimate_tokens(m) for m in recent) > max_tokens:
164 return windowed_history(history, max_tokens=max_tokens)
165 return pair + recent
168def overflow(history: list[ChatMessage], *, max_tokens: int) -> list[ChatMessage]:
169 """The oldest turns that do not fit *max_tokens*, i.e. what compaction folds away."""
170 kept = windowed_history(history, max_tokens=max_tokens)
171 dropped = len(history) - len(kept)
172 return history[:dropped] if dropped > 0 else []
175def compaction_due(history: list[ChatMessage], summary: str, *, max_tokens: int) -> bool:
176 """Whether the conversation has filled enough of its budget to compact.
178 Deliberately below the limit: waiting until the prompt no longer fits means
179 compacting on every subsequent turn, because folding just the overflow leaves
180 it full again. See COMPACT_TRIGGER_FRACTION.
181 """
182 used = sum(estimate_tokens(m) for m in history)
183 used += sum(estimate_tokens(m) for m in summary_messages(summary))
184 return used > max_tokens * COMPACT_TRIGGER_FRACTION
187def foldable(history: list[ChatMessage]) -> list[ChatMessage]:
188 """Everything compaction folds into notes: all but the newest exchanges.
190 Clearing this much is what buys headroom. The tail stays verbatim because a
191 chat has to answer the question just asked; a summary of it is not the same
192 thing, which is where a plain "summarize everything" would go wrong.
194 No ``keep`` parameter: COMPACT_KEEP_RECENT is the policy, and a knob nobody
195 turns is just a second place for it to disagree with itself.
197 The boundary is aligned forward to a user message so the kept window opens
198 a turn. A history can be odd-length (an interrupted turn persists a user
199 message with no reply), and cutting a fixed count would then leave the
200 window starting on an assistant reply, so the assembled prompt would run
201 user, assistant, assistant -- the non-alternating shape the summary pair
202 and windowed_history both take care to avoid.
203 """
204 keep = COMPACT_KEEP_RECENT
205 if len(history) <= keep:
206 return []
207 cut = len(history) - keep
208 for i in range(cut, len(history)):
209 if history[i]["role"] == "user":
210 return history[:i]
211 # No user message in the tail at all: keep the plain boundary rather than
212 # folding the entire history away.
213 return history[:cut]
216def summary_cap(ctx_target: int) -> int:
217 """How many tokens a summary may spend, scaled to the model it is written for."""
218 return max(_SUMMARY_MIN_TOKENS, min(COMPACT_MAX_TOKENS, ctx_target // _SUMMARY_CTX_FRACTION))
221def batch_overflow(dropped: list[ChatMessage], *, ctx_target: int) -> list[list[ChatMessage]]:
222 """Split *dropped* into batches whose summarize prompt each fit *ctx_target*.
224 Compaction usually nibbles a pair at a time, but switching models does not:
225 dropping from a 32k model to a 2k one turns tens of thousands of tokens into
226 overflow at once. Summarizing that in a single call would overflow the very
227 model being compacted for, and the failure path would keep the old summary,
228 losing every one of those turns. Folding batch by batch keeps each call
229 inside the current window, so a switch costs several small calls instead of
230 one impossible one.
232 A lone turn larger than a whole batch is truncated rather than sent to
233 certain failure: half its text summarized beats the entire turn dropped.
234 """
235 # Proportional, not a fixed pad: the estimate error and the server-side
236 # chat-template cost both grow with content (a raw-room batch measured
237 # ~2666 real tokens against a 2048 window and stranded every turn).
238 room = max(
239 _MIN_BATCH_TOKENS,
240 int(
241 (ctx_target - summary_cap(ctx_target) - _PROMPT_OVERHEAD_TOKENS)
242 * _ESTIMATE_SAFETY_FRACTION
243 ),
244 )
245 batches: list[list[ChatMessage]] = []
246 current: list[ChatMessage] = []
247 current_tokens = 0
248 for message in dropped:
249 cost = estimate_tokens(message)
250 if cost > room:
251 if current:
252 batches.append(current)
253 current, current_tokens = [], 0
254 batches.append([_truncated(message, room)])
255 continue
256 if current and current_tokens + cost > room:
257 batches.append(current)
258 current, current_tokens = [], 0
259 current.append(message)
260 current_tokens += cost
261 if current:
262 batches.append(current)
263 return batches
266def plan_compaction(dropped: list[ChatMessage], *, ctx_target: int) -> CompactionPlan:
267 """Decide what one compaction condenses, keeping its cost bounded.
269 Beyond MAX_COMPACT_CALLS batches the oldest turns are stranded rather than
270 folded: a 2k-context model cannot carry a 100k conversation whatever we
271 spend, so buying a marginally better note with minutes of stall is a bad
272 trade. The most recent slice is the part still worth remembering, and the
273 caller reports the stranded count instead of pretending it survived.
274 """
275 batches = batch_overflow(dropped, ctx_target=ctx_target)
276 if len(batches) <= MAX_COMPACT_CALLS:
277 return CompactionPlan(batches=batches, stranded=0)
278 return CompactionPlan(
279 batches=batches[-MAX_COMPACT_CALLS:],
280 stranded=sum(len(b) for b in batches[:-MAX_COMPACT_CALLS]),
281 )
284def merge_notes(previous_summary: str, notes: list[str]) -> str:
285 """Join carried-forward notes with fresh per-batch notes, oldest first."""
286 parts = [part.strip() for part in [previous_summary, *notes] if part.strip()]
287 return "\n".join(parts)
290def _truncated(message: ChatMessage, max_tokens: int) -> ChatMessage:
291 """A copy of *message* clipped to roughly *max_tokens*, marked as clipped."""
292 keep = chars_for_tokens(max_tokens)
293 return {"role": message["role"], "content": message["content"][:keep] + " […clipped]"}