Coverage for haystack/hooks/compaction/sliding_window.py: 100%
67 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from typing import Any
7from haystack.core.serialization import default_to_dict
8from haystack.dataclasses import ChatMessage, ChatRole
9from haystack.hooks.compaction.types import Compactor
10from haystack.hooks.compaction.utils import (
11 _COMPACTION_META_KEY,
12 _current_agent_step_groups,
13 _historical_turn_groups,
14 _is_compaction_message,
15 _latest_user_index,
16 _leading_system_end,
17 _messages_at,
18 _messages_except,
19)
20from haystack.token_counters import TokenCounter
21from haystack.utils.experimental import _experimental
23# Recorded as the strategy on every message this compactor produces, so a later run can recognize its own notes.
24_STRATEGY = "sliding_window"
26# Placeholder a custom omission note may include to have the number of removed messages substituted in.
27_NUM_REMOVED_PLACEHOLDER = "{num_removed}"
29_DEFAULT_OMISSION_NOTE = (
30 f"[{_NUM_REMOVED_PLACEHOLDER} earlier messages were removed from this conversation to free up context and cannot "
31 f"be recovered.]"
32)
35def _is_compaction_note(message: ChatMessage) -> bool:
36 """Whether a message is an omission note this strategy left in place of removed history."""
37 return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER)
40def _flatten(groups: list[list[int]]) -> list[int]:
41 """Join index groups into a single ordered list of indices."""
42 return [index for group in groups for index in group]
45def _removable_groups(
46 messages: list[ChatMessage], system_end: int, task_index: int | None
47) -> tuple[list[list[int]], list[list[int]]]:
48 """
49 Group the two stretches of conversation compaction is allowed to remove.
51 :param messages: The full conversation, ordered oldest to newest.
52 :param system_end: The end of the leading system-message block.
53 :param task_index: The index of the user message anchoring the current task, or None when there is none.
54 :returns: Index groups for the complete historical turns preceding the current task, and index groups for the
55 current task's own Agent steps. Both are ordered oldest group first. A group is the unit of removal: it is
56 kept or removed entire, which is what keeps a tool call with its results and an assistant reply with the user
57 message it answers.
58 """
59 # An earlier compaction's note is left out of its turn, so keeping the turn folds that note into the note this
60 # compaction leaves behind.
61 historical_groups = [
62 [index for index in group if not _is_compaction_note(message=messages[index])]
63 for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index)
64 ]
65 agent_step_groups = _current_agent_step_groups(messages=messages, system_end=system_end, task_index=task_index)
66 return historical_groups, agent_step_groups
69def _first_group_to_keep(
70 messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter
71) -> int:
72 """
73 Return the position in `groups` to start keeping from.
75 Groups are added up from newest to oldest and counting stops at the first group that does not fit.
77 :param messages: The full conversation containing the messages referenced by `groups`.
78 :param groups: Ordered index groups to choose from, the oldest group first is at position 0.
79 :param available_tokens: The token budget available for these groups.
80 :param token_counter: The `TokenCounter` used to measure each group.
81 :returns: The position in `groups` to start keeping from: `len(groups)` when nothing fits, `0` when it all fits.
82 """
83 # We count backwards so first_kept starts such that nothing would be kept
84 first_kept = len(groups)
85 while first_kept > 0:
86 # Calculate the tokens consumed by the group that would be kept next
87 group_tokens = token_counter.count(messages=_messages_at(messages=messages, indices=groups[first_kept - 1]))
88 # If this group does not fit, we stop
89 if group_tokens > available_tokens:
90 break
91 available_tokens -= group_tokens
92 first_kept -= 1
93 return first_kept
96def _first_turn_and_step_to_keep(
97 messages: list[ChatMessage],
98 historical_groups: list[list[int]],
99 agent_step_groups: list[list[int]],
100 available_tokens: int,
101 token_counter: TokenCounter,
102 min_keep_steps: int,
103) -> tuple[int, int]:
104 """
105 Return which historical turn and which Agent step of the current task to start keeping from.
107 :param messages: The full conversation containing the messages referenced by both group lists.
108 :param historical_groups: Index groups for the complete historical turns preceding the current task, oldest first.
109 :param agent_step_groups: Index groups for the current task's Agent steps, oldest first.
110 :param available_tokens: The token budget left once the protected context is paid for.
111 :param token_counter: The `TokenCounter` used to measure the groups.
112 :param min_keep_steps: The fewest recent Agent steps to keep, even when they exceed the budget.
113 :returns: The position in `historical_groups` and the position in `agent_step_groups` to start keeping from.
114 Either is the length of its list when nothing from it is kept.
115 """
116 current_task_tokens = token_counter.count(
117 messages=_messages_at(messages=messages, indices=_flatten(groups=agent_step_groups))
118 )
119 if current_task_tokens > available_tokens:
120 # The current task alone overruns the budget, so every historical turn is dropped and the current task's own
121 # oldest steps trimmed until what remains fits.
122 first_kept_step = _first_group_to_keep(
123 messages=messages, groups=agent_step_groups, available_tokens=available_tokens, token_counter=token_counter
124 )
125 # The newest steps are kept regardless of the budget.
126 return len(historical_groups), min(first_kept_step, max(len(agent_step_groups) - min_keep_steps, 0))
128 # The entire current task fits, so every step stays and the rest of the budget goes on the newest turns that fit.
129 first_kept_turn = _first_group_to_keep(
130 messages=messages,
131 groups=historical_groups,
132 available_tokens=available_tokens - current_task_tokens,
133 token_counter=token_counter,
134 )
135 return first_kept_turn, 0
138def _task_and_step_split(
139 messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int
140) -> tuple[list[ChatMessage], int, list[ChatMessage]]:
141 """
142 Split a conversation into the messages to keep and the messages to remove.
144 Leading system messages and the latest real user message are always kept. Historical turns are kept in full when
145 they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time
146 while keeping each assistant message together with its tool results.
148 :param messages: The full conversation to split, ordered oldest to newest.
149 :param target_tokens: The token budget for the messages that are kept.
150 :param token_counter: The `TokenCounter` used to decide which historical turns and current-task steps fit.
151 :param min_keep_steps: The fewest recent current-task Agent steps to keep, even if they exceed the target token
152 budget.
153 :returns: A tuple containing:
155 1. The messages to keep, ordered oldest to newest.
156 2. The position in that list where an omission note belongs, which is where the removed messages used to sit:
157 directly after the leading system messages when only historical turns were removed, and directly after the
158 user message anchoring the current task when the task's own steps were removed.
159 3. Every message selected for removal.
160 """
161 # The landmarks the split is built around: the Agent instructions, and the user message anchoring the current task.
162 system_end = _leading_system_end(messages=messages)
163 task_index = _latest_user_index(messages=messages)
164 task_indices = [task_index] if task_index is not None else []
166 # The two stretches compaction may remove. A group is the unit of removal, so a turn or a step is never split.
167 historical_groups, agent_step_groups = _removable_groups(
168 messages=messages, system_end=system_end, task_index=task_index
169 )
171 # The instructions and the current task are never removed, so they come off the budget first.
172 protected = _messages_at(messages=messages, indices=[*range(system_end), *task_indices])
173 first_kept_turn, first_kept_step = _first_turn_and_step_to_keep(
174 messages=messages,
175 historical_groups=historical_groups,
176 agent_step_groups=agent_step_groups,
177 available_tokens=target_tokens - token_counter.count(messages=protected),
178 token_counter=token_counter,
179 min_keep_steps=min_keep_steps,
180 )
182 # What survives, laid out in conversation order: the instructions, the turns that fit, the task, then its steps.
183 kept_turn_indices = _flatten(groups=historical_groups[first_kept_turn:])
184 kept_step_indices = _flatten(groups=agent_step_groups[first_kept_step:])
185 kept_indices = [*range(system_end), *kept_turn_indices, *task_indices, *kept_step_indices]
187 # The note stands in for what was dropped, so it goes where those messages used to sit. Either right after the
188 # instructions when the historical turns were trimmed, or right after the task anchor when its own steps were.
189 note_index = system_end if first_kept_step == 0 else system_end + len(kept_turn_indices) + len(task_indices)
190 return (
191 _messages_at(messages=messages, indices=kept_indices),
192 note_index,
193 _messages_except(messages=messages, indices=kept_indices),
194 )
197@_experimental
198class SlidingWindowCompactor(Compactor):
199 """
200 Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows.
202 Leading system messages and the latest user message are protected. Historical turns are kept in full when they fit,
203 and the current task's history is kept in complete Agent steps, where a step is an assistant message together
204 with all immediately following tool results.
206 An `omission_note` is left where the removed messages used to sit: directly after the leading system messages when
207 only historical turns were removed, and directly after the latest user message when the current task's own steps
208 were removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement.
210 <!-- test-ignore -->
211 ```python
212 from haystack.components.agents import Agent
213 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
214 from haystack.hooks.compaction import CompactionHook, SlidingWindowCompactor
216 hook = CompactionHook(
217 compactor=SlidingWindowCompactor(), context_window=400_000, compact_at=0.7, compact_to=0.4
218 )
219 agent = Agent(
220 chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.4-nano"),
221 tools=[web_search],
222 hooks={"before_llm": [hook]},
223 )
224 ```
225 """
227 def __init__(self, *, min_keep_steps: int = 1, omission_note: str | None = _DEFAULT_OMISSION_NOTE) -> None:
228 """
229 Initialize the compactor.
231 :param min_keep_steps: The fewest complete recent Agent steps to keep even when they exceed the target. A step
232 is an assistant message and all immediately following tool results. `0` allows all completed steps to be
233 removed when none fit.
234 :param omission_note: The user message left in place of what was removed, or None to remove the messages
235 silently. Include `{num_removed}` to have the number of removed messages substituted in.
236 :raises ValueError: If `min_keep_steps` is negative.
237 """
238 if min_keep_steps < 0:
239 raise ValueError(f"`min_keep_steps` must be at least 0, got {min_keep_steps}.")
240 self.min_keep_steps = min_keep_steps
241 self.omission_note = omission_note
243 def compact(
244 self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter
245 ) -> list[ChatMessage] | None:
246 """
247 Drop older history while preserving the task anchor and a complete recent conversation window.
249 :param messages: The conversation to compact, oldest to newest.
250 :param target_tokens: The size the kept conversation should come in under.
251 :param token_counter: The `TokenCounter` to measure messages with.
252 :returns: The conversation that survived, with an omission note if configured standing where the removed
253 messages used to sit; or None when there is nothing to remove but an earlier note.
254 """
255 if token_counter.count(messages=messages) <= target_tokens:
256 return None
257 kept, note_index, removable = _task_and_step_split(
258 messages=messages,
259 target_tokens=target_tokens,
260 token_counter=token_counter,
261 min_keep_steps=self.min_keep_steps,
262 )
263 if not removable:
264 return None
265 if not self.omission_note:
266 return kept
267 # Swapping an earlier note for a new one frees nothing, so decline instead of rewriting the conversation again
268 # on every following step.
269 if all(_is_compaction_note(message=message) for message in removable):
270 return None
272 # We prefer user over system since not all providers support multiple system messages
273 note = ChatMessage.from_user(
274 # `replace` rather than `format`, so a note carrying braces of its own cannot raise.
275 self.omission_note.replace(_NUM_REMOVED_PLACEHOLDER, str(len(removable))),
276 meta={
277 _COMPACTION_META_KEY: {
278 "strategy": _STRATEGY,
279 "removed_messages": len(removable),
280 "kept_messages": len(kept),
281 }
282 },
283 )
284 return [*kept[:note_index], note, *kept[note_index:]]
286 def to_dict(self) -> dict[str, Any]:
287 """
288 Serialize the compactor.
290 :returns: A dictionary representation of the compactor.
291 """
292 return default_to_dict(self, min_keep_steps=self.min_keep_steps, omission_note=self.omission_note)