From 5b0ecf5cfb5c3db1d524b8882533c85bee622a7c Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 13:25:07 +0200 Subject: [PATCH 01/38] Add support for keeping old user-assistant turns --- haystack/hooks/compaction/sliding_window.py | 105 +++++++++++++----- ...t-context-compaction-3258c08dec9d2b34.yaml | 12 +- test/hooks/compaction/test_sliding_window.py | 83 +++++++++++++- 3 files changed, 169 insertions(+), 31 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 9cba0dc924c..97c7e81240e 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -44,59 +44,114 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: return None +def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: + """ + Return spans for complete user turns in a bounded section of conversation history. + + Each turn begins with a real user message and continues up to, but does not include, the next real user message. + This groups a user's request with every assistant step and tool result produced in response to it. + + :param messages: The full conversation to analyze, ordered oldest to newest. + :param start: The inclusive index at which to begin looking for historical turns. + :param end: The exclusive index at which to stop. This is normally the current task's user-message index. + :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to + `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. + """ + # Compaction notes use the user role for provider compatibility, but they do not begin a new conversation turn. + # Ignoring marked messages here also lets a subsequent compaction fold an old note into its replacement. + user_indices = [ + index + for index in range(start, end) + if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta + ] + + # A real user message closes the preceding turn and starts the next one. The final historical turn extends to the + # supplied boundary, which is typically where the protected current task begins. + return [ + (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) + for position, index in enumerate(user_indices) + ] + + def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int ) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage], int]: - """Split messages into the protected task context, removable history, and retained Agent steps.""" + """Split messages into the protected prefix, removable history, and retained conversation window.""" # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) # Find the latest user message to use as the current task anchor. task_index = _latest_user_index(messages=messages) task = [messages[task_index]] if task_index is not None else [] - # Find complete Agent steps after the current task, or after the system messages when there is no task anchor. - steps = _agent_step_spans(messages=messages, start=(task_index + 1) if task_index is not None else system_end) + step_start = (task_index + 1) if task_index is not None else system_end + # Current-task steps can be removed individually. Earlier user/assistant exchanges are kept as complete turns so + # an assistant reply is never retained without the user message it answers. + steps = _agent_step_spans(messages=messages, start=step_start) + historical_end = task_index if task_index is not None else system_end + historical_turns = _historical_turn_spans(messages, system_end, historical_end) # Protect the Agent instructions and current task from removal. protected = [*messages[:system_end], *task] - # The remaining token budget to retain recent Agent steps. + # The remaining token budget after protecting the instructions and current task. available_tokens = target_tokens - token_counter.count(messages=protected) - # Work backwards through the steps, keeping as many as fit in the remaining budget. - kept_step_start = len(steps) - while kept_step_start > 0: - start, end = steps[kept_step_start - 1] - step_tokens = token_counter.count(messages=messages[start:end]) - # Stop at the first step that does not fit - if step_tokens > available_tokens: - break - available_tokens -= step_tokens - kept_step_start -= 1 + kept_turn_start = len(historical_turns) + all_step_tokens = token_counter.count(messages=[message for start, end in steps for message in messages[start:end]]) + if all_step_tokens <= available_tokens: + # Historical turns are considered only when the entire current task fits. This ensures that compaction removes + # every older turn before it starts trimming individual steps from the task the Agent is actively working on. + kept_step_start = 0 + available_tokens -= all_step_tokens + while kept_turn_start > 0: + start, end = historical_turns[kept_turn_start - 1] + turn = [message for message in messages[start:end] if _COMPACTION_META_KEY not in message.meta] + turn_tokens = token_counter.count(messages=turn) + if turn_tokens > available_tokens: + break + available_tokens -= turn_tokens + kept_turn_start -= 1 + else: + # Even after dropping every historical turn, the current task is too large. Work backwards through its Agent + # steps and retain the most recent complete suffix that fits. + kept_step_start = len(steps) + while kept_step_start > 0: + start, end = steps[kept_step_start - 1] + step_tokens = token_counter.count(messages=messages[start:end]) + if step_tokens > available_tokens: + break + available_tokens -= step_tokens + kept_step_start -= 1 # Enforce the minimum number of complete steps, even when they exceed the target token budget. kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) - kept_spans = steps[kept_step_start:] + kept_step_spans = steps[kept_step_start:] + kept_turn_spans = historical_turns[kept_turn_start:] # Record every protected or retained message index; equal ChatMessages can appear more than once in the list. kept_indices = {*range(system_end)} if task_index is not None: kept_indices.add(task_index) - for start, end in kept_spans: - kept_indices.update(range(start, end)) + for start, end in [*kept_turn_spans, *kept_step_spans]: + kept_indices.update(index for index in range(start, end) if _COMPACTION_META_KEY not in messages[index].meta) # Everything outside the protected context and retained steps can be removed. removable = [message for index, message in enumerate(messages) if index not in kept_indices] - # Flatten the retained spans back into the message list expected by the compactor. - kept_steps = [message for start, end in kept_spans for message in messages[start:end]] - return protected, removable, kept_steps, len(kept_spans) + kept_turns = [ + message + for start, end in kept_turn_spans + for message in messages[start:end] + if _COMPACTION_META_KEY not in message.meta + ] + kept_steps = [message for start, end in kept_step_spans for message in messages[start:end]] + return [*messages[:system_end], *kept_turns, *task], removable, kept_steps, len(steps) - kept_step_start @_experimental class SlidingWindowCompactor(Compactor): """ - Keeps the Agent's instructions, current task, and as many complete recent steps as the target allows. + Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows. - Leading system messages and the latest user message are protected. Recent history is retained in complete Agent - steps, where a step is an assistant message together with all immediately following tool results. An - `omission_note` is left in place of what was removed. + Leading system messages and the latest user message are protected. Earlier user/assistant turns are retained when + they fit, and the current task's history is retained in complete Agent steps, where a step is an assistant message + together with all immediately following tool results. An `omission_note` is left in place of what was removed. ```python from haystack.components.agents import Agent @@ -134,7 +189,7 @@ def compact( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter ) -> list[ChatMessage] | None: """ - Drop older history while preserving the task anchor and a recent window of complete Agent steps. + Drop older history while preserving the task anchor and a complete recent conversation window. :param messages: The conversation to compact, oldest to newest. :param target_tokens: The size the retained conversation should come in under. diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index 8e0a4810a64..c71c803f853 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -5,8 +5,9 @@ features: shortens the conversation when it reaches a configured fraction of the model's context window. The first built-in strategy, ``SlidingWindowCompactor``, preserves leading system messages, the latest user task, - and as many complete recent Agent steps as the target allows. It replaces removed history with a short omission - note. + and as much complete recent conversation as the target allows. It removes earlier user/assistant turns as whole + units first. Only when removing all earlier turns is insufficient does it remove individual Agent steps from the + current task. It replaces removed history with a short omission note. .. code-block:: python @@ -30,9 +31,10 @@ features: tool schemas. Leave headroom above ``compact_at`` for the next reply and its tool results. ``SlidingWindowCompactor`` treats an assistant message and its following tool results as one step, so a tool call is - never separated from its results. It may retain slightly more history than the requested target when preserving the - current task or the configured minimum number of steps. Compaction is lossy: removed messages cannot be recovered - or summarized by this strategy. Implement the ``Compactor`` protocol to provide a custom strategy. + never separated from its results. Historical turns are also kept whole, so an assistant reply is not retained + without the user message it answers. It may retain slightly more history than the requested target when preserving + the current task or the configured minimum number of steps. Compaction is lossy: removed messages cannot be + recovered or summarized by this strategy. Implement the ``Compactor`` protocol to provide a custom strategy. ``CompactionHook`` and ``SlidingWindowCompactor`` emit an ``ExperimentalWarning`` and may change without a deprecation cycle. diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index d395a32f795..635f4407d73 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -6,7 +6,7 @@ from haystack.dataclasses import ChatMessage, ChatRole from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import FakeCounter, count_markers, long_conversation, tool_call, tool_result @@ -18,6 +18,46 @@ COUNTER = FakeCounter() +class TestHistoricalTurnSpans: + def test_groups_each_user_message_with_its_assistant_steps_and_tool_results(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("first task"), + tool_call("c1"), + tool_result("first result", call_id="c1"), + ChatMessage.from_assistant("first answer"), + ChatMessage.from_user("second task"), + ChatMessage.from_assistant("second answer"), + ] + spans = _historical_turn_spans(messages=messages, start=1, end=len(messages)) + assert spans == [(1, 5), (5, 7)] + assert messages[slice(*spans[0])] == messages[1:5] + assert messages[slice(*spans[1])] == messages[5:7] + + def test_only_returns_turns_within_the_requested_bounds(self): + messages = [ + ChatMessage.from_user("outside"), + ChatMessage.from_assistant("outside answer"), + ChatMessage.from_user("inside"), + ChatMessage.from_assistant("inside answer"), + ChatMessage.from_user("current task"), + ] + assert _historical_turn_spans(messages=messages, start=2, end=4) == [(2, 4)] + + def test_compaction_note_does_not_start_a_new_turn(self): + note = ChatMessage.from_user( + "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ) + messages = [ + ChatMessage.from_user("task"), + ChatMessage.from_assistant("first step"), + note, + ChatMessage.from_assistant("second step"), + ChatMessage.from_user("next task"), + ] + assert _historical_turn_spans(messages=messages, start=0, end=len(messages)) == [(0, 4), (4, 5)] + + class TestSlidingWindowCompactor: def test_replaces_the_middle_with_an_omission_note(self): messages = long_conversation() @@ -48,6 +88,47 @@ def test_a_roomier_target_keeps_more(self): assert len(tight) == 4 assert len(roomy) == 8 + def test_keeps_complete_recent_user_assistant_turns_that_fit(self): + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user(text="old question"), + ChatMessage.from_assistant(text="old answer"), + ChatMessage.from_user(text="recent question"), + ChatMessage.from_assistant(text="recent answer"), + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Drops one historical turn + expected = [messages[0], *messages[3:]] + target_tokens = COUNTER.count(expected) + compacted = SlidingWindowCompactor(omission_note=None).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER + ) + assert compacted == expected + + def test_drops_historical_context_and_one_current_task_step_to_reach_target(self): + system_message = ChatMessage.from_system(text="rules") + historical_turn = [ + ChatMessage.from_user(text="old question"), + tool_call("old-call"), + tool_result(result="old result", call_id="old-call"), + ChatMessage.from_assistant(text="old final answer"), + ] + current_task = [ + ChatMessage.from_user(text="current task"), + tool_call("current-call-1"), + tool_result(result="large intermediate result " * 100, call_id="current-call-1"), + tool_call("current-call-2"), + tool_result(result="latest result", call_id="current-call-2"), + ] + messages = [system_message, *historical_turn, *current_task] + # The target is small enough that the historical context and one step in the current task must be removed. + target_tokens = 52 + compacted = SlidingWindowCompactor(omission_note=None).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER + ) + assert compacted == [system_message, current_task[0], *current_task[-2:]] + def test_returns_none_when_the_conversation_already_fits(self): assert ( SlidingWindowCompactor().compact(messages=long_conversation(), target_tokens=100_000, token_counter=COUNTER) From c8323e244d3484b429bc75169b3f29613c58405e Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 13:39:38 +0200 Subject: [PATCH 02/38] refactoring --- haystack/hooks/compaction/sliding_window.py | 179 ++++++++++++++----- test/hooks/compaction/test_sliding_window.py | 14 ++ 2 files changed, 148 insertions(+), 45 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 97c7e81240e..f3ed181bacb 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -73,10 +73,120 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> ] +def _messages_from_spans( + messages: list[ChatMessage], spans: list[tuple[int, int]], *, skip_compaction_notes: bool = False +) -> list[ChatMessage]: + """Flatten message spans, optionally excluding messages produced by an earlier compaction.""" + return [ + message + for start, end in spans + for message in messages[start:end] + if not skip_compaction_notes or _COMPACTION_META_KEY not in message.meta + ] + + +def _fitting_suffix_start( + messages: list[ChatMessage], + spans: list[tuple[int, int]], + available_tokens: int, + token_counter: TokenCounter, + *, + skip_compaction_notes: bool = False, +) -> int: + """ + Return the first span in the newest contiguous suffix that fits the token budget. + + Spans are measured from newest to oldest. Retention stops as soon as a span does not fit, ensuring that an older + span is never kept after a newer one has been removed. + + :param messages: The full conversation containing the messages referenced by `spans`. + :param spans: Ordered `(start_index, end_index)` pairs to consider for retention. Both indices refer to `messages`, + and `end_index` is exclusive. + :param available_tokens: The token budget available for retaining messages from `spans`. + :param token_counter: The `TokenCounter` used to measure each span. + :param skip_compaction_notes: Whether messages produced by an earlier compaction are excluded from token counting. + :returns: The index in `spans` at which the retained suffix begins. If no span fits, returns `len(spans)`; if every + span fits, returns `0`. + """ + kept_start = len(spans) + while kept_start > 0: + span_messages = _messages_from_spans( + messages=messages, spans=[spans[kept_start - 1]], skip_compaction_notes=skip_compaction_notes + ) + span_tokens = token_counter.count(messages=span_messages) + if span_tokens > available_tokens: + break + available_tokens -= span_tokens + kept_start -= 1 + return kept_start + + +def _retained_span_starts( + messages: list[ChatMessage], + protected: list[ChatMessage], + historical_turns: list[tuple[int, int]], + steps: list[tuple[int, int]], + target_tokens: int, + token_counter: TokenCounter, + min_keep_steps: int, +) -> tuple[int, int]: + """Return the first retained historical turn and current-task step.""" + available_tokens = target_tokens - token_counter.count(messages=protected) + all_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=steps, skip_compaction_notes=False) + ) + kept_turn_start = len(historical_turns) + if all_step_tokens <= available_tokens: + # Historical turns are considered only when the entire current task fits. This ensures that compaction removes + # every older turn before it starts trimming individual steps from the task the Agent is actively working on. + kept_step_start = 0 + kept_turn_start = _fitting_suffix_start( + messages=messages, + spans=historical_turns, + available_tokens=available_tokens - all_step_tokens, + token_counter=token_counter, + skip_compaction_notes=True, + ) + else: + # Even after dropping every historical turn, the current task is too large. Retain the most recent complete + # suffix of Agent steps that fits. + kept_step_start = _fitting_suffix_start( + messages=messages, + spans=steps, + available_tokens=available_tokens, + token_counter=token_counter, + skip_compaction_notes=False, + ) + kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) + return kept_turn_start, kept_step_start + + def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int ) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage], int]: - """Split messages into the protected prefix, removable history, and retained conversation window.""" + """ + Split a conversation into the messages kept before and after an omission note and the messages to remove. + + Leading system messages and the latest real user message are always retained. Historical user turns are retained + whole when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed + individually while preserving complete assistant/tool-result groups. + + :param messages: The full conversation to split, ordered oldest to newest. + :param target_tokens: The target token budget for the retained messages. + :param token_counter: The `TokenCounter` used to decide which historical turns and current-task steps fit. + :param min_keep_steps: The minimum number of recent current-task Agent steps to retain, even if they exceed the + target token budget. + :returns: A tuple containing: + + 1. Messages retained before the omission-note position. + 2. Every message selected for removal. + 3. Messages retained after the omission-note position. + 4. The number of retained current-task Agent steps. + + When only historical turns are removed, the first and third elements place the note immediately before the + current task. When current-task steps are also removed, they place it after the latest user message and before + the retained current-task steps. + """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) # Find the latest user message to use as the current task anchor. @@ -87,41 +197,19 @@ def _task_and_step_split( # an assistant reply is never retained without the user message it answers. steps = _agent_step_spans(messages=messages, start=step_start) historical_end = task_index if task_index is not None else system_end - historical_turns = _historical_turn_spans(messages, system_end, historical_end) + historical_turns = _historical_turn_spans(messages=messages, start=system_end, end=historical_end) # Protect the Agent instructions and current task from removal. protected = [*messages[:system_end], *task] - # The remaining token budget after protecting the instructions and current task. - available_tokens = target_tokens - token_counter.count(messages=protected) - kept_turn_start = len(historical_turns) - all_step_tokens = token_counter.count(messages=[message for start, end in steps for message in messages[start:end]]) - if all_step_tokens <= available_tokens: - # Historical turns are considered only when the entire current task fits. This ensures that compaction removes - # every older turn before it starts trimming individual steps from the task the Agent is actively working on. - kept_step_start = 0 - available_tokens -= all_step_tokens - while kept_turn_start > 0: - start, end = historical_turns[kept_turn_start - 1] - turn = [message for message in messages[start:end] if _COMPACTION_META_KEY not in message.meta] - turn_tokens = token_counter.count(messages=turn) - if turn_tokens > available_tokens: - break - available_tokens -= turn_tokens - kept_turn_start -= 1 - else: - # Even after dropping every historical turn, the current task is too large. Work backwards through its Agent - # steps and retain the most recent complete suffix that fits. - kept_step_start = len(steps) - while kept_step_start > 0: - start, end = steps[kept_step_start - 1] - step_tokens = token_counter.count(messages=messages[start:end]) - if step_tokens > available_tokens: - break - available_tokens -= step_tokens - kept_step_start -= 1 - - # Enforce the minimum number of complete steps, even when they exceed the target token budget. - kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) + kept_turn_start, kept_step_start = _retained_span_starts( + messages=messages, + protected=protected, + historical_turns=historical_turns, + steps=steps, + target_tokens=target_tokens, + token_counter=token_counter, + min_keep_steps=min_keep_steps, + ) kept_step_spans = steps[kept_step_start:] kept_turn_spans = historical_turns[kept_turn_start:] @@ -134,14 +222,15 @@ def _task_and_step_split( # Everything outside the protected context and retained steps can be removed. removable = [message for index, message in enumerate(messages) if index not in kept_indices] - kept_turns = [ - message - for start, end in kept_turn_spans - for message in messages[start:end] - if _COMPACTION_META_KEY not in message.meta - ] - kept_steps = [message for start, end in kept_step_spans for message in messages[start:end]] - return [*messages[:system_end], *kept_turns, *task], removable, kept_steps, len(steps) - kept_step_start + kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) + kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) + if kept_step_start == 0: + kept_before_note = [*messages[:system_end], *kept_turns] + kept_after_note = [*task, *kept_steps] + else: + kept_before_note = [*messages[:system_end], *task] + kept_after_note = kept_steps + return kept_before_note, removable, kept_after_note, len(steps) - kept_step_start @_experimental @@ -199,7 +288,7 @@ def compact( """ if token_counter.count(messages) <= target_tokens: return None - protected, removable, kept_steps, kept_step_count = _task_and_step_split( + kept_before_note, removable, kept_after_note, kept_step_count = _task_and_step_split( messages=messages, target_tokens=target_tokens, token_counter=token_counter, @@ -208,7 +297,7 @@ def compact( if not removable: return None if not self.omission_note: - return [*protected, *kept_steps] + return [*kept_before_note, *kept_after_note] # We prefer user over system since not all providers support multiple system messages note = ChatMessage.from_user( @@ -218,7 +307,7 @@ def compact( _COMPACTION_META_KEY: { "strategy": "sliding_window", "removed_messages": len(removable), - "kept_messages": len(protected) + len(kept_steps), + "kept_messages": len(kept_before_note) + len(kept_after_note), "kept_steps": kept_step_count, } }, @@ -226,7 +315,7 @@ def compact( # The note costs tokens of its own, so it is only worth leaving behind if what it stands in for is bigger. if token_counter.count([note]) >= token_counter.count(removable): return None - return [*protected, note, *kept_steps] + return [*kept_before_note, note, *kept_after_note] def to_dict(self) -> dict[str, Any]: """ diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 635f4407d73..4ba9627283e 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -76,6 +76,20 @@ def test_replaces_the_middle_with_an_omission_note(self): # A user message, not a system one, so providers that hoist system messages cannot move it out of position. assert compacted[2].is_from(role=ChatRole.USER) + def test_places_omission_note_before_current_task_when_only_historical_turns_are_removed(self): + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user(text="old question"), + ChatMessage.from_assistant(text="old answer " * 100), + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current answer"), + ] + compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=20, token_counter=COUNTER) + assert compacted is not None + assert compacted[0] == messages[0] + assert _COMPACTION_META_KEY in compacted[1].meta + assert compacted[2:] == messages[3:] + def test_a_roomier_target_keeps_more(self): messages = [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="task")] for index in range(4): From b78e3f1d48a67d9005ca5325680b29ab4458085f Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 13:58:05 +0200 Subject: [PATCH 03/38] making the logic less insane --- haystack/hooks/compaction/sliding_window.py | 100 ++++++++++++-------- 1 file changed, 62 insertions(+), 38 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index f3ed181bacb..fe8636182dc 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -121,44 +121,55 @@ def _fitting_suffix_start( return kept_start -def _retained_span_starts( +def _get_turn_start( messages: list[ChatMessage], - protected: list[ChatMessage], + available_tokens: int, historical_turns: list[tuple[int, int]], - steps: list[tuple[int, int]], - target_tokens: int, + current_agent_steps: list[tuple[int, int]], token_counter: TokenCounter, - min_keep_steps: int, -) -> tuple[int, int]: - """Return the first retained historical turn and current-task step.""" - available_tokens = target_tokens - token_counter.count(messages=protected) - all_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=steps, skip_compaction_notes=False) +) -> int: + """Return the start of the retained historical turns, or `len(historical_turns)` if none fit.""" + all_current_agent_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) ) kept_turn_start = len(historical_turns) - if all_step_tokens <= available_tokens: + if all_current_agent_step_tokens <= available_tokens: # Historical turns are considered only when the entire current task fits. This ensures that compaction removes # every older turn before it starts trimming individual steps from the task the Agent is actively working on. - kept_step_start = 0 kept_turn_start = _fitting_suffix_start( messages=messages, spans=historical_turns, - available_tokens=available_tokens - all_step_tokens, + available_tokens=available_tokens - all_current_agent_step_tokens, token_counter=token_counter, skip_compaction_notes=True, ) - else: - # Even after dropping every historical turn, the current task is too large. Retain the most recent complete - # suffix of Agent steps that fits. - kept_step_start = _fitting_suffix_start( - messages=messages, - spans=steps, - available_tokens=available_tokens, - token_counter=token_counter, - skip_compaction_notes=False, - ) - kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) - return kept_turn_start, kept_step_start + return kept_turn_start + + +def _get_step_start( + messages: list[ChatMessage], + available_tokens: int, + current_agent_steps: list[tuple[int, int]], + token_counter: TokenCounter, + min_keep_steps: int, +) -> int: + """Return the start of the retained current-task steps, or `len(current_agent_steps)` if none fit.""" + all_current_agent_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) + ) + if all_current_agent_step_tokens <= available_tokens: + # Since all current-task steps fit, we can retain all of them. + return 0 + # Even after dropping every historical turn, the current task is too large. Retain the most recent complete + # suffix of Agent steps that fits. + kept_step_start = _fitting_suffix_start( + messages=messages, + spans=current_agent_steps, + available_tokens=available_tokens, + token_counter=token_counter, + skip_compaction_notes=False, + ) + return min(kept_step_start, max(len(current_agent_steps) - min_keep_steps, 0)) def _task_and_step_split( @@ -189,30 +200,43 @@ def _task_and_step_split( """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) + # Find the latest user message to use as the current task anchor. task_index = _latest_user_index(messages=messages) task = [messages[task_index]] if task_index is not None else [] - step_start = (task_index + 1) if task_index is not None else system_end - # Current-task steps can be removed individually. Earlier user/assistant exchanges are kept as complete turns so - # an assistant reply is never retained without the user message it answers. - steps = _agent_step_spans(messages=messages, start=step_start) + + # Find the complete Agent steps that follow the current task anchor. + current_task_step_start = (task_index + 1) if task_index is not None else system_end + current_agent_steps = _agent_step_spans(messages=messages, start=current_task_step_start) + + # Find all complete historical turns (i.e. user-assistant) that precede the current task. historical_end = task_index if task_index is not None else system_end historical_turns = _historical_turn_spans(messages=messages, start=system_end, end=historical_end) - # Protect the Agent instructions and current task from removal. - protected = [*messages[:system_end], *task] - kept_turn_start, kept_step_start = _retained_span_starts( + # Calculate the token count of the protected context (leading system messages and current task) + protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + + # Get the start of where we should retain historical turns. + # If no turns are kept the value of this is `len(historical_turns)`, which is the same as `historical_end`. + kept_turn_start = _get_turn_start( messages=messages, - protected=protected, + available_tokens=target_tokens - protected_tokens, historical_turns=historical_turns, - steps=steps, - target_tokens=target_tokens, + current_agent_steps=current_agent_steps, token_counter=token_counter, - min_keep_steps=min_keep_steps, ) - kept_step_spans = steps[kept_step_start:] kept_turn_spans = historical_turns[kept_turn_start:] + # Get the start of where we should retain current-task steps + kept_step_start = _get_step_start( + messages=messages, + available_tokens=target_tokens - protected_tokens, + current_agent_steps=current_agent_steps, + token_counter=token_counter, + min_keep_steps=min_keep_steps, + ) + kept_step_spans = current_agent_steps[kept_step_start:] + # Record every protected or retained message index; equal ChatMessages can appear more than once in the list. kept_indices = {*range(system_end)} if task_index is not None: @@ -230,7 +254,7 @@ def _task_and_step_split( else: kept_before_note = [*messages[:system_end], *task] kept_after_note = kept_steps - return kept_before_note, removable, kept_after_note, len(steps) - kept_step_start + return kept_before_note, removable, kept_after_note, len(current_agent_steps) - kept_step_start @_experimental From 1565f0c2e95c7e0ff95d33706153f8f929a61f16 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:07:13 +0200 Subject: [PATCH 04/38] more logic refactoring --- haystack/hooks/compaction/sliding_window.py | 34 ++++++++++----------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index fe8636182dc..00640f7d226 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -74,7 +74,7 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> def _messages_from_spans( - messages: list[ChatMessage], spans: list[tuple[int, int]], *, skip_compaction_notes: bool = False + messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False ) -> list[ChatMessage]: """Flatten message spans, optionally excluding messages produced by an earlier compaction.""" return [ @@ -90,8 +90,7 @@ def _fitting_suffix_start( spans: list[tuple[int, int]], available_tokens: int, token_counter: TokenCounter, - *, - skip_compaction_notes: bool = False, + skip_compaction_notes: bool, ) -> int: """ Return the first span in the newest contiguous suffix that fits the token budget. @@ -124,14 +123,11 @@ def _fitting_suffix_start( def _get_turn_start( messages: list[ChatMessage], available_tokens: int, + all_current_agent_step_tokens: int, historical_turns: list[tuple[int, int]], - current_agent_steps: list[tuple[int, int]], token_counter: TokenCounter, ) -> int: """Return the start of the retained historical turns, or `len(historical_turns)` if none fit.""" - all_current_agent_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) - ) kept_turn_start = len(historical_turns) if all_current_agent_step_tokens <= available_tokens: # Historical turns are considered only when the entire current task fits. This ensures that compaction removes @@ -149,14 +145,12 @@ def _get_turn_start( def _get_step_start( messages: list[ChatMessage], available_tokens: int, + all_current_agent_step_tokens: int, current_agent_steps: list[tuple[int, int]], token_counter: TokenCounter, min_keep_steps: int, ) -> int: """Return the start of the retained current-task steps, or `len(current_agent_steps)` if none fit.""" - all_current_agent_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) - ) if all_current_agent_step_tokens <= available_tokens: # Since all current-task steps fit, we can retain all of them. return 0 @@ -216,38 +210,44 @@ def _task_and_step_split( # Calculate the token count of the protected context (leading system messages and current task) protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + # Calculate the size of the current task's Agent steps + all_current_agent_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) + ) + # Get the start of where we should retain historical turns. # If no turns are kept the value of this is `len(historical_turns)`, which is the same as `historical_end`. kept_turn_start = _get_turn_start( messages=messages, available_tokens=target_tokens - protected_tokens, + all_current_agent_step_tokens=all_current_agent_step_tokens, historical_turns=historical_turns, - current_agent_steps=current_agent_steps, token_counter=token_counter, ) kept_turn_spans = historical_turns[kept_turn_start:] + kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) # Get the start of where we should retain current-task steps kept_step_start = _get_step_start( messages=messages, available_tokens=target_tokens - protected_tokens, + all_current_agent_step_tokens=all_current_agent_step_tokens, current_agent_steps=current_agent_steps, token_counter=token_counter, min_keep_steps=min_keep_steps, ) kept_step_spans = current_agent_steps[kept_step_start:] + kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) - # Record every protected or retained message index; equal ChatMessages can appear more than once in the list. + # Record every index that we are keeping kept_indices = {*range(system_end)} if task_index is not None: kept_indices.add(task_index) for start, end in [*kept_turn_spans, *kept_step_spans]: - kept_indices.update(index for index in range(start, end) if _COMPACTION_META_KEY not in messages[index].meta) - - # Everything outside the protected context and retained steps can be removed. + kept_indices.update(index for index in range(start, end)) + # Remove everything else that's not in the kept indices removable = [message for index, message in enumerate(messages) if index not in kept_indices] - kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) - kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) + if kept_step_start == 0: kept_before_note = [*messages[:system_end], *kept_turns] kept_after_note = [*task, *kept_steps] From a3027490e5ee6c8094ef3bc680465b957370419d Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:29:03 +0200 Subject: [PATCH 05/38] simplifications --- haystack/hooks/compaction/sliding_window.py | 237 +++++++++---------- test/hooks/compaction/test_sliding_window.py | 25 +- 2 files changed, 135 insertions(+), 127 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 00640f7d226..d2b95090c90 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -73,124 +73,124 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> ] -def _messages_from_spans( +def _index_groups( messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False -) -> list[ChatMessage]: - """Flatten message spans, optionally excluding messages produced by an earlier compaction.""" +) -> list[list[int]]: + """ + Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. + """ return [ - message + [ + index + for index in range(start, end) + if not skip_compaction_notes or _COMPACTION_META_KEY not in messages[index].meta + ] for start, end in spans - for message in messages[start:end] - if not skip_compaction_notes or _COMPACTION_META_KEY not in message.meta ] -def _fitting_suffix_start( - messages: list[ChatMessage], - spans: list[tuple[int, int]], - available_tokens: int, - token_counter: TokenCounter, - skip_compaction_notes: bool, +def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages at the given indices, in conversation order.""" + return [messages[index] for index in indices] + + +def _flatten(groups: list[list[int]]) -> list[int]: + """Join index groups into a single ordered list of indices.""" + return [index for group in groups for index in group] + + +def _first_group_to_keep( + messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter ) -> int: """ - Return the first span in the newest contiguous suffix that fits the token budget. - - Spans are measured from newest to oldest. Retention stops as soon as a span does not fit, ensuring that an older - span is never kept after a newer one has been removed. - - :param messages: The full conversation containing the messages referenced by `spans`. - :param spans: Ordered `(start_index, end_index)` pairs to consider for retention. Both indices refer to `messages`, - and `end_index` is exclusive. - :param available_tokens: The token budget available for retaining messages from `spans`. - :param token_counter: The `TokenCounter` used to measure each span. - :param skip_compaction_notes: Whether messages produced by an earlier compaction are excluded from token counting. - :returns: The index in `spans` at which the retained suffix begins. If no span fits, returns `len(spans)`; if every - span fits, returns `0`. + Return the oldest group that the budget can still pay for, working backwards from the newest. + + Groups are added up from newest to oldest and counting stops at the first group that does not fit, so what is kept + is always a run of groups at the end of the list. An older group is never kept once a newer one has been dropped, + which would leave a hole in the conversation. + + :param messages: The full conversation containing the messages referenced by `groups`. + :param groups: Ordered index groups to choose from, oldest group first. + :param available_tokens: The token budget available for these groups. + :param token_counter: The `TokenCounter` used to measure each group. + :returns: The position in `groups` to start keeping from: `len(groups)` when nothing fits, `0` when it all fits. """ - kept_start = len(spans) - while kept_start > 0: - span_messages = _messages_from_spans( - messages=messages, spans=[spans[kept_start - 1]], skip_compaction_notes=skip_compaction_notes - ) - span_tokens = token_counter.count(messages=span_messages) - if span_tokens > available_tokens: + first_kept = len(groups) + while first_kept > 0: + group_tokens = token_counter.count(messages=_messages_at(messages=messages, indices=groups[first_kept - 1])) + if group_tokens > available_tokens: break - available_tokens -= span_tokens - kept_start -= 1 - return kept_start + available_tokens -= group_tokens + first_kept -= 1 + return first_kept -def _get_turn_start( +def _first_turn_and_step_to_keep( messages: list[ChatMessage], + turn_groups: list[list[int]], + step_groups: list[list[int]], available_tokens: int, - all_current_agent_step_tokens: int, - historical_turns: list[tuple[int, int]], token_counter: TokenCounter, -) -> int: - """Return the start of the retained historical turns, or `len(historical_turns)` if none fit.""" - kept_turn_start = len(historical_turns) - if all_current_agent_step_tokens <= available_tokens: - # Historical turns are considered only when the entire current task fits. This ensures that compaction removes - # every older turn before it starts trimming individual steps from the task the Agent is actively working on. - kept_turn_start = _fitting_suffix_start( - messages=messages, - spans=historical_turns, - available_tokens=available_tokens - all_current_agent_step_tokens, - token_counter=token_counter, - skip_compaction_notes=True, + min_keep_steps: int, +) -> tuple[int, int]: + """ + Return which historical turn and which Agent step of the current task to start keeping from. + + :param messages: The full conversation containing the messages referenced by both group lists. + :param turn_groups: Index groups for the complete historical turns preceding the current task, oldest first. + :param step_groups: Index groups for the current task's Agent steps, oldest first. + :param available_tokens: The token budget left once the protected context is paid for. + :param token_counter: The `TokenCounter` used to measure the groups. + :param min_keep_steps: The fewest recent Agent steps to keep, even when they exceed the budget. + :returns: The position in `turn_groups` and the position in `step_groups` to start keeping from. Either is the + length of its list when nothing from it is kept. + """ + current_task_tokens = token_counter.count( + messages=_messages_at(messages=messages, indices=_flatten(groups=step_groups)) + ) + if current_task_tokens > available_tokens: + # The current task alone overruns the budget, so every historical turn is dropped and + # the current task's own oldest steps trimmed until what remains fits. + first_kept_step = _first_group_to_keep( + messages=messages, groups=step_groups, available_tokens=available_tokens, token_counter=token_counter ) - return kept_turn_start + # The newest steps are kept regardless of the budget. + return len(turn_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) - -def _get_step_start( - messages: list[ChatMessage], - available_tokens: int, - all_current_agent_step_tokens: int, - current_agent_steps: list[tuple[int, int]], - token_counter: TokenCounter, - min_keep_steps: int, -) -> int: - """Return the start of the retained current-task steps, or `len(current_agent_steps)` if none fit.""" - if all_current_agent_step_tokens <= available_tokens: - # Since all current-task steps fit, we can retain all of them. - return 0 - # Even after dropping every historical turn, the current task is too large. Retain the most recent complete - # suffix of Agent steps that fits. - kept_step_start = _fitting_suffix_start( + # The whole current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. + first_kept_turn = _first_group_to_keep( messages=messages, - spans=current_agent_steps, - available_tokens=available_tokens, + groups=turn_groups, + available_tokens=available_tokens - current_task_tokens, token_counter=token_counter, - skip_compaction_notes=False, ) - return min(kept_step_start, max(len(current_agent_steps) - min_keep_steps, 0)) + return first_kept_turn, 0 def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int -) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage], int]: +) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage]]: """ Split a conversation into the messages kept before and after an omission note and the messages to remove. - Leading system messages and the latest real user message are always retained. Historical user turns are retained - whole when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed - individually while preserving complete assistant/tool-result groups. + Leading system messages and the latest real user message are always kept. Historical user turns are kept whole when + they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time + while keeping each assistant message together with its tool results. :param messages: The full conversation to split, ordered oldest to newest. - :param target_tokens: The target token budget for the retained messages. + :param target_tokens: The token budget for the messages that are kept. :param token_counter: The `TokenCounter` used to decide which historical turns and current-task steps fit. - :param min_keep_steps: The minimum number of recent current-task Agent steps to retain, even if they exceed the - target token budget. + :param min_keep_steps: The fewest recent current-task Agent steps to keep, even if they exceed the target token + budget. :returns: A tuple containing: - 1. Messages retained before the omission-note position. - 2. Every message selected for removal. - 3. Messages retained after the omission-note position. - 4. The number of retained current-task Agent steps. + 1. Messages kept before the omission-note position. + 2. Messages kept after the omission-note position. + 3. Every message selected for removal. - When only historical turns are removed, the first and third elements place the note immediately before the - current task. When current-task steps are also removed, they place it after the latest user message and before - the retained current-task steps. + When only historical turns are removed, the first two elements place the note immediately before the current + task. When current-task steps are also removed, they place it after the latest user message and before the + current-task steps that survived. """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) @@ -199,62 +199,48 @@ def _task_and_step_split( task_index = _latest_user_index(messages=messages) task = [messages[task_index]] if task_index is not None else [] - # Find the complete Agent steps that follow the current task anchor. - current_task_step_start = (task_index + 1) if task_index is not None else system_end - current_agent_steps = _agent_step_spans(messages=messages, start=current_task_step_start) + # Group the complete Agent steps that follow the current task anchor. + step_start_index = (task_index + 1) if task_index is not None else system_end + step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start_index)) - # Find all complete historical turns (i.e. user-assistant) that precede the current task. + # Group the complete historical turns (i.e. user-assistant) that precede the current task. An earlier compaction's + # note is left out of its turn, so keeping the turn folds that note into the note this compaction leaves behind. historical_end = task_index if task_index is not None else system_end - historical_turns = _historical_turn_spans(messages=messages, start=system_end, end=historical_end) - - # Calculate the token count of the protected context (leading system messages and current task) - protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) - - # Calculate the size of the current task's Agent steps - all_current_agent_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) - ) - - # Get the start of where we should retain historical turns. - # If no turns are kept the value of this is `len(historical_turns)`, which is the same as `historical_end`. - kept_turn_start = _get_turn_start( + turn_groups = _index_groups( messages=messages, - available_tokens=target_tokens - protected_tokens, - all_current_agent_step_tokens=all_current_agent_step_tokens, - historical_turns=historical_turns, - token_counter=token_counter, + spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), + skip_compaction_notes=True, ) - kept_turn_spans = historical_turns[kept_turn_start:] - kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) - # Get the start of where we should retain current-task steps - kept_step_start = _get_step_start( + # The Agent instructions and the current task are never removed, so they are paid for out of the target first. + protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + first_kept_turn, first_kept_step = _first_turn_and_step_to_keep( messages=messages, + turn_groups=turn_groups, + step_groups=step_groups, available_tokens=target_tokens - protected_tokens, - all_current_agent_step_tokens=all_current_agent_step_tokens, - current_agent_steps=current_agent_steps, token_counter=token_counter, min_keep_steps=min_keep_steps, ) - kept_step_spans = current_agent_steps[kept_step_start:] - kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) + kept_turn_indices = _flatten(groups=turn_groups[first_kept_turn:]) + kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) - # Record every index that we are keeping - kept_indices = {*range(system_end)} + # A message survives only by being protected or by falling in a group we are keeping; everything else goes. + kept_indices = {*range(system_end), *kept_turn_indices, *kept_step_indices} if task_index is not None: kept_indices.add(task_index) - for start, end in [*kept_turn_spans, *kept_step_spans]: - kept_indices.update(index for index in range(start, end)) - # Remove everything else that's not in the kept indices removable = [message for index, message in enumerate(messages) if index not in kept_indices] - if kept_step_start == 0: - kept_before_note = [*messages[:system_end], *kept_turns] - kept_after_note = [*task, *kept_steps] + if first_kept_step == 0: + # The current task is untouched, so the note stands in for the older turns and belongs in front of the task. + kept_before_note = [*messages[:system_end], *_messages_at(messages=messages, indices=kept_turn_indices)] + kept_after_note = [*task, *_messages_at(messages=messages, indices=kept_step_indices)] else: + # Steps were cut from the current task, which means every historical turn was already dropped, so the note + # goes between the task and the steps that survived it. kept_before_note = [*messages[:system_end], *task] - kept_after_note = kept_steps - return kept_before_note, removable, kept_after_note, len(current_agent_steps) - kept_step_start + kept_after_note = _messages_at(messages=messages, indices=kept_step_indices) + return kept_before_note, kept_after_note, removable @_experimental @@ -312,7 +298,7 @@ def compact( """ if token_counter.count(messages) <= target_tokens: return None - kept_before_note, removable, kept_after_note, kept_step_count = _task_and_step_split( + kept_before_note, kept_after_note, removable = _task_and_step_split( messages=messages, target_tokens=target_tokens, token_counter=token_counter, @@ -332,7 +318,6 @@ def compact( "strategy": "sliding_window", "removed_messages": len(removable), "kept_messages": len(kept_before_note) + len(kept_after_note), - "kept_steps": kept_step_count, } }, ) diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 4ba9627283e..32754f81ebe 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -70,7 +70,6 @@ def test_replaces_the_middle_with_an_omission_note(self): "strategy": "sliding_window", "removed_messages": 2, "kept_messages": 4, - "kept_steps": 1, } assert compacted[2].text == _DEFAULT_OMISSION_NOTE.replace("{num_removed}", "2") # A user message, not a system one, so providers that hoist system messages cannot move it out of position. @@ -143,6 +142,30 @@ def test_drops_historical_context_and_one_current_task_step_to_reach_target(self ) assert compacted == [system_message, current_task[0], *current_task[-2:]] + def test_folds_an_earlier_note_inside_a_retained_turn_and_counts_it_as_removed(self): + note = ChatMessage.from_user( + "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ) + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user(text="old question " * 200), + ChatMessage.from_assistant(text="old answer"), + ChatMessage.from_user(text="recent question"), + ChatMessage.from_assistant(text="recent answer"), + note, + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Room for everything but the oldest turn, so the turn holding the earlier note is retained around it. + compacted = SlidingWindowCompactor().compact( + messages=messages, target_tokens=COUNTER.count([messages[0], *messages[3:]]), token_counter=COUNTER + ) + assert compacted is not None + # The earlier note is replaced by the new one rather than surviving alongside it, and it counts as removed. + assert count_markers(messages=compacted) == 1 + assert compacted == [messages[0], *messages[3:5], compacted[3], *messages[6:]] + assert compacted[3].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 + def test_returns_none_when_the_conversation_already_fits(self): assert ( SlidingWindowCompactor().compact(messages=long_conversation(), target_tokens=100_000, token_counter=COUNTER) From a8a1f7d9987cd80bac70f61fd4ffb0dbe14bc7ce Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:40:17 +0200 Subject: [PATCH 06/38] more simplification --- haystack/hooks/compaction/sliding_window.py | 49 ++++++++------------ test/hooks/compaction/test_sliding_window.py | 23 --------- 2 files changed, 20 insertions(+), 52 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index d2b95090c90..f18d1fe07a4 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -169,9 +169,9 @@ def _first_turn_and_step_to_keep( def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int -) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage]]: +) -> tuple[list[ChatMessage], int, list[ChatMessage]]: """ - Split a conversation into the messages kept before and after an omission note and the messages to remove. + Split a conversation into the messages to keep and the messages to remove. Leading system messages and the latest real user message are always kept. Historical user turns are kept whole when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time @@ -184,13 +184,10 @@ def _task_and_step_split( budget. :returns: A tuple containing: - 1. Messages kept before the omission-note position. - 2. Messages kept after the omission-note position. + 1. The messages to keep, ordered oldest to newest. + 2. The position in that list where an omission note belongs, immediately before the current task when only + historical turns were removed and immediately before the surviving steps when the task itself was trimmed. 3. Every message selected for removal. - - When only historical turns are removed, the first two elements place the note immediately before the current - task. When current-task steps are also removed, they place it after the latest user message and before the - current-task steps that survived. """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) @@ -224,6 +221,9 @@ def _task_and_step_split( ) kept_turn_indices = _flatten(groups=turn_groups[first_kept_turn:]) kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) + kept_turns = _messages_at(messages=messages, indices=kept_turn_indices) + kept_steps = _messages_at(messages=messages, indices=kept_step_indices) + kept = [*messages[:system_end], *kept_turns, *task, *kept_steps] # A message survives only by being protected or by falling in a group we are keeping; everything else goes. kept_indices = {*range(system_end), *kept_turn_indices, *kept_step_indices} @@ -231,16 +231,10 @@ def _task_and_step_split( kept_indices.add(task_index) removable = [message for index, message in enumerate(messages) if index not in kept_indices] - if first_kept_step == 0: - # The current task is untouched, so the note stands in for the older turns and belongs in front of the task. - kept_before_note = [*messages[:system_end], *_messages_at(messages=messages, indices=kept_turn_indices)] - kept_after_note = [*task, *_messages_at(messages=messages, indices=kept_step_indices)] - else: - # Steps were cut from the current task, which means every historical turn was already dropped, so the note - # goes between the task and the steps that survived it. - kept_before_note = [*messages[:system_end], *task] - kept_after_note = _messages_at(messages=messages, indices=kept_step_indices) - return kept_before_note, kept_after_note, removable + # The note stands in for the newest thing that was dropped. That is the current task's own steps when those were + # cut, in which case every historical turn went too and `kept_turns` is empty; otherwise it is the older turns. + note_index = system_end + len(kept_turns) + (len(task) if first_kept_step > 0 else 0) + return kept, note_index, removable @_experimental @@ -291,14 +285,14 @@ def compact( Drop older history while preserving the task anchor and a complete recent conversation window. :param messages: The conversation to compact, oldest to newest. - :param target_tokens: The size the retained conversation should come in under. + :param target_tokens: The size the kept conversation should come in under. :param token_counter: The `TokenCounter` to measure messages with. - :returns: The protected context, an omission note if configured, and the retained steps; or None when there is - nothing worth removing. + :returns: The protected context, an omission note if configured, and the steps that survived; or None when + there is nothing to remove. """ - if token_counter.count(messages) <= target_tokens: + if token_counter.count(messages=messages) <= target_tokens: return None - kept_before_note, kept_after_note, removable = _task_and_step_split( + kept, note_index, removable = _task_and_step_split( messages=messages, target_tokens=target_tokens, token_counter=token_counter, @@ -307,7 +301,7 @@ def compact( if not removable: return None if not self.omission_note: - return [*kept_before_note, *kept_after_note] + return kept # We prefer user over system since not all providers support multiple system messages note = ChatMessage.from_user( @@ -317,14 +311,11 @@ def compact( _COMPACTION_META_KEY: { "strategy": "sliding_window", "removed_messages": len(removable), - "kept_messages": len(kept_before_note) + len(kept_after_note), + "kept_messages": len(kept), } }, ) - # The note costs tokens of its own, so it is only worth leaving behind if what it stands in for is bigger. - if token_counter.count([note]) >= token_counter.count(removable): - return None - return [*kept_before_note, note, *kept_after_note] + return [*kept[:note_index], note, *kept[note_index:]] def to_dict(self) -> dict[str, Any]: """ diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 32754f81ebe..bafacf91205 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -268,29 +268,6 @@ def test_repeated_compaction_folds_the_previous_note(self): assert second[0].text == "rules" assert second[1].text == "start" - @pytest.mark.parametrize( - ("removable", "worth_replacing"), - [ - pytest.param("a", False, id="cheaper-than-the-note"), - pytest.param("x" * 4000, True, id="dearer-than-the-note"), - ], - ) - def test_a_cut_is_made_only_when_the_note_costs_less_than_what_it_replaces(self, removable, worth_replacing): - # A note is not free, so what matters is the size of what goes rather than how many messages it is: one long - # tool result is worth replacing, one short message is not. - messages = [ - ChatMessage.from_system(text="rules"), - ChatMessage.from_user(text="task"), - ChatMessage.from_assistant(text=removable), - ChatMessage.from_assistant(text="latest step"), - ] - compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) - assert (compacted is not None) is worth_replacing - # Without a note there is nothing to pay for, so the same cut is always worth making. - assert SlidingWindowCompactor(omission_note=None).compact( - messages=messages, target_tokens=SMALLEST, token_counter=COUNTER - ) == [*messages[:2], messages[-1]] - def test_keeping_no_steps_still_preserves_the_current_task(self): messages = long_conversation() compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( From f321e691748fae33b4125cffe45d61d53cad8e26 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:52:28 +0200 Subject: [PATCH 07/38] update reno and add more dev comments --- haystack/hooks/compaction/sliding_window.py | 15 ++++++++------- ...agent-context-compaction-3258c08dec9d2b34.yaml | 8 +++++--- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index f18d1fe07a4..d3961abd32f 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -103,21 +103,22 @@ def _first_group_to_keep( messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter ) -> int: """ - Return the oldest group that the budget can still pay for, working backwards from the newest. + Return the position in `groups` to start keeping from. - Groups are added up from newest to oldest and counting stops at the first group that does not fit, so what is kept - is always a run of groups at the end of the list. An older group is never kept once a newer one has been dropped, - which would leave a hole in the conversation. + Groups are added up from newest to oldest and counting stops at the first group that does not fit. :param messages: The full conversation containing the messages referenced by `groups`. - :param groups: Ordered index groups to choose from, oldest group first. + :param groups: Ordered index groups to choose from, the oldest group first is at position 0. :param available_tokens: The token budget available for these groups. :param token_counter: The `TokenCounter` used to measure each group. :returns: The position in `groups` to start keeping from: `len(groups)` when nothing fits, `0` when it all fits. """ + # We count backwards so first_kept starts such that nothing would be kept first_kept = len(groups) while first_kept > 0: + # Calculate the tokens consumed by the group that would be kept next group_tokens = token_counter.count(messages=_messages_at(messages=messages, indices=groups[first_kept - 1])) + # If this group does not fit, we stop if group_tokens > available_tokens: break available_tokens -= group_tokens @@ -149,8 +150,8 @@ def _first_turn_and_step_to_keep( messages=_messages_at(messages=messages, indices=_flatten(groups=step_groups)) ) if current_task_tokens > available_tokens: - # The current task alone overruns the budget, so every historical turn is dropped and - # the current task's own oldest steps trimmed until what remains fits. + # The current task alone overruns the budget, so every historical turn is dropped and the current task's own + # oldest steps trimmed until what remains fits. first_kept_step = _first_group_to_keep( messages=messages, groups=step_groups, available_tokens=available_tokens, token_counter=token_counter ) diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index c71c803f853..1c7d53ec068 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -32,9 +32,11 @@ features: ``SlidingWindowCompactor`` treats an assistant message and its following tool results as one step, so a tool call is never separated from its results. Historical turns are also kept whole, so an assistant reply is not retained - without the user message it answers. It may retain slightly more history than the requested target when preserving - the current task or the configured minimum number of steps. Compaction is lossy: removed messages cannot be - recovered or summarized by this strategy. Implement the ``Compactor`` protocol to provide a custom strategy. + without the user message it answers. It can also land above the requested target rather than under it, because + leading system messages and the current task are never removed and ``min_keep_steps`` holds on to the newest Agent + steps whatever their size, so a long system prompt or one large tool result can leave the conversation well over + the target. Compaction is lossy: removed messages cannot be recovered or summarized by this strategy. Implement + the ``Compactor`` protocol to provide a custom strategy. ``CompactionHook`` and ``SlidingWindowCompactor`` emit an ``ExperimentalWarning`` and may change without a deprecation cycle. From a8c878a9a437bb1d6c46e6a7e1374b9c7d75ab19 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 10:50:26 +0200 Subject: [PATCH 08/38] PR comments --- haystack/hooks/compaction/sliding_window.py | 42 ++++---- ...t-context-compaction-3258c08dec9d2b34.yaml | 5 +- test/hooks/compaction/helpers.py | 7 +- test/hooks/compaction/test_hooks.py | 12 +-- test/hooks/compaction/test_sliding_window.py | 95 +++++++++---------- 5 files changed, 80 insertions(+), 81 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index d3961abd32f..750bb24de06 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -128,7 +128,7 @@ def _first_group_to_keep( def _first_turn_and_step_to_keep( messages: list[ChatMessage], - turn_groups: list[list[int]], + historical_groups: list[list[int]], step_groups: list[list[int]], available_tokens: int, token_counter: TokenCounter, @@ -138,12 +138,12 @@ def _first_turn_and_step_to_keep( Return which historical turn and which Agent step of the current task to start keeping from. :param messages: The full conversation containing the messages referenced by both group lists. - :param turn_groups: Index groups for the complete historical turns preceding the current task, oldest first. + :param historical_groups: Index groups for the complete historical turns preceding the current task, oldest first. :param step_groups: Index groups for the current task's Agent steps, oldest first. :param available_tokens: The token budget left once the protected context is paid for. :param token_counter: The `TokenCounter` used to measure the groups. :param min_keep_steps: The fewest recent Agent steps to keep, even when they exceed the budget. - :returns: The position in `turn_groups` and the position in `step_groups` to start keeping from. Either is the + :returns: The position in `historical_groups` and the position in `step_groups` to start keeping from. Either is the length of its list when nothing from it is kept. """ current_task_tokens = token_counter.count( @@ -156,12 +156,12 @@ def _first_turn_and_step_to_keep( messages=messages, groups=step_groups, available_tokens=available_tokens, token_counter=token_counter ) # The newest steps are kept regardless of the budget. - return len(turn_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) + return len(historical_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) # The whole current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. first_kept_turn = _first_group_to_keep( messages=messages, - groups=turn_groups, + groups=historical_groups, available_tokens=available_tokens - current_task_tokens, token_counter=token_counter, ) @@ -186,8 +186,9 @@ def _task_and_step_split( :returns: A tuple containing: 1. The messages to keep, ordered oldest to newest. - 2. The position in that list where an omission note belongs, immediately before the current task when only - historical turns were removed and immediately before the surviving steps when the task itself was trimmed. + 2. The position in that list where an omission note belongs, which is where the removed messages used to sit: + directly after the leading system messages when only historical turns were removed, and directly after the + user message anchoring the current task when the task's own steps were removed. 3. Every message selected for removal. """ # Find the leading system messages that contain the Agent instructions. @@ -204,7 +205,7 @@ def _task_and_step_split( # Group the complete historical turns (i.e. user-assistant) that precede the current task. An earlier compaction's # note is left out of its turn, so keeping the turn folds that note into the note this compaction leaves behind. historical_end = task_index if task_index is not None else system_end - turn_groups = _index_groups( + historical_groups = _index_groups( messages=messages, spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), skip_compaction_notes=True, @@ -214,13 +215,13 @@ def _task_and_step_split( protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) first_kept_turn, first_kept_step = _first_turn_and_step_to_keep( messages=messages, - turn_groups=turn_groups, + historical_groups=historical_groups, step_groups=step_groups, available_tokens=target_tokens - protected_tokens, token_counter=token_counter, min_keep_steps=min_keep_steps, ) - kept_turn_indices = _flatten(groups=turn_groups[first_kept_turn:]) + kept_turn_indices = _flatten(groups=historical_groups[first_kept_turn:]) kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) kept_turns = _messages_at(messages=messages, indices=kept_turn_indices) kept_steps = _messages_at(messages=messages, indices=kept_step_indices) @@ -232,9 +233,10 @@ def _task_and_step_split( kept_indices.add(task_index) removable = [message for index, message in enumerate(messages) if index not in kept_indices] - # The note stands in for the newest thing that was dropped. That is the current task's own steps when those were - # cut, in which case every historical turn went too and `kept_turns` is empty; otherwise it is the older turns. - note_index = system_end + len(kept_turns) + (len(task) if first_kept_step > 0 else 0) + # The note stands in for what was dropped, so it goes where the dropped messages used to sit. Either right after + # the leading system messages when the historical turns were trimmed, or right after the user message that anchors + # the current task when its own Agent steps were trimmed. Both positions are counted off the layout of `kept`. + note_index = system_end if first_kept_step == 0 else system_end + len(kept_turns) + len(task) return kept, note_index, removable @@ -243,9 +245,13 @@ class SlidingWindowCompactor(Compactor): """ Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows. - Leading system messages and the latest user message are protected. Earlier user/assistant turns are retained when - they fit, and the current task's history is retained in complete Agent steps, where a step is an assistant message - together with all immediately following tool results. An `omission_note` is left in place of what was removed. + Leading system messages and the latest user message are protected. Earlier user/assistant turns are kept when they + fit, and the current task's history is kept in complete Agent steps, where a step is an assistant message together + with all immediately following tool results. + + An `omission_note` is left where the removed messages used to sit: directly after the leading system messages when + only earlier turns were removed, and directly after the latest user message when the current task's own steps were + removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement. ```python from haystack.components.agents import Agent @@ -288,8 +294,8 @@ def compact( :param messages: The conversation to compact, oldest to newest. :param target_tokens: The size the kept conversation should come in under. :param token_counter: The `TokenCounter` to measure messages with. - :returns: The protected context, an omission note if configured, and the steps that survived; or None when - there is nothing to remove. + :returns: The conversation that survived, with an omission note if configured standing where the removed + messages used to sit; or None when there is nothing to remove. """ if token_counter.count(messages=messages) <= target_tokens: return None diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index 1c7d53ec068..6caa6298451 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -7,7 +7,10 @@ features: The first built-in strategy, ``SlidingWindowCompactor``, preserves leading system messages, the latest user task, and as much complete recent conversation as the target allows. It removes earlier user/assistant turns as whole units first. Only when removing all earlier turns is insufficient does it remove individual Agent steps from the - current task. It replaces removed history with a short omission note. + current task. It replaces removed history with a short omission note, left where the removed messages used to sit: + directly after the leading system messages when only earlier turns were removed, and directly after the latest user + message when the current task's own steps were removed. Only one note is ever present, because a later compaction + folds an earlier one into its replacement. .. code-block:: python diff --git a/test/hooks/compaction/helpers.py b/test/hooks/compaction/helpers.py index 637e6005d98..6966aade7f5 100644 --- a/test/hooks/compaction/helpers.py +++ b/test/hooks/compaction/helpers.py @@ -67,12 +67,11 @@ def make_state(messages: list[ChatMessage], **data: Any) -> State: return State(schema=_SCHEMA, data={**base, **data}) -def long_conversation() -> list[ChatMessage]: +def fresh_conversation_with_two_steps() -> list[ChatMessage]: """ - Six messages: a system prefix, a user turn, then two tool round-trips. + A system prefix and a first user task with two Agent steps behind it, so there are no earlier turns to remove. - The results are padded so that removing them saves more than an omission note costs, which is what a compactor - weighs before leaving one behind. + The tool results are padded so that dropping a step is a saving worth making. """ return [ ChatMessage.from_system("rules"), diff --git a/test/hooks/compaction/test_hooks.py b/test/hooks/compaction/test_hooks.py index 4a46630bc9a..1dbeb56be36 100644 --- a/test/hooks/compaction/test_hooks.py +++ b/test/hooks/compaction/test_hooks.py @@ -18,7 +18,7 @@ from test.hooks.compaction.helpers import ( FakeCounter, count_markers, - long_conversation, + fresh_conversation_with_two_steps, make_state, tool_call, tool_result, @@ -165,7 +165,7 @@ class TestCompactionHook: ) def test_trigger(self, context_tokens, should_compact): compactor = _RecordingCompactor() - _hook(compactor).run(make_state(messages=long_conversation(), context_tokens=context_tokens)) + _hook(compactor).run(make_state(messages=fresh_conversation_with_two_steps(), context_tokens=context_tokens)) assert compactor.calls == (["compact"] if should_compact else []) def test_fires_without_reported_usage_by_counting_locally(self): @@ -201,7 +201,7 @@ def test_subtracts_provider_overhead_from_the_target(self): # The reported count also covers tool schemas and template overhead, which a compactor cannot remove. Here that # overhead alone exceeds the target, so the compactor is told to cut the messages as far as it is allowed. compactor = _RecordingCompactor() - _hook(compactor).run(make_state(messages=long_conversation(), context_tokens=800)) + _hook(compactor).run(make_state(messages=fresh_conversation_with_two_steps(), context_tokens=800)) assert compactor.targets[0] == 0 def test_warns_when_the_token_counter_exceeds_the_context_estimate(self, caplog): @@ -213,7 +213,7 @@ def test_warns_when_the_token_counter_exceeds_the_context_estimate(self, caplog) def test_rewrites_messages_and_re_estimates_context_tokens(self): counter = FakeCounter() hook = _hook(token_counter=counter) - messages = long_conversation() + messages = fresh_conversation_with_two_steps() original_context_tokens = 800 estimated = _estimated_context_tokens( messages=messages, context_tokens=original_context_tokens, token_counter=counter @@ -244,7 +244,7 @@ def test_preserves_the_no_usage_sentinel_after_compaction(self): assert state.data["context_tokens"] == 0 def test_leaves_the_conversation_alone_when_the_compactor_declines(self): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() state = make_state(messages, context_tokens=800) _hook(_RecordingCompactor(result=None)).run(state=state) assert state.data["messages"] == messages @@ -312,7 +312,7 @@ class TestCompactionHookAsync: @pytest.mark.asyncio async def test_run_async_uses_the_async_compaction_path(self): compactor = _RecordingCompactor() - await _hook(compactor).run_async(make_state(long_conversation(), context_tokens=800)) + await _hook(compactor).run_async(make_state(fresh_conversation_with_two_steps(), context_tokens=800)) assert compactor.calls == ["compact_async"] @pytest.mark.asyncio diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index bafacf91205..9745b7b1aa9 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -8,7 +8,13 @@ from haystack.hooks.compaction import SlidingWindowCompactor from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans from haystack.hooks.compaction.utils import _COMPACTION_META_KEY -from test.hooks.compaction.helpers import FakeCounter, count_markers, long_conversation, tool_call, tool_result +from test.hooks.compaction.helpers import ( + FakeCounter, + count_markers, + fresh_conversation_with_two_steps, + tool_call, + tool_result, +) pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") @@ -59,67 +65,43 @@ def test_compaction_note_does_not_start_a_new_turn(self): class TestSlidingWindowCompactor: - def test_replaces_the_middle_with_an_omission_note(self): - messages = long_conversation() + def test_replaces_oldest_agent_step(self): + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) assert compacted is not None - # The instructions and task survive, the note stands in for the older step, and the latest step stays intact. - assert compacted[:2] == messages[:2] - assert compacted[3:] == messages[4:] + # Only the oldest agent step is removed which is now where the omission note is. + assert compacted == [*messages[:2], compacted[2], *messages[4:]] + + # Test omission note assert compacted[2].meta[_COMPACTION_META_KEY] == { "strategy": "sliding_window", "removed_messages": 2, "kept_messages": 4, } assert compacted[2].text == _DEFAULT_OMISSION_NOTE.replace("{num_removed}", "2") - # A user message, not a system one, so providers that hoist system messages cannot move it out of position. assert compacted[2].is_from(role=ChatRole.USER) - def test_places_omission_note_before_current_task_when_only_historical_turns_are_removed(self): + def test_replaces_oldest_historical_turn(self): messages = [ ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="old question"), ChatMessage.from_assistant(text="old answer " * 100), - ChatMessage.from_user(text="current task"), - ChatMessage.from_assistant(text="current answer"), - ] - compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=20, token_counter=COUNTER) - assert compacted is not None - assert compacted[0] == messages[0] - assert _COMPACTION_META_KEY in compacted[1].meta - assert compacted[2:] == messages[3:] - - def test_a_roomier_target_keeps_more(self): - messages = [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="task")] - for index in range(4): - messages.extend([tool_call(f"c{index}"), tool_result(result="x" * 400, call_id=f"c{index}")]) - compactor = SlidingWindowCompactor(omission_note=None) - tight = compactor.compact(messages=messages, target_tokens=60, token_counter=COUNTER) - roomy = compactor.compact(messages=messages, target_tokens=350, token_counter=COUNTER) - assert tight is not None - assert roomy is not None - assert len(tight) == 4 - assert len(roomy) == 8 - - def test_keeps_complete_recent_user_assistant_turns_that_fit(self): - messages = [ - ChatMessage.from_system(text="rules"), - ChatMessage.from_user(text="old question"), - ChatMessage.from_assistant(text="old answer"), ChatMessage.from_user(text="recent question"), ChatMessage.from_assistant(text="recent answer"), ChatMessage.from_user(text="current task"), - ChatMessage.from_assistant(text="current step"), + ChatMessage.from_assistant(text="current answer"), ] - # Drops one historical turn - expected = [messages[0], *messages[3:]] - target_tokens = COUNTER.count(expected) - compacted = SlidingWindowCompactor(omission_note=None).compact( + # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn. + target_tokens = 30 + compacted = SlidingWindowCompactor().compact( messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) - assert compacted == expected + assert compacted is not None + # The oldest historical turn is removed and replaced with an omission note. + assert compacted == [messages[0], compacted[1], *messages[3:]] + assert _COMPACTION_META_KEY in compacted[1].meta - def test_drops_historical_context_and_one_current_task_step_to_reach_target(self): + def test_drops_all_historical_turns_and_oldest_agent_step(self): system_message = ChatMessage.from_system(text="rules") historical_turn = [ ChatMessage.from_user(text="old question"), @@ -147,33 +129,40 @@ def test_folds_an_earlier_note_inside_a_retained_turn_and_counts_it_as_removed(s "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} ) messages = [ + # System ChatMessage.from_system(text="rules"), + # Historical ChatMessage.from_user(text="old question " * 200), ChatMessage.from_assistant(text="old answer"), ChatMessage.from_user(text="recent question"), ChatMessage.from_assistant(text="recent answer"), note, + # Current Task ChatMessage.from_user(text="current task"), ChatMessage.from_assistant(text="current step"), ] - # Room for everything but the oldest turn, so the turn holding the earlier note is retained around it. + # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn, so + # the turn holding the earlier note survives around it. + target_tokens = 30 compacted = SlidingWindowCompactor().compact( - messages=messages, target_tokens=COUNTER.count([messages[0], *messages[3:]]), token_counter=COUNTER + messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) assert compacted is not None + assert compacted == [messages[0], compacted[1], *messages[3:5], *messages[6:]] # The earlier note is replaced by the new one rather than surviving alongside it, and it counts as removed. assert count_markers(messages=compacted) == 1 - assert compacted == [messages[0], *messages[3:5], compacted[3], *messages[6:]] - assert compacted[3].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 + assert compacted[1].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 def test_returns_none_when_the_conversation_already_fits(self): assert ( - SlidingWindowCompactor().compact(messages=long_conversation(), target_tokens=100_000, token_counter=COUNTER) + SlidingWindowCompactor().compact( + messages=fresh_conversation_with_two_steps(), target_tokens=100_000, token_counter=COUNTER + ) is None ) def test_omission_note_can_be_turned_off(self): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(omission_note=None).compact( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) @@ -192,7 +181,7 @@ def test_omission_note_can_be_turned_off(self): ) def test_omission_note_can_be_customized(self, note, expected): compacted = SlidingWindowCompactor(omission_note=note).compact( - messages=long_conversation(), target_tokens=SMALLEST, token_counter=COUNTER + messages=fresh_conversation_with_two_steps(), target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None assert compacted[2].text == expected @@ -232,7 +221,7 @@ def test_keeps_a_parallel_tool_call_together_with_all_results(self): @pytest.mark.parametrize(("min_keep_steps", "expected"), [(0, 0), (1, 1), (2, 2), (20, 2)]) def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, expected): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(min_keep_steps=min_keep_steps, omission_note=None).compact( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) @@ -257,7 +246,9 @@ def test_preserves_the_latest_user_task_when_compacting_input_history(self): def test_repeated_compaction_folds_the_previous_note(self): compactor = SlidingWindowCompactor() - first = compactor.compact(messages=long_conversation(), target_tokens=SMALLEST, token_counter=COUNTER) + first = compactor.compact( + messages=fresh_conversation_with_two_steps(), target_tokens=SMALLEST, token_counter=COUNTER + ) assert first is not None # Simulate two more turns arriving on top of the already-compacted conversation. grown = [*first, tool_call("c3"), tool_result(result="third result", call_id="c3")] @@ -269,7 +260,7 @@ def test_repeated_compaction_folds_the_previous_note(self): assert second[1].text == "start" def test_keeping_no_steps_still_preserves_the_current_task(self): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) @@ -292,7 +283,7 @@ class TestSlidingWindowCompactorAsync: async def test_compact_async_matches_compact(self): # `SlidingWindowCompactor` does no I/O, so it relies on the protocol's default `compact_async`. compactor = SlidingWindowCompactor() - messages = long_conversation() + messages = fresh_conversation_with_two_steps() assert await compactor.compact_async( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) == compactor.compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) From d83cc0aff79bfecfdc19662411455ba905ce1d10 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 11:35:47 +0200 Subject: [PATCH 09/38] refactoring tests and improve previous compaction note detection --- haystack/hooks/compaction/sliding_window.py | 23 ++- test/hooks/compaction/test_sliding_window.py | 198 +++++++++++++------ 2 files changed, 159 insertions(+), 62 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 750bb24de06..37227923091 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -11,6 +11,9 @@ from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental +# Recorded as the strategy on every message this compactor produces, so a later run can recognize its own notes. +_STRATEGY = "sliding_window" + # Placeholder a custom omission note may include to have the number of removed messages substituted in. _NUM_REMOVED_PLACEHOLDER = "{num_removed}" @@ -44,6 +47,18 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: return None +def _is_compaction_note(message: ChatMessage) -> bool: + """ + Whether a message is an omission note this strategy left in place of removed history. + + Every compactor marks what it produces with the same meta key, including the tool results + `ToolResultPruningCompactor` rewrites into a placeholder. Matching on the role and the strategy keeps those out: + they are still part of the conversation and have to travel with the turn they belong to. + """ + marker = message.meta.get(_COMPACTION_META_KEY) + return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY + + def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: """ Return spans for complete user turns in a bounded section of conversation history. @@ -80,11 +95,7 @@ def _index_groups( Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. """ return [ - [ - index - for index in range(start, end) - if not skip_compaction_notes or _COMPACTION_META_KEY not in messages[index].meta - ] + [index for index in range(start, end) if not (skip_compaction_notes and _is_compaction_note(messages[index]))] for start, end in spans ] @@ -316,7 +327,7 @@ def compact( self.omission_note.replace(_NUM_REMOVED_PLACEHOLDER, str(len(removable))), meta={ _COMPACTION_META_KEY: { - "strategy": "sliding_window", + "strategy": _STRATEGY, "removed_messages": len(removable), "kept_messages": len(kept), } diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 9745b7b1aa9..332332b71b6 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -4,9 +4,9 @@ import pytest -from haystack.dataclasses import ChatMessage, ChatRole +from haystack.dataclasses import ChatMessage, ChatRole, ToolCall from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans, _is_compaction_note from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import ( FakeCounter, @@ -24,6 +24,51 @@ COUNTER = FakeCounter() +class TestIsCompactionNote: + @pytest.mark.parametrize( + ("message", "expected"), + [ + pytest.param( + ChatMessage.from_user( + text="Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), + True, + id="note-this-strategy-left", + ), + # A pruned result carries the same meta key but is still part of the conversation, so it is not a note. + pytest.param( + ChatMessage.from_tool( + tool_result="[Tool result removed to free up context.]", + origin=ToolCall(tool_name="search", arguments={}, id="c1"), + meta={_COMPACTION_META_KEY: {"strategy": "tool_result_pruning", "original_tokens": 180}}, + ), + False, + id="tool-result-another-strategy-pruned", + ), + # Another strategy's note is not this one's to fold away, so it is left where it is. + pytest.param( + ChatMessage.from_user( + text="A summary of what came before.", meta={_COMPACTION_META_KEY: {"strategy": "summarization"}} + ), + False, + id="note-another-strategy-left", + ), + pytest.param( + ChatMessage.from_system(text="rules", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}}), + False, + id="system-message", + ), + pytest.param( + ChatMessage.from_user(text="odd", meta={_COMPACTION_META_KEY: "sliding_window"}), + False, + id="marker-that-is-not-a-mapping", + ), + ], + ) + def test_only_matches_sliding_window_omission_message(self, message, expected): + assert _is_compaction_note(message=message) is expected + + class TestHistoricalTurnSpans: def test_groups_each_user_message_with_its_assistant_steps_and_tool_results(self): messages = [ @@ -51,20 +96,36 @@ def test_only_returns_turns_within_the_requested_bounds(self): assert _historical_turn_spans(messages=messages, start=2, end=4) == [(2, 4)] def test_compaction_note_does_not_start_a_new_turn(self): - note = ChatMessage.from_user( - "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} - ) messages = [ + # Historical turns + ChatMessage.from_user( + "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), ChatMessage.from_user("task"), ChatMessage.from_assistant("first step"), - note, - ChatMessage.from_assistant("second step"), ChatMessage.from_user("next task"), + ChatMessage.from_assistant("second step"), ] - assert _historical_turn_spans(messages=messages, start=0, end=len(messages)) == [(0, 4), (4, 5)] + # The note is skipped which is why the first span starts at 1 + assert _historical_turn_spans(messages=messages, start=0, end=len(messages)) == [(1, 3), (3, 5)] class TestSlidingWindowCompactor: + def test_replaces_all_historical_turns(self): + messages = [ + # System + ChatMessage.from_system(text="rules"), + # Historical turn + ChatMessage.from_user(text="old question " * 100), + ChatMessage.from_assistant(text="old answer"), + # Current task + ChatMessage.from_user(text="current task"), + ] + compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + assert compacted == [messages[0], messages[-1]] + def test_replaces_oldest_agent_step(self): messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) @@ -124,42 +185,87 @@ def test_drops_all_historical_turns_and_oldest_agent_step(self): ) assert compacted == [system_message, current_task[0], *current_task[-2:]] - def test_folds_an_earlier_note_inside_a_retained_turn_and_counts_it_as_removed(self): - note = ChatMessage.from_user( - "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} - ) + def test_replaces_earlier_note_and_oldest_historical_turn(self): messages = [ # System ChatMessage.from_system(text="rules"), + # The note an earlier compaction left, which sits at the top of the historical turns. + ChatMessage.from_user( + text="Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), # Historical ChatMessage.from_user(text="old question " * 200), ChatMessage.from_assistant(text="old answer"), ChatMessage.from_user(text="recent question"), ChatMessage.from_assistant(text="recent answer"), - note, # Current Task ChatMessage.from_user(text="current task"), ChatMessage.from_assistant(text="current step"), ] - # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn, so - # the turn holding the earlier note survives around it. + # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn. target_tokens = 30 compacted = SlidingWindowCompactor().compact( messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) assert compacted is not None - assert compacted == [messages[0], compacted[1], *messages[3:5], *messages[6:]] - # The earlier note is replaced by the new one rather than surviving alongside it, and it counts as removed. + # The earlier note goes with the turn it stood in front of, and the new note takes its place. + assert compacted == [messages[0], compacted[1], *messages[4:]] assert count_markers(messages=compacted) == 1 + # The earlier note is counted among the removed, alongside the two messages of the oldest turn. assert compacted[1].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 - def test_returns_none_when_the_conversation_already_fits(self): - assert ( - SlidingWindowCompactor().compact( - messages=fresh_conversation_with_two_steps(), target_tokens=100_000, token_counter=COUNTER - ) - is None + def test_replaces_earlier_note_and_oldest_agent_step(self): + messages = [ + # System + ChatMessage.from_system(text="rules"), + # Current Task + ChatMessage.from_user(text="current task"), + # The note an earlier compaction left, which sits right after the task when its own steps were trimmed. + ChatMessage.from_user( + text="Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), + tool_call("c1"), + tool_result(result="first result", call_id="c1"), + tool_call("c2"), + tool_result(result="second result", call_id="c2"), + ] + compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) + assert compacted is not None + # The earlier note goes with the step it stood in front of, and the new note takes its place. + assert compacted == [*messages[:2], compacted[2], *messages[5:]] + assert count_markers(messages=compacted) == 1 + # The earlier note is counted among the removed, alongside the two messages of the oldest step. + assert compacted[2].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 + + def test_keeps_a_pruned_tool_result_inside_a_kept_turn(self): + messages = [ + # System + ChatMessage.from_system(text="rules"), + # Historical, dropped to make room + ChatMessage.from_user(text="ancient question " * 200), + ChatMessage.from_assistant(text="ancient answer"), + # Historical, kept + ChatMessage.from_user(text="old question"), + tool_call("old"), + # A result `ToolResultPruningCompactor` already pruned, which carries the same meta key as an omission note + # but is part of the conversation rather than standing in for removed history. + ChatMessage.from_tool( + tool_result="[Tool result removed to free up context.]", + origin=ToolCall(tool_name="search", arguments={}, id="old"), + meta={_COMPACTION_META_KEY: {"strategy": "tool_result_pruning", "original_tokens": 180}}, + ), + # Current Task + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Enough for the instructions, the kept turn, and the task with its step, but not the padded oldest turn. + target_tokens = 45 + compacted = SlidingWindowCompactor(omission_note=None).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) + assert compacted is not None + # The pruned result is not an omission note, so it stays with the turn and its tool call keeps its answer. + assert compacted == [messages[0], *messages[3:]] def test_omission_note_can_be_turned_off(self): messages = fresh_conversation_with_two_steps() @@ -187,18 +293,25 @@ def test_omission_note_can_be_customized(self, note, expected): assert compacted[2].text == expected @pytest.mark.parametrize( - "messages", + ("messages", "target_tokens"), [ - pytest.param([], id="empty"), - pytest.param([ChatMessage.from_system(text="a"), ChatMessage.from_system(text="b")], id="only-system"), + # The conversation is already under the target, so there is nothing to do. + pytest.param(fresh_conversation_with_two_steps(), 100_000, id="conversation-already-fits"), + # Over the target, but everything that is left is protected, so there is nothing the compactor may remove. pytest.param( - [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="hi")], id="nothing-outside-window" + [ChatMessage.from_system(text="a"), ChatMessage.from_system(text="b")], SMALLEST, id="only-system" + ), + pytest.param( + [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="hi")], + SMALLEST, + id="only-system-and-task", ), ], ) - def test_returns_none_when_there_is_nothing_to_remove(self, messages): + def test_returns_none_when_there_is_nothing_to_remove(self, messages, target_tokens): assert ( - SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) is None + SlidingWindowCompactor().compact(messages=messages, target_tokens=target_tokens, token_counter=COUNTER) + is None ) def test_keeps_a_parallel_tool_call_together_with_all_results(self): @@ -232,33 +345,6 @@ def test_rejects_negative_min_keep_steps(self): with pytest.raises(ValueError, match="`min_keep_steps` must be at least 0"): SlidingWindowCompactor(min_keep_steps=-1) - def test_preserves_the_latest_user_task_when_compacting_input_history(self): - messages = [ - ChatMessage.from_system(text="rules"), - ChatMessage.from_user(text="old question " * 100), - ChatMessage.from_assistant(text="old answer"), - ChatMessage.from_user(text="current task"), - ] - compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( - messages=messages, target_tokens=SMALLEST, token_counter=COUNTER - ) - assert compacted == [messages[0], messages[-1]] - - def test_repeated_compaction_folds_the_previous_note(self): - compactor = SlidingWindowCompactor() - first = compactor.compact( - messages=fresh_conversation_with_two_steps(), target_tokens=SMALLEST, token_counter=COUNTER - ) - assert first is not None - # Simulate two more turns arriving on top of the already-compacted conversation. - grown = [*first, tool_call("c3"), tool_result(result="third result", call_id="c3")] - second = compactor.compact(messages=grown, target_tokens=SMALLEST, token_counter=COUNTER) - assert second is not None - # The original task stays anchored and the earlier compaction note is folded into the new one. - assert count_markers(messages=second) == 1 - assert second[0].text == "rules" - assert second[1].text == "start" - def test_keeping_no_steps_still_preserves_the_current_task(self): messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( From 59271e07941ff9e98b5ce0205ced7f17bdb8c796 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 11:45:55 +0200 Subject: [PATCH 10/38] PR comments --- haystack/hooks/compaction/sliding_window.py | 91 +++++++++++++-------- 1 file changed, 56 insertions(+), 35 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 37227923091..f1d0cfbbd7f 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -101,15 +101,51 @@ def _index_groups( def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages at the given indices, in conversation order.""" + """Return the messages at the given indices, in the order the indices are given.""" return [messages[index] for index in indices] +def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages the given indices leave out, in conversation order.""" + left_out = set(indices) + return [message for index, message in enumerate(messages) if index not in left_out] + + def _flatten(groups: list[list[int]]) -> list[int]: """Join index groups into a single ordered list of indices.""" return [index for group in groups for index in group] +def _removable_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> tuple[list[list[int]], list[list[int]]]: + """ + Group the two stretches of conversation compaction is allowed to remove. + + :param messages: The full conversation, ordered oldest to newest. + :param system_end: The end of the leading system-message block. + :param task_index: The index of the user message anchoring the current task, or None when there is none. + :returns: Index groups for the complete historical turns preceding the current task, and index groups for the + current task's own Agent steps. Both are ordered oldest group first. A group is the unit of removal: it is + kept or removed entire, which is what keeps a tool call with its results and an assistant reply with the user + message it answers. + """ + # Steps belong to the current task, so they start after its anchor, or after the instructions when the + # conversation has no user message to anchor on. + step_start = (task_index + 1) if task_index is not None else system_end + step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start)) + + # An earlier compaction's note is left out of its turn, so keeping the turn folds that note into the note this + # compaction leaves behind. + historical_end = task_index if task_index is not None else system_end + historical_groups = _index_groups( + messages=messages, + spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), + skip_compaction_notes=True, + ) + return historical_groups, step_groups + + def _first_group_to_keep( messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter ) -> int: @@ -202,53 +238,38 @@ def _task_and_step_split( user message anchoring the current task when the task's own steps were removed. 3. Every message selected for removal. """ - # Find the leading system messages that contain the Agent instructions. + # The landmarks the split is built around: the Agent instructions, and the user message anchoring the current task. system_end = _leading_system_end(messages=messages) - - # Find the latest user message to use as the current task anchor. task_index = _latest_user_index(messages=messages) - task = [messages[task_index]] if task_index is not None else [] + task_indices = [task_index] if task_index is not None else [] - # Group the complete Agent steps that follow the current task anchor. - step_start_index = (task_index + 1) if task_index is not None else system_end - step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start_index)) + # The two stretches compaction may remove. A group is the unit of removal, so a turn or a step is never split. + historical_groups, step_groups = _removable_groups(messages=messages, system_end=system_end, task_index=task_index) - # Group the complete historical turns (i.e. user-assistant) that precede the current task. An earlier compaction's - # note is left out of its turn, so keeping the turn folds that note into the note this compaction leaves behind. - historical_end = task_index if task_index is not None else system_end - historical_groups = _index_groups( - messages=messages, - spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), - skip_compaction_notes=True, - ) - - # The Agent instructions and the current task are never removed, so they are paid for out of the target first. - protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + # The instructions and the current task are never removed, so they come off the budget first. + protected = _messages_at(messages=messages, indices=[*range(system_end), *task_indices]) first_kept_turn, first_kept_step = _first_turn_and_step_to_keep( messages=messages, historical_groups=historical_groups, step_groups=step_groups, - available_tokens=target_tokens - protected_tokens, + available_tokens=target_tokens - token_counter.count(messages=protected), token_counter=token_counter, min_keep_steps=min_keep_steps, ) + + # What survives, laid out in conversation order: the instructions, the turns that fit, the task, then its steps. kept_turn_indices = _flatten(groups=historical_groups[first_kept_turn:]) kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) - kept_turns = _messages_at(messages=messages, indices=kept_turn_indices) - kept_steps = _messages_at(messages=messages, indices=kept_step_indices) - kept = [*messages[:system_end], *kept_turns, *task, *kept_steps] - - # A message survives only by being protected or by falling in a group we are keeping; everything else goes. - kept_indices = {*range(system_end), *kept_turn_indices, *kept_step_indices} - if task_index is not None: - kept_indices.add(task_index) - removable = [message for index, message in enumerate(messages) if index not in kept_indices] - - # The note stands in for what was dropped, so it goes where the dropped messages used to sit. Either right after - # the leading system messages when the historical turns were trimmed, or right after the user message that anchors - # the current task when its own Agent steps were trimmed. Both positions are counted off the layout of `kept`. - note_index = system_end if first_kept_step == 0 else system_end + len(kept_turns) + len(task) - return kept, note_index, removable + kept_indices = [*range(system_end), *kept_turn_indices, *task_indices, *kept_step_indices] + + # The note stands in for what was dropped, so it goes where those messages used to sit. Either right after the + # instructions when the historical turns were trimmed, or right after the task anchor when its own steps were. + note_index = system_end if first_kept_step == 0 else system_end + len(kept_turn_indices) + len(task_indices) + return ( + _messages_at(messages=messages, indices=kept_indices), + note_index, + _messages_except(messages=messages, indices=kept_indices), + ) @_experimental From 6a8c8037984e285b85366081640f2c928041f843 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 11:52:07 +0200 Subject: [PATCH 11/38] update docs pages --- .../agents-1/compaction.mdx | 4 +-- .../agents-1/compaction/compaction-hook.mdx | 2 +- .../compaction/sliding-window-compactor.mdx | 27 +++++++++---------- 3 files changed, 16 insertions(+), 17 deletions(-) diff --git a/docs-website/docs/pipeline-components/agents-1/compaction.mdx b/docs-website/docs/pipeline-components/agents-1/compaction.mdx index b8562acc820..792a8a231d8 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction.mdx @@ -72,7 +72,7 @@ Compactors receive the current messages, a target token count, and the same toke | Compactor | Strategy | Trade-off | | --- | --- | --- | -| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions, latest user task, and recent complete Agent steps while removing older history. | Fast and local, but discarded information is not summarized. | +| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions and latest user task, keeps earlier user/assistant turns whole while they fit, and trims the current task's own Agent steps only when that is not enough. | Fast and local, but discarded information is not summarized. | | [`ToolResultPruningCompactor`](compaction/tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while preserving tool-call/result structure. | Retains the shape of the run and recent results, but removes the content of pruned results. | ## Combining compaction strategies @@ -108,7 +108,7 @@ agent = Agent( ) ``` -If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes older complete Agent steps. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. +If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes earlier turns and then, if needed, the current task's oldest Agent steps. A result the pruning compactor already replaced with a placeholder stays with the turn it belongs to, so its tool call keeps an answer. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. ### Creating a custom compactor diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx index 0b8a27a66a4..db5394c1d74 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx @@ -90,7 +90,7 @@ The compactor controls what information survives: | Compactor | Strategy | | --- | --- | -| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and the most recent complete Agent steps, dropping older history. | +| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and as much complete recent conversation as fits, dropping earlier turns whole before it trims the task's own steps. | | [`ToolResultPruningCompactor`](tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while keeping recent results intact. | You can also implement the `Compactor` protocol for a custom strategy. See [Context Compaction](../compaction.mdx#creating-a-custom-compactor) for its requirements. diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx index 7de7773be7e..d27a358a733 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx @@ -2,12 +2,12 @@ title: "SlidingWindowCompactor" id: sliding-window-compactor slug: "/sliding-window-compactor" -description: "Use SlidingWindowCompactor to remove older Agent history while preserving the current task and recent complete steps." +description: "Use SlidingWindowCompactor to remove older Agent history while preserving the current task and as much recent complete conversation as fits." --- # SlidingWindowCompactor -`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as many recent complete Agent steps as the token target allows. +`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as much complete recent conversation as the token target allows. It removes earlier user/assistant turns as whole units first, and only trims the current task's own Agent steps when removing every earlier turn is not enough. :::warning[Experimental] @@ -51,17 +51,19 @@ compaction_hook = CompactionHook( ## How the sliding window is selected -The compactor divides a conversation into protected context, removable history, and recent Agent steps: +The compactor divides a conversation into protected context, earlier turns, and the current task's Agent steps: 1. It preserves all leading system messages as the Agent's instructions. 2. It preserves the latest user message as the current task. -3. It groups each assistant message and all immediately following tool-result messages into one complete Agent step. -4. Working backwards from the latest step, it keeps as many complete steps as fit within the target. -5. It replaces the removable middle history with an omission note, unless the note is disabled. +3. It groups the history before that task into complete user turns, each running from one user message up to the next. +4. It groups each assistant message and all immediately following tool-result messages into one complete Agent step. +5. Working backwards from the newest, it keeps as many whole earlier turns as fit within the target. +6. Only when the current task alone still exceeds the target does it begin removing that task's own steps, oldest first. +7. It replaces what it removed with an omission note, unless the note is disabled. -Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. +Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. Keeping whole turns likewise ensures an assistant reply is never retained without the user message it answers. -The target is a goal rather than a guarantee. Protected messages and the configured minimum number of recent steps take precedence when they already exceed the available token budget. +The target is a goal rather than a guarantee, and the conversation can end up above it rather than below. Leading system messages and the current task are never removed, and `min_keep_steps` holds on to the newest Agent steps whatever their size, so a long system prompt or a single large tool result can leave the conversation well over the target. ## Configuration @@ -76,16 +78,13 @@ The target is a goal rather than a guarantee. Protected messages and the configu An omission note tells the model that earlier context is missing. Without one, the shortened conversation can appear complete and the model may repeat work or behave as though it still has the removed information. -The compactor only inserts the note when it costs fewer tokens than the messages it replaces. Otherwise, it returns `None` and leaves the conversation unchanged. Repeated compactions fold the previous note into the newly removed block, leaving a single current note. +The note is left where the removed messages used to sit: directly after the leading system messages when only earlier turns were removed, and directly after the latest user message when the current task's own steps were removed. Repeated compactions fold an earlier note into the new one, so the conversation carries at most one. -Compaction metadata is stored on the note, including the strategy name and the numbers of removed messages, retained messages, and retained steps. +Compaction metadata is stored on the note, including the strategy name and the numbers of removed and retained messages. ## When the conversation is unchanged The compactor returns `None` without changing the conversation when: - The conversation already fits within `target_tokens`. -- There is no removable history outside the protected messages and retained steps. -- The configured omission note would cost at least as many tokens as the history it replaces. - -Setting `omission_note=None` removes the final condition because no replacement note needs to fit. +- There is no removable history outside the protected messages and the history it retained. From 466729a56c49f5a6bdde1139c24c1256b2f98c96 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 12:54:35 +0200 Subject: [PATCH 12/38] fix docs --- .../agents-1/compaction.mdx | 4 ++-- .../agents-1/compaction/compaction-hook.mdx | 2 +- .../compaction/sliding-window-compactor.mdx | 12 +++++----- haystack/hooks/compaction/sliding_window.py | 12 +++++----- ...t-context-compaction-3258c08dec9d2b34.yaml | 24 +++++++++---------- 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/docs-website/docs/pipeline-components/agents-1/compaction.mdx b/docs-website/docs/pipeline-components/agents-1/compaction.mdx index 792a8a231d8..8a76a1185d6 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction.mdx @@ -72,7 +72,7 @@ Compactors receive the current messages, a target token count, and the same toke | Compactor | Strategy | Trade-off | | --- | --- | --- | -| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions and latest user task, keeps earlier user/assistant turns whole while they fit, and trims the current task's own Agent steps only when that is not enough. | Fast and local, but discarded information is not summarized. | +| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions and latest user task, keeps complete historical turns while they fit, and trims the current task's own Agent steps only when that is not enough. | Fast and local, but discarded information is not summarized. | | [`ToolResultPruningCompactor`](compaction/tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while preserving tool-call/result structure. | Retains the shape of the run and recent results, but removes the content of pruned results. | ## Combining compaction strategies @@ -108,7 +108,7 @@ agent = Agent( ) ``` -If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes earlier turns and then, if needed, the current task's oldest Agent steps. A result the pruning compactor already replaced with a placeholder stays with the turn it belongs to, so its tool call keeps an answer. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. +If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes historical turns and then, if needed, the current task's oldest Agent steps. A result the pruning compactor already replaced with a placeholder stays with the historical turn it belongs to, so its tool call keeps an answer. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. ### Creating a custom compactor diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx index db5394c1d74..8837f6d2b51 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx @@ -90,7 +90,7 @@ The compactor controls what information survives: | Compactor | Strategy | | --- | --- | -| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and as much complete recent conversation as fits, dropping earlier turns whole before it trims the task's own steps. | +| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and as much complete recent conversation as fits, removing complete historical turns before it trims the task's own steps. | | [`ToolResultPruningCompactor`](tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while keeping recent results intact. | You can also implement the `Compactor` protocol for a custom strategy. See [Context Compaction](../compaction.mdx#creating-a-custom-compactor) for its requirements. diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx index d27a358a733..e66e5c915fe 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx @@ -7,7 +7,7 @@ description: "Use SlidingWindowCompactor to remove older Agent history while pre # SlidingWindowCompactor -`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as much complete recent conversation as the token target allows. It removes earlier user/assistant turns as whole units first, and only trims the current task's own Agent steps when removing every earlier turn is not enough. +`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as much complete recent conversation as the token target allows. It removes complete historical turns first, and only trims the current task's own Agent steps when removing every historical turn is not enough. :::warning[Experimental] @@ -51,17 +51,17 @@ compaction_hook = CompactionHook( ## How the sliding window is selected -The compactor divides a conversation into protected context, earlier turns, and the current task's Agent steps: +The compactor divides a conversation into protected context, historical turns, and the current task's Agent steps: 1. It preserves all leading system messages as the Agent's instructions. 2. It preserves the latest user message as the current task. -3. It groups the history before that task into complete user turns, each running from one user message up to the next. +3. It groups the history before that task into complete historical turns, each running from one user message up to the next. 4. It groups each assistant message and all immediately following tool-result messages into one complete Agent step. -5. Working backwards from the newest, it keeps as many whole earlier turns as fit within the target. +5. Working backwards from the newest, it keeps as many complete historical turns as fit within the target. 6. Only when the current task alone still exceeds the target does it begin removing that task's own steps, oldest first. 7. It replaces what it removed with an omission note, unless the note is disabled. -Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. Keeping whole turns likewise ensures an assistant reply is never retained without the user message it answers. +Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. Historical turns are kept or removed in full for the same reason: an assistant reply is never retained without the user message it answers. The target is a goal rather than a guarantee, and the conversation can end up above it rather than below. Leading system messages and the current task are never removed, and `min_keep_steps` holds on to the newest Agent steps whatever their size, so a long system prompt or a single large tool result can leave the conversation well over the target. @@ -78,7 +78,7 @@ The target is a goal rather than a guarantee, and the conversation can end up ab An omission note tells the model that earlier context is missing. Without one, the shortened conversation can appear complete and the model may repeat work or behave as though it still has the removed information. -The note is left where the removed messages used to sit: directly after the leading system messages when only earlier turns were removed, and directly after the latest user message when the current task's own steps were removed. Repeated compactions fold an earlier note into the new one, so the conversation carries at most one. +The note is left where the removed messages used to sit: directly after the leading system messages when only historical turns were removed, and directly after the latest user message when the current task's own steps were removed. Repeated compactions fold an earlier note into the new one, so the conversation carries at most one. Compaction metadata is stored on the note, including the strategy name and the numbers of removed and retained messages. diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index f1d0cfbbd7f..bd780e83d10 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -205,7 +205,7 @@ def _first_turn_and_step_to_keep( # The newest steps are kept regardless of the budget. return len(historical_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) - # The whole current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. + # The entire current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. first_kept_turn = _first_group_to_keep( messages=messages, groups=historical_groups, @@ -221,7 +221,7 @@ def _task_and_step_split( """ Split a conversation into the messages to keep and the messages to remove. - Leading system messages and the latest real user message are always kept. Historical user turns are kept whole when + Leading system messages and the latest real user message are always kept. Historical turns are kept in full when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time while keeping each assistant message together with its tool results. @@ -277,13 +277,13 @@ class SlidingWindowCompactor(Compactor): """ Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows. - Leading system messages and the latest user message are protected. Earlier user/assistant turns are kept when they - fit, and the current task's history is kept in complete Agent steps, where a step is an assistant message together + Leading system messages and the latest user message are protected. Historical turns are kept in full when they fit, + and the current task's history is kept in complete Agent steps, where a step is an assistant message together with all immediately following tool results. An `omission_note` is left where the removed messages used to sit: directly after the leading system messages when - only earlier turns were removed, and directly after the latest user message when the current task's own steps were - removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement. + only historical turns were removed, and directly after the latest user message when the current task's own steps + were removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement. ```python from haystack.components.agents import Agent diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index 6caa6298451..4f6955996ca 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -5,12 +5,12 @@ features: shortens the conversation when it reaches a configured fraction of the model's context window. The first built-in strategy, ``SlidingWindowCompactor``, preserves leading system messages, the latest user task, - and as much complete recent conversation as the target allows. It removes earlier user/assistant turns as whole - units first. Only when removing all earlier turns is insufficient does it remove individual Agent steps from the - current task. It replaces removed history with a short omission note, left where the removed messages used to sit: - directly after the leading system messages when only earlier turns were removed, and directly after the latest user - message when the current task's own steps were removed. Only one note is ever present, because a later compaction - folds an earlier one into its replacement. + and as much complete recent conversation as the target allows. It removes complete historical turns first, and only + when removing every historical turn is insufficient does it remove individual Agent steps from the current task. It + replaces removed history with a short omission note, left where the removed messages used to sit: directly after + the leading system messages when only historical turns were removed, and directly after the latest user message + when the current task's own steps were removed. Only one note is ever present, because a later compaction folds an + earlier one into its replacement. .. code-block:: python @@ -34,12 +34,12 @@ features: tool schemas. Leave headroom above ``compact_at`` for the next reply and its tool results. ``SlidingWindowCompactor`` treats an assistant message and its following tool results as one step, so a tool call is - never separated from its results. Historical turns are also kept whole, so an assistant reply is not retained - without the user message it answers. It can also land above the requested target rather than under it, because - leading system messages and the current task are never removed and ``min_keep_steps`` holds on to the newest Agent - steps whatever their size, so a long system prompt or one large tool result can leave the conversation well over - the target. Compaction is lossy: removed messages cannot be recovered or summarized by this strategy. Implement - the ``Compactor`` protocol to provide a custom strategy. + never separated from its results. Historical turns are likewise kept or removed in full, so an assistant reply is + not retained without the user message it answers. It can also land above the requested target rather than under it, + because leading system messages and the current task are never removed and ``min_keep_steps`` holds on to the + newest Agent steps whatever their size, so a long system prompt or one large tool result can leave the conversation + well over the target. Compaction is lossy: removed messages cannot be recovered or summarized by this strategy. + Implement the ``Compactor`` protocol to provide a custom strategy. ``CompactionHook`` and ``SlidingWindowCompactor`` emit an ``ExperimentalWarning`` and may change without a deprecation cycle. From abac3069202aca0b84cb18a5a5e7f4d999a05a6e Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 12:59:39 +0200 Subject: [PATCH 13/38] changes --- haystack/hooks/compaction/sliding_window.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index bd780e83d10..c916e5958ef 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -48,13 +48,7 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: def _is_compaction_note(message: ChatMessage) -> bool: - """ - Whether a message is an omission note this strategy left in place of removed history. - - Every compactor marks what it produces with the same meta key, including the tool results - `ToolResultPruningCompactor` rewrites into a placeholder. Matching on the role and the strategy keeps those out: - they are still part of the conversation and have to travel with the turn they belong to. - """ + """Whether a message is an omission note this strategy left in place of removed history.""" marker = message.meta.get(_COMPACTION_META_KEY) return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY @@ -72,8 +66,8 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. """ - # Compaction notes use the user role for provider compatibility, but they do not begin a new conversation turn. - # Ignoring marked messages here also lets a subsequent compaction fold an old note into its replacement. + # Reject any user-role message an earlier compaction produced, whichever strategy made it: none of them are user + # requests, so none of them begin a turn. Leaving them out also lets this compaction fold an old note away. user_indices = [ index for index in range(start, end) From e7f505e88368e8757e02eee14d720aa964e1e9ef Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:22:31 +0200 Subject: [PATCH 14/38] First pass at adding summarization compactor --- haystack/components/generators/chat/utils.py | 46 ++ haystack/hooks/compaction/__init__.py | 2 + haystack/hooks/compaction/sliding_window.py | 100 +-- haystack/hooks/compaction/summarization.py | 588 ++++++++++++++++++ haystack/hooks/compaction/utils.py | 62 ++ pydoc/hooks_api.yml | 2 +- ...marization-compactor-91b6be6855f478df.yaml | 6 + test/components/generators/chat/test_utils.py | 46 ++ test/hooks/compaction/test_sliding_window.py | 4 +- test/hooks/compaction/test_summarization.py | 249 ++++++++ 10 files changed, 1017 insertions(+), 88 deletions(-) create mode 100644 haystack/components/generators/chat/utils.py create mode 100644 haystack/hooks/compaction/summarization.py create mode 100644 releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml create mode 100644 test/components/generators/chat/test_utils.py create mode 100644 test/hooks/compaction/test_summarization.py diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py new file mode 100644 index 00000000000..9a1a6d5b824 --- /dev/null +++ b/haystack/components/generators/chat/utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.components.generators.chat.types import ChatGenerator + +_CHAT_COMPLETIONS_GENERATORS = {"OpenAIChatGenerator", "AzureOpenAIChatGenerator"} +_RESPONSES_GENERATORS = {"OpenAIResponsesChatGenerator", "AzureOpenAIResponsesChatGenerator"} + + +def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | None: + """Return the output-token parameter used by a known built-in Chat Generator.""" + class_names = {cls.__name__ for cls in type(chat_generator).__mro__} + if class_names & _RESPONSES_GENERATORS: + return "max_output_tokens" + if class_names & _CHAT_COMPLETIONS_GENERATORS: + return "max_completion_tokens" + return None + + +def _resolve_output_token_limit(chat_generator: ChatGenerator, default_limit: int) -> tuple[int, dict[str, Any] | None]: + """ + Resolve an effective output-token limit and runtime kwargs for a Chat Generator. + + A recognized limit configured directly on a built-in generator wins and is not repeated at runtime. When the + built-in generator has no configured limit, the default is returned as its provider-specific runtime setting. + Unknown generators receive no runtime setting because the ChatGenerator protocol does not standardize the key. + + :param chat_generator: The generator whose output should be limited. + :param default_limit: The positive fallback output-token limit. + :returns: The effective limit and provider-specific runtime generation kwargs, or None for no runtime kwargs. + """ + limit_key = _generator_output_token_limit_key(chat_generator=chat_generator) + if limit_key is None: + return default_limit, None + + configured = getattr(chat_generator, "generation_kwargs", None) + if isinstance(configured, dict) and limit_key in configured: + value = configured[limit_key] + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value, None + # The generator owns this setting. Do not silently replace an invalid value; let it report the problem. + return default_limit, None + return default_limit, {limit_key: default_limit} diff --git a/haystack/hooks/compaction/__init__.py b/haystack/hooks/compaction/__init__.py index 878515a6d1f..fb0ae0a1630 100644 --- a/haystack/hooks/compaction/__init__.py +++ b/haystack/hooks/compaction/__init__.py @@ -10,6 +10,7 @@ _import_structure = { "hooks": ["CompactionHook"], "sliding_window": ["SlidingWindowCompactor"], + "summarization": ["SummarizationCompactor"], "tool_result_pruning": ["ToolResultPruningCompactor"], "types": ["Compactor"], } @@ -17,6 +18,7 @@ if TYPE_CHECKING: from .hooks import CompactionHook as CompactionHook from .sliding_window import SlidingWindowCompactor as SlidingWindowCompactor + from .summarization import SummarizationCompactor as SummarizationCompactor from .tool_result_pruning import ToolResultPruningCompactor as ToolResultPruningCompactor from .types import Compactor as Compactor else: diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index c916e5958ef..b3b90c44604 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -7,7 +7,15 @@ from haystack.core.serialization import default_to_dict from haystack.dataclasses import ChatMessage, ChatRole from haystack.hooks.compaction.types import Compactor -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _agent_step_spans +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _current_step_groups, + _historical_turn_groups, + _is_compaction_message, + _latest_user_index, + _leading_system_end, + _messages_except, +) from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental @@ -23,75 +31,9 @@ ) -def _leading_system_end(messages: list[ChatMessage]) -> int: - """Return the end of the leading system-message block.""" - for index, message in enumerate(messages): - # Find the first non-leading system message or a system message produced by compaction - if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: - return index - return len(messages) - - -def _latest_user_index(messages: list[ChatMessage]) -> int | None: - """ - Return the latest user message not produced by compaction. - - :param messages: The conversation to analyze, oldest to newest. - """ - # We loop backwards to find the latest user message - for index in reversed(range(len(messages))): - message = messages[index] - # Find the latest user message that was not produced by a previous compaction - if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: - return index - return None - - def _is_compaction_note(message: ChatMessage) -> bool: """Whether a message is an omission note this strategy left in place of removed history.""" - marker = message.meta.get(_COMPACTION_META_KEY) - return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY - - -def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: - """ - Return spans for complete user turns in a bounded section of conversation history. - - Each turn begins with a real user message and continues up to, but does not include, the next real user message. - This groups a user's request with every assistant step and tool result produced in response to it. - - :param messages: The full conversation to analyze, ordered oldest to newest. - :param start: The inclusive index at which to begin looking for historical turns. - :param end: The exclusive index at which to stop. This is normally the current task's user-message index. - :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to - `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. - """ - # Reject any user-role message an earlier compaction produced, whichever strategy made it: none of them are user - # requests, so none of them begin a turn. Leaving them out also lets this compaction fold an old note away. - user_indices = [ - index - for index in range(start, end) - if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta - ] - - # A real user message closes the preceding turn and starts the next one. The final historical turn extends to the - # supplied boundary, which is typically where the protected current task begins. - return [ - (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) - for position, index in enumerate(user_indices) - ] - - -def _index_groups( - messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False -) -> list[list[int]]: - """ - Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. - """ - return [ - [index for index in range(start, end) if not (skip_compaction_notes and _is_compaction_note(messages[index]))] - for start, end in spans - ] + return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER) def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: @@ -99,12 +41,6 @@ def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMe return [messages[index] for index in indices] -def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages the given indices leave out, in conversation order.""" - left_out = set(indices) - return [message for index, message in enumerate(messages) if index not in left_out] - - def _flatten(groups: list[list[int]]) -> list[int]: """Join index groups into a single ordered list of indices.""" return [index for group in groups for index in group] @@ -124,19 +60,13 @@ def _removable_groups( kept or removed entire, which is what keeps a tool call with its results and an assistant reply with the user message it answers. """ - # Steps belong to the current task, so they start after its anchor, or after the instructions when the - # conversation has no user message to anchor on. - step_start = (task_index + 1) if task_index is not None else system_end - step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start)) - # An earlier compaction's note is left out of its turn, so keeping the turn folds that note into the note this # compaction leaves behind. - historical_end = task_index if task_index is not None else system_end - historical_groups = _index_groups( - messages=messages, - spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), - skip_compaction_notes=True, - ) + historical_groups = [ + [index for index in group if not _is_compaction_note(message=messages[index])] + for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + ] + step_groups = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) return historical_groups, step_groups diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py new file mode 100644 index 00000000000..b466a7ee798 --- /dev/null +++ b/haystack/hooks/compaction/summarization.py @@ -0,0 +1,588 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Awaitable, Callable +from typing import Any + +from haystack import logging +from haystack.components.generators.chat.types import ChatGenerator +from haystack.components.generators.chat.utils import _resolve_output_token_limit +from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict +from haystack.dataclasses import ChatMessage +from haystack.hooks.compaction.types import Compactor +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _current_step_groups, + _historical_turn_groups, + _is_compaction_message, + _latest_user_index, + _leading_system_end, + _messages_except, +) +from haystack.token_counters import TokenCounter +from haystack.token_counters.utils import _render_message +from haystack.utils.async_utils import _execute_component_async +from haystack.utils.deserialization import deserialize_component_inplace +from haystack.utils.experimental import _experimental + +logger = logging.getLogger(__name__) + +_STRATEGY = "summarization" + +_DEFAULT_SUMMARY_INSTRUCTION = """You are compacting part of a conversation between a user and an AI agent so the \ +agent can keep working with fewer tokens. Write a self-contained summary that preserves: +- The user's goal, requirements, constraints, and preferences. +- Decisions and the reasoning behind them. +- Work already completed and important tool results. +- Exact file paths, URLs, identifiers, and references to stored data. +- Unresolved work and the immediate next step. + +Fold any existing blocks into one summary. Record only what the conversation shows. Do not \ +infer or add advice. Use plain prose or short bullets, and do not address the user.""" + + +def _indices(messages: list[ChatMessage], start: int, end: int, *, summaries: bool) -> list[int]: + """Return summary or non-summary indices in a bounded part of a conversation.""" + return [ + index + for index in range(start, end) + if _is_compaction_message(message=messages[index], strategy=_STRATEGY) is summaries + ] + + +def _raw_historical_turn_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> list[list[int]]: + """Return shared historical-turn groups with this strategy's summaries filtered out.""" + return [ + [index for index in group if not _is_compaction_message(message=messages[index], strategy=_STRATEGY)] + for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + ] + + +def _groups_to_summarize( + messages: list[ChatMessage], + groups: list[list[int]], + target_tokens: int, + summary_budget: int, + token_counter: TokenCounter, +) -> list[int]: + """Select the fewest oldest groups that should make room for a summary of the configured size.""" + selected: list[int] = [] + for group in groups: + selected.extend(group) + if ( + token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + summary_budget + <= target_tokens + ): + break + return selected + + +def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMessage: + """Build a marked summary message.""" + body = f"\n{text.strip()}\n" + meta = {_COMPACTION_META_KEY: {"strategy": _STRATEGY, "summarized_messages": summarized_messages, "source": source}} + return ChatMessage.from_user(text=body, meta=meta) + + +def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: + """Replace possibly non-contiguous selected messages with one summary at their first position.""" + selected = set(indices) + insertion_index = min(indices) + compacted: list[ChatMessage] = [] + for index, message in enumerate(messages): + if index == insertion_index: + compacted.append(summary) + if index not in selected: + compacted.append(message) + return compacted + + +@_experimental +class SummarizationCompactor(Compactor): + """ + Condenses old historical turns first, then old steps from the Agent's current task. + + Leading system messages and the latest real user message are always kept. Historical turns are summarized in full, + oldest first. Summaries normally accumulate so they are not repeatedly rewritten; if every historical turn has + already been summarized and more space is needed, those historical summaries are folded into one before any + current-task steps are summarized. An assistant message and all immediately following tool results form one step, + so tool calls are never separated from their results. + + Each summary is requested within `max_summary_tokens`. Built-in OpenAI Chat Completions and Responses generators + receive the corresponding runtime output limit unless they already configure one in their `generation_kwargs`, in + which case the generator's setting wins. Other generators receive the same limit as prompt guidance and the actual + result is measured before it is accepted. + + ```python + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIResponsesChatGenerator + from haystack.hooks.compaction import CompactionHook, SummarizationCompactor + + summary_generator = OpenAIResponsesChatGenerator(model="gpt-5.4-nano") + hook = CompactionHook( + compactor=SummarizationCompactor(chat_generator=summary_generator), + context_window=400_000, + compact_at=0.7, + compact_to=0.4, + ) + agent = Agent(chat_generator=agent_generator, tools=[web_search], hooks={"before_llm": [hook]}) + ``` + """ + + def __init__( + self, + chat_generator: ChatGenerator, + *, + min_keep_steps: int = 1, + max_summary_tokens: int = 1024, + summary_instruction: str | None = None, + raise_on_failure: bool = False, + ) -> None: + """ + Initialize the compactor. + + :param chat_generator: The Chat Generator used to write summaries. Its configured output-token limit takes + precedence over `max_summary_tokens` when the generator exposes a recognized setting. + :param min_keep_steps: The fewest complete recent Agent steps to keep, even when they exceed the target. + :param max_summary_tokens: The output-token budget reserved for each summary. Known built-in generators receive + the corresponding runtime generation setting unless one is already configured on the generator. + :param summary_instruction: An instruction replacing the built-in summary prompt. + :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is + logged and any successful partial compaction is returned. + """ + if min_keep_steps < 0: + raise ValueError(f"`min_keep_steps` must be at least 0, got {min_keep_steps}.") + if max_summary_tokens < 1: + raise ValueError(f"`max_summary_tokens` must be a positive number of tokens, got {max_summary_tokens}.") + self.chat_generator = chat_generator + self.min_keep_steps = min_keep_steps + self.max_summary_tokens = max_summary_tokens + self.summary_instruction = summary_instruction or _DEFAULT_SUMMARY_INSTRUCTION + self.raise_on_failure = raise_on_failure + + def compact( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Return a progressively summarized conversation, or None when no useful reduction is possible. + + :param messages: The conversation to compact, ordered oldest to newest. + :param target_tokens: The token budget the compacted messages should aim to fit. + :param token_counter: The counter used both to plan compaction and verify generated summaries. + :returns: A smaller replacement conversation, or None when nothing was reduced. + """ + if token_counter.count(messages=messages) <= target_tokens: + return None + budget, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + + def generate(prompt: list[ChatMessage]) -> dict[str, Any]: + kwargs: dict[str, Any] = {"messages": prompt} + if generation_kwargs is not None: + kwargs["generation_kwargs"] = generation_kwargs + return self.chat_generator.run(**kwargs) + + return self._compact( + original=messages, + target_tokens=target_tokens, + token_counter=token_counter, + budget=budget, + generate=generate, + ) + + async def compact_async( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Asynchronously return a progressively summarized conversation. + + :param messages: The conversation to compact, ordered oldest to newest. + :param target_tokens: The token budget the compacted messages should aim to fit. + :param token_counter: The counter used both to plan compaction and verify generated summaries. + :returns: A smaller replacement conversation, or None when nothing was reduced. + """ + if token_counter.count(messages=messages) <= target_tokens: + return None + budget, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + + async def generate(prompt: list[ChatMessage]) -> dict[str, Any]: + kwargs: dict[str, Any] = {"messages": prompt} + if generation_kwargs is not None: + kwargs["generation_kwargs"] = generation_kwargs + return await _execute_component_async(component_instance=self.chat_generator, **kwargs) + + return await self._compact_async( + original=messages, + target_tokens=target_tokens, + token_counter=token_counter, + budget=budget, + generate=generate, + ) + + def _prompt(self, messages: list[ChatMessage], budget: int) -> list[ChatMessage]: + """Build the bounded summarization instruction and rendered source transcript.""" + transcript = "\n".join(_render_message(message=message) for message in messages) + instruction = ( + f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately {budget} tokens. " + "Prioritize completeness within that limit so the response is not cut off." + ) + return [ + ChatMessage.from_system(text=instruction), + ChatMessage.from_user(text=f"\n{transcript}\n"), + ] + + def _apply_result( + self, + messages: list[ChatMessage], + indices: list[int], + source: str, + result: dict[str, Any], + token_counter: TokenCounter, + ) -> list[ChatMessage]: + """Validate a generator reply and replace its source messages when the result is smaller.""" + replies = result.get("replies") or [] + text = replies[-1].text if replies else None + if not text or not text.strip(): + raise RuntimeError("The Chat Generator returned no text to use as a conversation summary.") + summary = _summary_message(text=text, summarized_messages=len(indices), source=source) + compacted = _replace_indices(messages=messages, indices=indices, summary=summary) + before = token_counter.count(messages=messages) + after = token_counter.count(messages=compacted) + if after >= before: + raise RuntimeError( + f"The generated summary did not reduce the conversation size ({before} tokens before and {after} " + "tokens after)." + ) + return compacted + + def _attempt( + self, + messages: list[ChatMessage], + indices: list[int], + source: str, + budget: int, + token_counter: TokenCounter, + generate: Callable[[list[ChatMessage]], dict[str, Any]], + ) -> list[ChatMessage] | None: + """Attempt one synchronous summary, applying the configured failure policy.""" + try: + result = generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) + return self._apply_result( + messages=messages, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + if self.raise_on_failure: + raise + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) + return None + + async def _attempt_async( + self, + messages: list[ChatMessage], + indices: list[int], + source: str, + budget: int, + token_counter: TokenCounter, + generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], + ) -> list[ChatMessage] | None: + """Attempt one asynchronous summary, applying the configured failure policy.""" + try: + result = await generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) + return self._apply_result( + messages=messages, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + if self.raise_on_failure: + raise + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) + return None + + def _compact( + self, + original: list[ChatMessage], + target_tokens: int, + token_counter: TokenCounter, + budget: int, + generate: Callable[[list[ChatMessage]], dict[str, Any]], + ) -> list[ChatMessage] | None: + """Run synchronous historical, consolidation, and current-step compaction tiers in order.""" + working = list(original) + + # First replace the fewest oldest raw historical turns expected to reach the target. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + groups = [ + group + for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) + if group + ] + if not groups: + break + selected = _groups_to_summarize( + messages=working, + groups=groups, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = self._attempt( + messages=working, + indices=selected, + source="historical_turns", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # If all historical turns are summaries and the target is still unmet, fold them into one summary. + if token_counter.count(messages=working) > target_tokens: + summaries = self._summary_indices(messages=working, source="historical_summaries") + if len(summaries) > 1: + compacted = self._attempt( + messages=working, + indices=summaries, + source="historical_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + if not eligible: + break + + # Fold accumulated current-task summaries before consuming more raw agent steps. + summaries = self._summary_indices(messages=working, source="current_task_summaries") + if len(summaries) > 1: + compacted = self._attempt( + messages=working, + indices=summaries, + source="current_task_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + selected = _groups_to_summarize( + messages=working, + groups=eligible, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = self._attempt( + messages=working, + indices=selected, + source="current_task_steps", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + break + working = compacted + return self._result(original=original, working=working, token_counter=token_counter) + + async def _compact_async( + self, + original: list[ChatMessage], + target_tokens: int, + token_counter: TokenCounter, + budget: int, + generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], + ) -> list[ChatMessage] | None: + """Run asynchronous historical, consolidation, and current-step compaction tiers in order.""" + working = list(original) + + # First replace the fewest oldest raw historical turns expected to reach the target. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + groups = [ + group + for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) + if group + ] + if not groups: + break + selected = _groups_to_summarize( + messages=working, + groups=groups, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = await self._attempt_async( + messages=working, + indices=selected, + source="historical_turns", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # If all historical turns are summaries and the target is still unmet, fold them into one summary. + if token_counter.count(messages=working) > target_tokens: + summaries = self._summary_indices(messages=working, source="historical_summaries") + if len(summaries) > 1: + compacted = await self._attempt_async( + messages=working, + indices=summaries, + source="historical_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + if not eligible: + break + + # Fold accumulated current-task summaries before consuming more raw agent steps. + summaries = self._summary_indices(messages=working, source="current_task_summaries") + if len(summaries) > 1: + compacted = await self._attempt_async( + messages=working, + indices=summaries, + source="current_task_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + selected = _groups_to_summarize( + messages=working, + groups=eligible, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = await self._attempt_async( + messages=working, + indices=selected, + source="current_task_steps", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + break + working = compacted + return self._result(original=original, working=working, token_counter=token_counter) + + def _summary_indices(self, messages: list[ChatMessage], source: str) -> list[int]: + """Return historical or current-task summary indices based on their conversation position.""" + system_end = _leading_system_end(messages=messages) + task_index = _latest_user_index(messages=messages) + end = task_index if task_index is not None else len(messages) + if source == "historical_summaries": + return _indices(messages=messages, start=system_end, end=end, summaries=True) + start = task_index + 1 if task_index is not None else system_end + return _indices(messages=messages, start=start, end=len(messages), summaries=True) + + @staticmethod + def _result( + original: list[ChatMessage], working: list[ChatMessage], token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """Return partial or complete progress only when it reduced the original conversation.""" + if token_counter.count(messages=working) < token_counter.count(messages=original): + return working + return None + + def warm_up(self) -> None: + """Warm up the Chat Generator that writes summaries.""" + if hasattr(self.chat_generator, "warm_up"): + self.chat_generator.warm_up() + + async def warm_up_async(self) -> None: + """Warm up the Chat Generator on the serving event loop.""" + warm_up_async = getattr(self.chat_generator, "warm_up_async", None) + if warm_up_async is not None: + await warm_up_async() + elif hasattr(self.chat_generator, "warm_up"): + self.chat_generator.warm_up() + + def close(self) -> None: + """Release the Chat Generator's resources.""" + if hasattr(self.chat_generator, "close"): + self.chat_generator.close() + + async def close_async(self) -> None: + """Release the Chat Generator's resources.""" + close_async = getattr(self.chat_generator, "close_async", None) + if close_async is not None: + await close_async() + elif hasattr(self.chat_generator, "close"): + self.chat_generator.close() + + def to_dict(self) -> dict[str, Any]: + """Serialize the compactor and its Chat Generator.""" + return default_to_dict( + self, + chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"), + min_keep_steps=self.min_keep_steps, + max_summary_tokens=self.max_summary_tokens, + summary_instruction=self.summary_instruction, + raise_on_failure=self.raise_on_failure, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "SummarizationCompactor": + """Deserialize the compactor and reconstruct its Chat Generator.""" + init_params = data.get("init_parameters", {}) + if init_params.get("chat_generator") is not None: + deserialize_component_inplace(data=init_params, key="chat_generator") + return default_from_dict(cls=cls, data=data) diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 6a6d1402c52..99b0fe79d8e 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -10,6 +10,36 @@ _COMPACTION_META_KEY = "context_compaction" +def _leading_system_end(messages: list[ChatMessage]) -> int: + """Return the end of the leading system-message block, excluding system messages created by compaction.""" + for index, message in enumerate(messages): + if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: + return index + return len(messages) + + +def _latest_user_index(messages: list[ChatMessage]) -> int | None: + """Return the latest user message not produced by compaction.""" + for index in reversed(range(len(messages))): + message = messages[index] + if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: + return index + return None + + +def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages the given indices leave out, in conversation order.""" + left_out = set(indices) + return [message for index, message in enumerate(messages) if index not in left_out] + + +def _is_compaction_message(message: ChatMessage, strategy: str, role: ChatRole | None = None) -> bool: + """Return whether a message was produced by a compaction strategy and optionally has the requested role.""" + marker = message.meta.get(_COMPACTION_META_KEY) + has_role = role is None or message.is_from(role=role) + return has_role and isinstance(marker, dict) and marker.get("strategy") == strategy + + def _last_assistant_index(messages: list[ChatMessage]) -> int: """Return the index of the last assistant message, or -1 if none exists.""" for index in reversed(range(len(messages))): @@ -47,6 +77,38 @@ def _agent_step_spans(messages: list[ChatMessage], start: int) -> list[tuple[int return spans +def _current_step_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """Return message-index groups for complete Agent steps belonging to the current task.""" + step_start = task_index + 1 if task_index is not None else system_end + return [list(range(start, end)) for start, end in _agent_step_spans(messages=messages, start=step_start)] + + +def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: + """ + Return spans for complete real-user turns in a bounded section of conversation history. + + Each turn begins with a user message not created by compaction and continues up to the next such message. + """ + user_indices = [ + index + for index in range(start, end) + if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta + ] + return [ + (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) + for position, index in enumerate(user_indices) + ] + + +def _historical_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """Return message-index groups for complete historical turns preceding the current task.""" + historical_end = task_index if task_index is not None else system_end + return [ + list(range(start, end)) + for start, end in _historical_turn_spans(messages=messages, start=system_end, end=historical_end) + ] + + def _estimated_context_tokens( messages: list[ChatMessage], context_tokens: int, token_counter: TokenCounter, tools: ToolsType | None = None ) -> int: diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml index 4fed6fb1452..fef10aad17e 100644 --- a/pydoc/hooks_api.yml +++ b/pydoc/hooks_api.yml @@ -1,7 +1,7 @@ loaders: - search_path: [../haystack/hooks] modules: ["protocol", "from_function", "compaction/hooks", "compaction/sliding_window", - "compaction/tool_result_pruning", + "compaction/summarization", "compaction/tool_result_pruning", "compaction/types/protocol", "human_in_the_loop/dataclasses", "human_in_the_loop/hooks", "human_in_the_loop/policies", "human_in_the_loop/strategies", "human_in_the_loop/user_interfaces", "tool_result_offloading/hooks", "tool_result_offloading/policies", "tool_result_offloading/stores", diff --git a/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml b/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml new file mode 100644 index 00000000000..d6125f531b5 --- /dev/null +++ b/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml @@ -0,0 +1,6 @@ +--- +features: + - | + Added the experimental ``SummarizationCompactor``. It uses a Chat Generator to condense the Agent's oldest + historical turns, accumulated summaries, and old current-task steps as needed to reach the context target while + preserving leading system messages, the latest user task, and complete recent tool-calling steps. diff --git a/test/components/generators/chat/test_utils.py b/test/components/generators/chat/test_utils.py new file mode 100644 index 00000000000..6d64d54dd26 --- /dev/null +++ b/test/components/generators/chat/test_utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from haystack.components.generators.chat import MockChatGenerator +from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _resolve_output_token_limit + + +class OpenAIChatGenerator(MockChatGenerator): + def __init__(self, generation_kwargs=None): + super().__init__("response") + self.generation_kwargs = generation_kwargs or {} + + +class OpenAIResponsesChatGenerator(MockChatGenerator): + def __init__(self, generation_kwargs=None): + super().__init__("response") + self.generation_kwargs = generation_kwargs or {} + + +class CustomGenerator(MockChatGenerator): + generation_kwargs = {"max_tokens": 7} + + +def test_resolves_chat_completions_limit(): + generator = OpenAIChatGenerator() + assert _generator_output_token_limit_key(generator) == "max_completion_tokens" + assert _resolve_output_token_limit(generator, 100) == (100, {"max_completion_tokens": 100}) + + +def test_resolves_responses_limit(): + generator = OpenAIResponsesChatGenerator() + assert _generator_output_token_limit_key(generator) == "max_output_tokens" + assert _resolve_output_token_limit(generator, 100) == (100, {"max_output_tokens": 100}) + + +def test_configured_generator_limit_wins_without_mutation(): + generator = OpenAIChatGenerator({"temperature": 0, "max_completion_tokens": 23}) + original = dict(generator.generation_kwargs) + assert _resolve_output_token_limit(generator, 100) == (23, None) + assert generator.generation_kwargs == original + + +def test_unknown_generator_receives_no_guessed_runtime_setting(): + assert _generator_output_token_limit_key(CustomGenerator()) is None + assert _resolve_output_token_limit(CustomGenerator(), 100) == (100, None) diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 332332b71b6..d847c50fe55 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -6,8 +6,8 @@ from haystack.dataclasses import ChatMessage, ChatRole, ToolCall from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans, _is_compaction_note -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _is_compaction_note +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _historical_turn_spans from test.hooks.compaction.helpers import ( FakeCounter, count_markers, diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py new file mode 100644 index 00000000000..62064ffcd6d --- /dev/null +++ b/test/hooks/compaction/test_summarization.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +import pytest + +from haystack.components.generators.chat import MockChatGenerator +from haystack.dataclasses import ChatMessage, ChatRole +from haystack.hooks.compaction import SummarizationCompactor +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result + +pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") +COUNTER = FakeCounter(chars_per_token=1) + + +def recording_generator(responses: list[str | Exception]) -> tuple[MockChatGenerator, list[dict[str, Any]]]: + """Build a MockChatGenerator whose response function records prompts and can raise queued failures.""" + queued = list(responses) + calls: list[dict[str, Any]] = [] + + def respond(messages: list[ChatMessage]) -> str: + calls.append({"messages": messages}) + response = queued.pop(0) + if isinstance(response, Exception): + raise response + return response + + return MockChatGenerator(response_fn=respond), calls + + +def summary(text: str, source: str) -> ChatMessage: + return ChatMessage.from_user( + f"\n{text}\n", + meta={_COMPACTION_META_KEY: {"strategy": "summarization", "source": source}}, + ) + + +def transcript(call: dict[str, Any]) -> str: + return call["messages"][-1].text + + +class TestSummarizationCompactor: + def test_summarizes_the_minimum_number_of_oldest_historical_turns(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + generator, calls = recording_generator(["short historical summary"]) + retained_without_oldest = [messages[0], *messages[3:]] + target = COUNTER.count(retained_without_oldest) + 100 + compacted = SummarizationCompactor(generator, max_summary_tokens=100).compact( + messages=messages, target_tokens=target, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 1 + assert "oldest question" in transcript(calls[0]) + assert "recent question" not in transcript(calls[0]) + assert compacted[0] == messages[0] + assert compacted[2:] == messages[3:] + assert messages == [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + + def test_reaches_current_steps_in_the_same_call_after_historical_context(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question " * 40), + ChatMessage.from_assistant("old answer " * 40), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 40, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + generator, calls = recording_generator(["history", "old step"]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert "old question" in transcript(calls[0]) + assert "old result" in transcript(calls[1]) + assert compacted[-2:] == messages[-2:] + assert [m.meta[_COMPACTION_META_KEY]["source"] for m in compacted if _COMPACTION_META_KEY in m.meta] == [ + "historical_turns", + "current_task_steps", + ] + + def test_consolidates_historical_summaries_before_current_steps(self): + messages = [ + ChatMessage.from_system("rules"), + summary("first history " * 20, "historical_turns"), + summary("second history " * 20, "historical_turns"), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + generator, calls = recording_generator(["combined history", "old step"]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert "first history" in transcript(calls[0]) + assert "old result" not in transcript(calls[0]) + assert "old result" in transcript(calls[1]) + + def test_consolidates_current_summaries_before_more_raw_steps(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + summary("first step summary " * 20, "current_task_steps"), + summary("second step summary " * 20, "current_task_steps"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + generator, calls = recording_generator(["combined steps", "old step"]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert "first step summary" in transcript(calls[0]) + assert "old result" in transcript(calls[1]) + assert compacted[-2:] == messages[-2:] + + def test_returns_partial_progress_when_a_later_tier_fails_by_default(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question " * 30), + ChatMessage.from_assistant("old answer " * 30), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("old step " * 30), + ChatMessage.from_assistant("new step"), + ] + generator, calls = recording_generator(["history", RuntimeError("provider unavailable")]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert any( + message.meta.get(_COMPACTION_META_KEY, {}).get("source") == "historical_turns" for message in compacted + ) + assert messages[-2:] == compacted[-2:] + + def test_raises_when_configured_and_summary_does_not_shrink_context(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old"), + ChatMessage.from_assistant("answer"), + ChatMessage.from_user("current"), + ] + generator, _ = recording_generator(["much longer summary " * 100]) + with pytest.raises(RuntimeError, match="did not reduce"): + SummarizationCompactor(generator, max_summary_tokens=1, raise_on_failure=True).compact(messages, 1, COUNTER) + + def test_returns_none_without_calling_generator_when_context_fits(self): + generator, calls = recording_generator(["unused"]) + messages = [ChatMessage.from_system("rules"), ChatMessage.from_user("task")] + assert ( + SummarizationCompactor(generator).compact(messages=messages, target_tokens=10_000, token_counter=COUNTER) + is None + ) + assert calls == [] + + def test_summary_is_a_user_message_with_compaction_metadata(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old " * 100), + ChatMessage.from_assistant("answer " * 100), + ChatMessage.from_user("task"), + ] + compacted = SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + assert compacted is not None + generated = compacted[1] + assert generated.is_from(ChatRole.USER) + assert generated.meta[_COMPACTION_META_KEY] == { + "strategy": "summarization", + "summarized_messages": 2, + "source": "historical_turns", + } + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"min_keep_steps": -1}, "`min_keep_steps` must be at least 0"), + ({"max_summary_tokens": 0}, "`max_summary_tokens` must be a positive"), + ], + ) + def test_rejects_invalid_settings(self, kwargs, match): + with pytest.raises(ValueError, match=match): + SummarizationCompactor(MockChatGenerator("summary"), **kwargs) + + def test_serde_round_trip(self): + compactor = SummarizationCompactor( + MockChatGenerator("summary"), + min_keep_steps=2, + max_summary_tokens=321, + summary_instruction="custom", + raise_on_failure=True, + ) + restored = SummarizationCompactor.from_dict(compactor.to_dict()) + assert isinstance(restored.chat_generator, MockChatGenerator) + assert restored.min_keep_steps == 2 + assert restored.max_summary_tokens == 321 + assert restored.summary_instruction == "custom" + assert restored.raise_on_failure is True + + +@pytest.mark.asyncio +async def test_async_compaction_uses_async_generator(): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old " * 100), + ChatMessage.from_assistant("answer " * 100), + ChatMessage.from_user("task"), + ] + generator, calls = recording_generator(["async summary"]) + compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + assert compacted is not None + assert len(calls) == 1 From 952250d90cca9539f79180cb473aaec3c4eae49c Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:45:56 +0200 Subject: [PATCH 15/38] refactor to make it more understandable --- haystack/hooks/compaction/sliding_window.py | 6 +- haystack/hooks/compaction/summarization.py | 510 +++++++------------- haystack/hooks/compaction/utils.py | 5 + test/hooks/compaction/test_summarization.py | 254 +++++----- 4 files changed, 325 insertions(+), 450 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index b3b90c44604..59a5a3cd6d4 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -14,6 +14,7 @@ _is_compaction_message, _latest_user_index, _leading_system_end, + _messages_at, _messages_except, ) from haystack.token_counters import TokenCounter @@ -36,11 +37,6 @@ def _is_compaction_note(message: ChatMessage) -> bool: return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER) -def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages at the given indices, in the order the indices are given.""" - return [messages[index] for index in indices] - - def _flatten(groups: list[list[int]]) -> list[int]: """Join index groups into a single ordered list of indices.""" return [index for group in groups for index in group] diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index b466a7ee798..2fa2add59a9 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -2,7 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Awaitable, Callable from typing import Any from haystack import logging @@ -18,18 +17,27 @@ _is_compaction_message, _latest_user_index, _leading_system_end, + _messages_at, _messages_except, ) from haystack.token_counters import TokenCounter -from haystack.token_counters.utils import _render_message +from haystack.token_counters.utils import _rendered_conversation from haystack.utils.async_utils import _execute_component_async from haystack.utils.deserialization import deserialize_component_inplace from haystack.utils.experimental import _experimental logger = logging.getLogger(__name__) +# Recorded as the strategy on every summary this compactor produces, so a later run can recognize its own summaries. _STRATEGY = "summarization" +# Recorded as the `source` on a summary, naming the stretch of conversation it stands in for. Compaction gives these up +# in order, so the Agent's current task is the last thing to go. +_HISTORICAL_TURNS = "historical_turns" +_HISTORICAL_SUMMARIES = "historical_summaries" +_CURRENT_TASK_SUMMARIES = "current_task_summaries" +_CURRENT_TASK_STEPS = "current_task_steps" + _DEFAULT_SUMMARY_INSTRUCTION = """You are compacting part of a conversation between a user and an AI agent so the \ agent can keep working with fewer tokens. Write a self-contained summary that preserves: - The user's goal, requirements, constraints, and preferences. @@ -42,53 +50,61 @@ infer or add advice. Use plain prose or short bullets, and do not address the user.""" -def _indices(messages: list[ChatMessage], start: int, end: int, *, summaries: bool) -> list[int]: - """Return summary or non-summary indices in a bounded part of a conversation.""" - return [ - index - for index in range(start, end) - if _is_compaction_message(message=messages[index], strategy=_STRATEGY) is summaries - ] +def _is_summary(message: ChatMessage) -> bool: + """Whether a message is a summary this strategy wrote.""" + return _is_compaction_message(message=message, strategy=_STRATEGY) + + +def _summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: + """Return the positions of this strategy's summaries within a bounded part of a conversation.""" + return [index for index in range(start, end) if _is_summary(message=messages[index])] -def _raw_historical_turn_groups( - messages: list[ChatMessage], system_end: int, task_index: int | None -) -> list[list[int]]: - """Return shared historical-turn groups with this strategy's summaries filtered out.""" - return [ - [index for index in group if not _is_compaction_message(message=messages[index], strategy=_STRATEGY)] +def _summarizable_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """ + Return the historical turns that still hold raw conversation, oldest turn first. + + Summaries this strategy already wrote are left out of their turn, so summarizing the turn folds that summary into + the summary this run produces. A turn that is nothing but summaries has nothing left to give up and is dropped. + """ + groups = [ + [index for index in group if not _is_summary(message=messages[index])] for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) ] + return [group for group in groups if group] def _groups_to_summarize( messages: list[ChatMessage], groups: list[list[int]], target_tokens: int, - summary_budget: int, + summary_tokens: int, token_counter: TokenCounter, ) -> list[int]: - """Select the fewest oldest groups that should make room for a summary of the configured size.""" + """ + Return the fewest oldest groups whose removal makes room for a summary of the expected size. + + Groups are taken oldest first and counting stops as soon as what remains, plus the summary that replaces them, + fits the target. When even taking all of them is not enough, all of them are returned. + """ selected: list[int] = [] for group in groups: selected.extend(group) - if ( - token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + summary_budget - <= target_tokens - ): + remaining = token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + if remaining + summary_tokens <= target_tokens: break return selected def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMessage: - """Build a marked summary message.""" + """Build the marked user message that stands in for the messages the summary replaced.""" body = f"\n{text.strip()}\n" meta = {_COMPACTION_META_KEY: {"strategy": _STRATEGY, "summarized_messages": summarized_messages, "source": source}} return ChatMessage.from_user(text=body, meta=meta) def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: - """Replace possibly non-contiguous selected messages with one summary at their first position.""" + """Replace the selected messages, which need not be contiguous, with one summary at the oldest one's position.""" selected = set(indices) insertion_index = min(indices) compacted: list[ChatMessage] = [] @@ -152,6 +168,7 @@ def __init__( :param summary_instruction: An instruction replacing the built-in summary prompt. :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is logged and any successful partial compaction is returned. + :raises ValueError: If `min_keep_steps` is negative or `max_summary_tokens` is not positive. """ if min_keep_steps < 0: raise ValueError(f"`min_keep_steps` must be at least 0, got {min_keep_steps}.") @@ -174,25 +191,28 @@ def compact( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ - if token_counter.count(messages=messages) <= target_tokens: - return None - budget, generation_kwargs = _resolve_output_token_limit( - chat_generator=self.chat_generator, default_limit=self.max_summary_tokens - ) - - def generate(prompt: list[ChatMessage]) -> dict[str, Any]: - kwargs: dict[str, Any] = {"messages": prompt} - if generation_kwargs is not None: - kwargs["generation_kwargs"] = generation_kwargs - return self.chat_generator.run(**kwargs) - - return self._compact( - original=messages, - target_tokens=target_tokens, - token_counter=token_counter, - budget=budget, - generate=generate, - ) + summary_tokens, run_kwargs = self._summary_limit() + working = list(messages) + while True: + plan = self._next_summary( + messages=working, + target_tokens=target_tokens, + token_counter=token_counter, + summary_tokens=summary_tokens, + ) + if plan is None: + break + indices, source = plan + prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + try: + result = self.chat_generator.run(messages=prompt, **run_kwargs) + working = self._apply_summary( + messages=working, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + self._report_failure(error=error) + break + return self._reduced(original=messages, working=working, token_counter=token_counter) async def compact_async( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter @@ -205,47 +225,129 @@ async def compact_async( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ + summary_tokens, run_kwargs = self._summary_limit() + working = list(messages) + while True: + plan = self._next_summary( + messages=working, + target_tokens=target_tokens, + token_counter=token_counter, + summary_tokens=summary_tokens, + ) + if plan is None: + break + indices, source = plan + prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + try: + result = await _execute_component_async( + component_instance=self.chat_generator, messages=prompt, **run_kwargs + ) + working = self._apply_summary( + messages=working, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + self._report_failure(error=error) + break + return self._reduced(original=messages, working=working, token_counter=token_counter) + + def _next_summary( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int + ) -> tuple[list[int], str] | None: + """ + Choose the next stretch of conversation to replace with a summary. + + Four tiers are tried in order, so the oldest and least useful context goes first and the Agent's current task + is given up last: + + 1. `_HISTORICAL_TURNS`: the fewest oldest raw turns that make room for a summary. + 2. `_HISTORICAL_SUMMARIES`: nothing raw is left in history, so fold its summaries into one. + 3. `_CURRENT_TASK_SUMMARIES`: fold the summaries earlier steps left behind before giving up more steps. + 4. `_CURRENT_TASK_STEPS`: the fewest oldest steps of the current task, keeping `min_keep_steps` of the newest. + + :param messages: The conversation as it stands, ordered oldest to newest. + :param target_tokens: The token budget the conversation should come in under. + :param token_counter: The counter used to measure candidate selections. + :param summary_tokens: The size a summary is expected to take, reserved when choosing how much to replace. + :returns: The message indices to summarize and the `source` to record on the resulting summary, or None when + the conversation already fits or nothing is left that may be given up. + """ if token_counter.count(messages=messages) <= target_tokens: return None - budget, generation_kwargs = _resolve_output_token_limit( - chat_generator=self.chat_generator, default_limit=self.max_summary_tokens - ) - async def generate(prompt: list[ChatMessage]) -> dict[str, Any]: - kwargs: dict[str, Any] = {"messages": prompt} - if generation_kwargs is not None: - kwargs["generation_kwargs"] = generation_kwargs - return await _execute_component_async(component_instance=self.chat_generator, **kwargs) + # The landmarks everything is measured against: the Agent's instructions, and the user message anchoring the + # current task. History runs from the instructions up to that anchor, the current task from the anchor on. + system_end = _leading_system_end(messages=messages) + task_index = _latest_user_index(messages=messages) + history_end = task_index if task_index is not None else system_end + task_start = task_index + 1 if task_index is not None else system_end + + turns = _summarizable_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + if turns: + oldest_turns = _groups_to_summarize( + messages=messages, + groups=turns, + target_tokens=target_tokens, + summary_tokens=summary_tokens, + token_counter=token_counter, + ) + return oldest_turns, _HISTORICAL_TURNS + + history_summaries = _summary_indices(messages=messages, start=system_end, end=history_end) + if len(history_summaries) > 1: + return history_summaries, _HISTORICAL_SUMMARIES + + # Only steps older than the `min_keep_steps` most recent ones may be given up. + steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) + eligible = steps[: max(len(steps) - self.min_keep_steps, 0)] + if not eligible: + return None + + task_summaries = _summary_indices(messages=messages, start=task_start, end=len(messages)) + if len(task_summaries) > 1: + return task_summaries, _CURRENT_TASK_SUMMARIES - return await self._compact_async( - original=messages, + oldest_steps = _groups_to_summarize( + messages=messages, + groups=eligible, target_tokens=target_tokens, + summary_tokens=summary_tokens, token_counter=token_counter, - budget=budget, - generate=generate, ) + return oldest_steps, _CURRENT_TASK_STEPS - def _prompt(self, messages: list[ChatMessage], budget: int) -> list[ChatMessage]: - """Build the bounded summarization instruction and rendered source transcript.""" - transcript = "\n".join(_render_message(message=message) for message in messages) + def _summary_limit(self) -> tuple[int, dict[str, Any]]: + """Return the token budget for one summary and the run kwargs, if any, that ask the generator to honor it.""" + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + return summary_tokens, {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + + def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: + """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" + transcript = _rendered_conversation(_messages_at(messages=messages, indices=indices)) instruction = ( - f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately {budget} tokens. " - "Prioritize completeness within that limit so the response is not cut off." + f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately " + f"{summary_tokens} tokens. Prioritize completeness within that limit so the response is not cut off." ) return [ ChatMessage.from_system(text=instruction), ChatMessage.from_user(text=f"\n{transcript}\n"), ] - def _apply_result( - self, + @staticmethod + def _apply_summary( messages: list[ChatMessage], indices: list[int], source: str, result: dict[str, Any], token_counter: TokenCounter, ) -> list[ChatMessage]: - """Validate a generator reply and replace its source messages when the result is smaller.""" + """ + Swap the selected messages for the generated summary. + + :raises RuntimeError: If the generator returned no usable text, or if the swap did not make the conversation + smaller, in which case keeping the raw messages is the better outcome. + """ replies = result.get("replies") or [] text = replies[-1].text if replies else None if not text or not text.strip(): @@ -261,283 +363,21 @@ def _apply_result( ) return compacted - def _attempt( - self, - messages: list[ChatMessage], - indices: list[int], - source: str, - budget: int, - token_counter: TokenCounter, - generate: Callable[[list[ChatMessage]], dict[str, Any]], - ) -> list[ChatMessage] | None: - """Attempt one synchronous summary, applying the configured failure policy.""" - try: - result = generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) - return self._apply_result( - messages=messages, indices=indices, source=source, result=result, token_counter=token_counter - ) - except Exception as error: - if self.raise_on_failure: - raise - logger.warning( - "Summarizing the conversation for context compaction failed; keeping the last successful result. " - "Error: {error}", - error=error, - ) - return None - - async def _attempt_async( - self, - messages: list[ChatMessage], - indices: list[int], - source: str, - budget: int, - token_counter: TokenCounter, - generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], - ) -> list[ChatMessage] | None: - """Attempt one asynchronous summary, applying the configured failure policy.""" - try: - result = await generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) - return self._apply_result( - messages=messages, indices=indices, source=source, result=result, token_counter=token_counter - ) - except Exception as error: - if self.raise_on_failure: - raise - logger.warning( - "Summarizing the conversation for context compaction failed; keeping the last successful result. " - "Error: {error}", - error=error, - ) - return None - - def _compact( - self, - original: list[ChatMessage], - target_tokens: int, - token_counter: TokenCounter, - budget: int, - generate: Callable[[list[ChatMessage]], dict[str, Any]], - ) -> list[ChatMessage] | None: - """Run synchronous historical, consolidation, and current-step compaction tiers in order.""" - working = list(original) - - # First replace the fewest oldest raw historical turns expected to reach the target. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - groups = [ - group - for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) - if group - ] - if not groups: - break - selected = _groups_to_summarize( - messages=working, - groups=groups, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = self._attempt( - messages=working, - indices=selected, - source="historical_turns", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # If all historical turns are summaries and the target is still unmet, fold them into one summary. - if token_counter.count(messages=working) > target_tokens: - summaries = self._summary_indices(messages=working, source="historical_summaries") - if len(summaries) > 1: - compacted = self._attempt( - messages=working, - indices=summaries, - source="historical_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - if not eligible: - break - - # Fold accumulated current-task summaries before consuming more raw agent steps. - summaries = self._summary_indices(messages=working, source="current_task_summaries") - if len(summaries) > 1: - compacted = self._attempt( - messages=working, - indices=summaries, - source="current_task_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - selected = _groups_to_summarize( - messages=working, - groups=eligible, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = self._attempt( - messages=working, - indices=selected, - source="current_task_steps", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - break - working = compacted - return self._result(original=original, working=working, token_counter=token_counter) - - async def _compact_async( - self, - original: list[ChatMessage], - target_tokens: int, - token_counter: TokenCounter, - budget: int, - generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], - ) -> list[ChatMessage] | None: - """Run asynchronous historical, consolidation, and current-step compaction tiers in order.""" - working = list(original) - - # First replace the fewest oldest raw historical turns expected to reach the target. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - groups = [ - group - for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) - if group - ] - if not groups: - break - selected = _groups_to_summarize( - messages=working, - groups=groups, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = await self._attempt_async( - messages=working, - indices=selected, - source="historical_turns", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # If all historical turns are summaries and the target is still unmet, fold them into one summary. - if token_counter.count(messages=working) > target_tokens: - summaries = self._summary_indices(messages=working, source="historical_summaries") - if len(summaries) > 1: - compacted = await self._attempt_async( - messages=working, - indices=summaries, - source="historical_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - if not eligible: - break - - # Fold accumulated current-task summaries before consuming more raw agent steps. - summaries = self._summary_indices(messages=working, source="current_task_summaries") - if len(summaries) > 1: - compacted = await self._attempt_async( - messages=working, - indices=summaries, - source="current_task_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - selected = _groups_to_summarize( - messages=working, - groups=eligible, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = await self._attempt_async( - messages=working, - indices=selected, - source="current_task_steps", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - break - working = compacted - return self._result(original=original, working=working, token_counter=token_counter) - - def _summary_indices(self, messages: list[ChatMessage], source: str) -> list[int]: - """Return historical or current-task summary indices based on their conversation position.""" - system_end = _leading_system_end(messages=messages) - task_index = _latest_user_index(messages=messages) - end = task_index if task_index is not None else len(messages) - if source == "historical_summaries": - return _indices(messages=messages, start=system_end, end=end, summaries=True) - start = task_index + 1 if task_index is not None else system_end - return _indices(messages=messages, start=start, end=len(messages), summaries=True) + def _report_failure(self, error: Exception) -> None: + """Re-raise a failed summarization or log it, so whatever compacted successfully so far is still returned.""" + if self.raise_on_failure: + raise error + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) @staticmethod - def _result( + def _reduced( original: list[ChatMessage], working: list[ChatMessage], token_counter: TokenCounter ) -> list[ChatMessage] | None: - """Return partial or complete progress only when it reduced the original conversation.""" + """Return partial or complete progress only when it made the original conversation smaller.""" if token_counter.count(messages=working) < token_counter.count(messages=original): return working return None diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 99b0fe79d8e..590ddb84b2d 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -27,6 +27,11 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: return None +def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages at the given indices, in the order the indices are given.""" + return [messages[index] for index in indices] + + def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: """Return the messages the given indices leave out, in conversation order.""" left_out = set(indices) diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py index 62064ffcd6d..c9d916fa56d 100644 --- a/test/hooks/compaction/test_summarization.py +++ b/test/hooks/compaction/test_summarization.py @@ -2,8 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from typing import Any - import pytest from haystack.components.generators.chat import MockChatGenerator @@ -13,140 +11,175 @@ from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") + +# A target of one token forces every tier to run, isolating the structural rules from sizing. +SMALLEST = 1 +# One character per token, so the padded messages below are obviously the expensive ones. COUNTER = FakeCounter(chars_per_token=1) -def recording_generator(responses: list[str | Exception]) -> tuple[MockChatGenerator, list[dict[str, Any]]]: - """Build a MockChatGenerator whose response function records prompts and can raise queued failures.""" +def summarizer(*responses: str | Exception) -> tuple[MockChatGenerator, list[str]]: + """ + A Chat Generator returning the given summaries in order, recording the prompt it received for each. + + An `Exception` among the responses is raised instead of answering, so a test can fail one summarization step. + """ queued = list(responses) - calls: list[dict[str, Any]] = [] + prompts: list[str] = [] def respond(messages: list[ChatMessage]) -> str: - calls.append({"messages": messages}) + prompts.append("\n".join(message.text or "" for message in messages)) response = queued.pop(0) if isinstance(response, Exception): raise response return response - return MockChatGenerator(response_fn=respond), calls + return MockChatGenerator(response_fn=respond), prompts def summary(text: str, source: str) -> ChatMessage: + """A summary an earlier compaction left behind, marked the way this compactor marks its own.""" return ChatMessage.from_user( f"\n{text}\n", meta={_COMPACTION_META_KEY: {"strategy": "summarization", "source": source}}, ) -def transcript(call: dict[str, Any]) -> str: - return call["messages"][-1].text +def sources(messages: list[ChatMessage]) -> list[str]: + """Which stretch of conversation each summary in `messages` stands in for, oldest first.""" + return [ + message.meta[_COMPACTION_META_KEY]["source"] for message in messages if _COMPACTION_META_KEY in message.meta + ] + + +def two_turns_and_a_task() -> list[ChatMessage]: + """A padded oldest turn, a short recent turn, and the current task with one step behind it.""" + return [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + + +def a_task_with_two_steps() -> list[ChatMessage]: + """The current task with a padded oldest step and a cheap newest one, and no history in front of it.""" + return [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] class TestSummarizationCompactor: - def test_summarizes_the_minimum_number_of_oldest_historical_turns(self): - messages = [ - ChatMessage.from_system("rules"), - ChatMessage.from_user("oldest question " * 30), - ChatMessage.from_assistant("oldest answer " * 30), - ChatMessage.from_user("recent question"), - ChatMessage.from_assistant("recent answer"), - ChatMessage.from_user("current task"), - ChatMessage.from_assistant("current step"), - ] - generator, calls = recording_generator(["short historical summary"]) - retained_without_oldest = [messages[0], *messages[3:]] - target = COUNTER.count(retained_without_oldest) + 100 + def test_summarizes_oldest_historical_turn(self): + messages = two_turns_and_a_task() + generator, prompts = summarizer("short historical summary") + # Room for everything but the padded oldest turn, plus the summary standing in for it. + target_tokens = COUNTER.count([messages[0], *messages[3:]]) + 100 + compacted = SummarizationCompactor(generator, max_summary_tokens=100).compact( - messages=messages, target_tokens=target, token_counter=COUNTER + messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 1 - assert "oldest question" in transcript(calls[0]) - assert "recent question" not in transcript(calls[0]) - assert compacted[0] == messages[0] - assert compacted[2:] == messages[3:] - assert messages == [ - ChatMessage.from_system("rules"), - ChatMessage.from_user("oldest question " * 30), - ChatMessage.from_assistant("oldest answer " * 30), - ChatMessage.from_user("recent question"), - ChatMessage.from_assistant("recent answer"), - ChatMessage.from_user("current task"), - ChatMessage.from_assistant("current step"), - ] + # Only the oldest turn was summarized, so the recent turn never reached the generator. + assert len(prompts) == 1 + assert "oldest question" in prompts[0] + assert "recent question" not in prompts[0] + assert compacted == [messages[0], compacted[1], *messages[3:]] + + def test_leaves_the_input_conversation_untouched(self): + messages = two_turns_and_a_task() + SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + assert messages == two_turns_and_a_task() - def test_reaches_current_steps_in_the_same_call_after_historical_context(self): + def test_summarizes_historical_turns_then_current_steps(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old question " * 40), ChatMessage.from_assistant("old answer " * 40), - ChatMessage.from_user("current task"), - tool_call("old"), - tool_result("old result " * 40, call_id="old"), - tool_call("new"), - tool_result("new result", call_id="new"), + *a_task_with_two_steps()[1:], ] - generator, calls = recording_generator(["history", "old step"]) + generator, prompts = summarizer("history", "old step") + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert "old question" in transcript(calls[0]) - assert "old result" in transcript(calls[1]) + assert "old question" in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["historical_turns", "current_task_steps"] + # The newest step is never given up. assert compacted[-2:] == messages[-2:] - assert [m.meta[_COMPACTION_META_KEY]["source"] for m in compacted if _COMPACTION_META_KEY in m.meta] == [ - "historical_turns", - "current_task_steps", - ] - def test_consolidates_historical_summaries_before_current_steps(self): + def test_folds_historical_summaries_before_current_steps(self): messages = [ ChatMessage.from_system("rules"), summary("first history " * 20, "historical_turns"), summary("second history " * 20, "historical_turns"), - ChatMessage.from_user("current task"), - tool_call("old"), - tool_result("old result " * 30, call_id="old"), - tool_call("new"), - tool_result("new result", call_id="new"), + *a_task_with_two_steps()[1:], ] - generator, calls = recording_generator(["combined history", "old step"]) + generator, prompts = summarizer("combined history", "old step") + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert "first history" in transcript(calls[0]) - assert "old result" not in transcript(calls[0]) - assert "old result" in transcript(calls[1]) + # The two historical summaries are folded into one before any current-task step is touched. + assert "first history" in prompts[0] and "old result" not in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["historical_summaries", "current_task_steps"] - def test_consolidates_current_summaries_before_more_raw_steps(self): + def test_folds_current_task_summaries_before_more_steps(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("current task"), summary("first step summary " * 20, "current_task_steps"), summary("second step summary " * 20, "current_task_steps"), - tool_call("old"), - tool_result("old result " * 30, call_id="old"), - tool_call("new"), - tool_result("new result", call_id="new"), + *a_task_with_two_steps()[2:], ] - generator, calls = recording_generator(["combined steps", "old step"]) + generator, prompts = summarizer("combined steps", "old step") + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert "first step summary" in transcript(calls[0]) - assert "old result" in transcript(calls[1]) + assert "first step summary" in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["current_task_summaries", "current_task_steps"] assert compacted[-2:] == messages[-2:] - def test_returns_partial_progress_when_a_later_tier_fails_by_default(self): + @pytest.mark.parametrize(("min_keep_steps", "expected"), [(0, 0), (1, 1), (2, 2), (20, 2)]) + def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, expected): + messages = a_task_with_two_steps() + generator, _ = summarizer("step summary") + compacted = SummarizationCompactor(generator, min_keep_steps=min_keep_steps, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + result = compacted or messages + assert sum(message.is_from(role=ChatRole.ASSISTANT) for message in result) == expected + + def test_custom_summary_instruction_replaces_the_default(self): + generator, prompts = summarizer("summary") + SummarizationCompactor(generator, summary_instruction="Only list file paths.", max_summary_tokens=1).compact( + messages=two_turns_and_a_task(), target_tokens=SMALLEST, token_counter=COUNTER + ) + assert "Only list file paths." in prompts[0] + assert "You are compacting part of a conversation" not in prompts[0] + + def test_keeps_partial_progress_when_a_summary_fails(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old question " * 30), @@ -155,39 +188,41 @@ def test_returns_partial_progress_when_a_later_tier_fails_by_default(self): ChatMessage.from_assistant("old step " * 30), ChatMessage.from_assistant("new step"), ] - generator, calls = recording_generator(["history", RuntimeError("provider unavailable")]) + generator, prompts = summarizer("history", RuntimeError("provider unavailable")) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert any( - message.meta.get(_COMPACTION_META_KEY, {}).get("source") == "historical_turns" for message in compacted - ) - assert messages[-2:] == compacted[-2:] + assert len(prompts) == 2 + # The history was summarized before the step summary failed, and that progress is kept. + assert sources(compacted) == ["historical_turns"] + assert compacted[-2:] == messages[-2:] - def test_raises_when_configured_and_summary_does_not_shrink_context(self): + def test_raises_when_a_summary_does_not_shrink_the_conversation(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old"), ChatMessage.from_assistant("answer"), ChatMessage.from_user("current"), ] - generator, _ = recording_generator(["much longer summary " * 100]) + compactor = SummarizationCompactor( + MockChatGenerator("much longer summary " * 100), max_summary_tokens=1, raise_on_failure=True + ) with pytest.raises(RuntimeError, match="did not reduce"): - SummarizationCompactor(generator, max_summary_tokens=1, raise_on_failure=True).compact(messages, 1, COUNTER) + compactor.compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) - def test_returns_none_without_calling_generator_when_context_fits(self): - generator, calls = recording_generator(["unused"]) + def test_returns_none_when_the_conversation_fits(self): + generator, prompts = summarizer("unused") messages = [ChatMessage.from_system("rules"), ChatMessage.from_user("task")] - assert ( - SummarizationCompactor(generator).compact(messages=messages, target_tokens=10_000, token_counter=COUNTER) - is None + compacted = SummarizationCompactor(generator).compact( + messages=messages, target_tokens=10_000, token_counter=COUNTER ) - assert calls == [] + assert compacted is None + assert prompts == [] - def test_summary_is_a_user_message_with_compaction_metadata(self): + def test_summary_is_a_marked_user_message(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old " * 100), @@ -195,12 +230,11 @@ def test_summary_is_a_user_message_with_compaction_metadata(self): ChatMessage.from_user("task"), ] compacted = SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - generated = compacted[1] - assert generated.is_from(ChatRole.USER) - assert generated.meta[_COMPACTION_META_KEY] == { + assert compacted[1].is_from(role=ChatRole.USER) + assert compacted[1].meta[_COMPACTION_META_KEY] == { "strategy": "summarization", "summarized_messages": 2, "source": "historical_turns", @@ -233,17 +267,17 @@ def test_serde_round_trip(self): assert restored.raise_on_failure is True -@pytest.mark.asyncio -async def test_async_compaction_uses_async_generator(): - messages = [ - ChatMessage.from_system("rules"), - ChatMessage.from_user("old " * 100), - ChatMessage.from_assistant("answer " * 100), - ChatMessage.from_user("task"), - ] - generator, calls = recording_generator(["async summary"]) - compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( - messages=messages, target_tokens=1, token_counter=COUNTER - ) - assert compacted is not None - assert len(calls) == 1 +class TestSummarizationCompactorAsync: + @pytest.mark.asyncio + async def test_compact_async_matches_compact(self): + messages = two_turns_and_a_task() + generator, prompts = summarizer("async summary") + + compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert len(prompts) == 1 + assert compacted == SummarizationCompactor(MockChatGenerator("async summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) From efd1d0df85c220b34a4c22dbe8118bffaed07f88 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:57:22 +0200 Subject: [PATCH 16/38] more refactoring --- haystack/hooks/compaction/summarization.py | 106 +++++++++++++-------- 1 file changed, 68 insertions(+), 38 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index 2fa2add59a9..7e2ed9053ee 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -55,18 +55,22 @@ def _is_summary(message: ChatMessage) -> bool: return _is_compaction_message(message=message, strategy=_STRATEGY) -def _summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: - """Return the positions of this strategy's summaries within a bounded part of a conversation.""" +def _previous_summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: + """Return the positions of the summaries an earlier compaction left in a bounded part of a conversation.""" return [index for index in range(start, end) if _is_summary(message=messages[index])] -def _summarizable_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: +def _raw_historical_turn_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> list[list[int]]: """ - Return the historical turns that still hold raw conversation, oldest turn first. + Return the historical turns that still hold raw, never-summarized conversation, oldest turn first. - Summaries this strategy already wrote are left out of their turn, so summarizing the turn folds that summary into - the summary this run produces. A turn that is nothing but summaries has nothing left to give up and is dropped. + Summaries an earlier compaction wrote are excluded, so summarizing a turn leaves them in place for + `_HISTORICAL_SUMMARIES` to fold later. The list is empty when there are no historical turns, or when every one of + them is already nothing but summaries. """ + # Strip the previous summaries out of each turn, then drop the turns that strip away to nothing. groups = [ [index for index in group if not _is_summary(message=messages[index])] for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) @@ -106,9 +110,11 @@ def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMe def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: """Replace the selected messages, which need not be contiguous, with one summary at the oldest one's position.""" selected = set(indices) + # The summary stands in for everything it replaced, so it takes the position of the oldest message it covers. insertion_index = min(indices) compacted: list[ChatMessage] = [] for index, message in enumerate(messages): + # Emit the summary before the message it displaces, so the surrounding conversation keeps its order. if index == insertion_index: compacted.append(summary) if index not in selected: @@ -191,11 +197,14 @@ def compact( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ + # How large each summary may be, and the run kwargs, if any, that hold the generator to it. summary_tokens, run_kwargs = self._summary_limit() - working = list(messages) + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. + compacted = messages while True: + # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( - messages=working, + messages=compacted, target_tokens=target_tokens, token_counter=token_counter, summary_tokens=summary_tokens, @@ -203,16 +212,21 @@ def compact( if plan is None: break indices, source = plan - prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + prompt = self._prompt(messages=compacted, indices=indices, summary_tokens=summary_tokens) try: + # Summarize that stretch and swap it in, so the next round plans against the smaller conversation. + # A generator error or a summary that does not shrink raises out of here. result = self.chat_generator.run(messages=prompt, **run_kwargs) - working = self._apply_summary( - messages=working, indices=indices, source=source, result=result, token_counter=token_counter + compacted = self._apply_summary( + messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) except Exception as error: + # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - return self._reduced(original=messages, working=working, token_counter=token_counter) + # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than + # the untouched input means real progress, whether or not the target was met. + return None if compacted is messages else compacted async def compact_async( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter @@ -225,11 +239,14 @@ async def compact_async( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ + # How large each summary may be, and the run kwargs, if any, that hold the generator to it. summary_tokens, run_kwargs = self._summary_limit() - working = list(messages) + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. + compacted = messages while True: + # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( - messages=working, + messages=compacted, target_tokens=target_tokens, token_counter=token_counter, summary_tokens=summary_tokens, @@ -237,18 +254,23 @@ async def compact_async( if plan is None: break indices, source = plan - prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + prompt = self._prompt(messages=compacted, indices=indices, summary_tokens=summary_tokens) try: + # Summarize that stretch and swap it in, so the next round plans against the smaller conversation. + # Only the generator call is awaited; planning and swapping are pure. result = await _execute_component_async( component_instance=self.chat_generator, messages=prompt, **run_kwargs ) - working = self._apply_summary( - messages=working, indices=indices, source=source, result=result, token_counter=token_counter + compacted = self._apply_summary( + messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) except Exception as error: + # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - return self._reduced(original=messages, working=working, token_counter=token_counter) + # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than + # the untouched input means real progress, whether or not the target was met. + return None if compacted is messages else compacted def _next_summary( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int @@ -271,6 +293,7 @@ def _next_summary( :returns: The message indices to summarize and the `source` to record on the resulting summary, or None when the conversation already fits or nothing is left that may be given up. """ + # Nothing to give up once the conversation fits. if token_counter.count(messages=messages) <= target_tokens: return None @@ -281,34 +304,39 @@ def _next_summary( history_end = task_index if task_index is not None else system_end task_start = task_index + 1 if task_index is not None else system_end - turns = _summarizable_turn_groups(messages=messages, system_end=system_end, task_index=task_index) - if turns: + # Tier 1. Raw history is the cheapest context to lose, so take the oldest turns that still hold any. + historical_turns = _raw_historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + if historical_turns: oldest_turns = _groups_to_summarize( messages=messages, - groups=turns, + groups=historical_turns, target_tokens=target_tokens, summary_tokens=summary_tokens, token_counter=token_counter, ) return oldest_turns, _HISTORICAL_TURNS - history_summaries = _summary_indices(messages=messages, start=system_end, end=history_end) + # Tier 2. History is nothing but summaries now, so the only room left there is in folding them into one. They + # are left to accumulate until this point so that they are not rewritten on every compaction. + history_summaries = _previous_summary_indices(messages=messages, start=system_end, end=history_end) if len(history_summaries) > 1: return history_summaries, _HISTORICAL_SUMMARIES - # Only steps older than the `min_keep_steps` most recent ones may be given up. - steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) - eligible = steps[: max(len(steps) - self.min_keep_steps, 0)] - if not eligible: + # History is exhausted, so the current task has to pay. Its `min_keep_steps` newest steps are off limits. + agent_steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) + eligible_steps = agent_steps[: max(len(agent_steps) - self.min_keep_steps, 0)] + if not eligible_steps: return None - task_summaries = _summary_indices(messages=messages, start=task_start, end=len(messages)) + # Tier 3. Fold the summaries earlier steps left behind before spending another raw step on the same space. + task_summaries = _previous_summary_indices(messages=messages, start=task_start, end=len(messages)) if len(task_summaries) > 1: return task_summaries, _CURRENT_TASK_SUMMARIES + # Tier 4. Last resort: give up the oldest steps of the task the Agent is working on right now. oldest_steps = _groups_to_summarize( messages=messages, - groups=eligible, + groups=eligible_steps, target_tokens=target_tokens, summary_tokens=summary_tokens, token_counter=token_counter, @@ -316,7 +344,18 @@ def _next_summary( return oldest_steps, _CURRENT_TASK_STEPS def _summary_limit(self) -> tuple[int, dict[str, Any]]: - """Return the token budget for one summary and the run kwargs, if any, that ask the generator to honor it.""" + """ + Work out how large one summary may be and how to hold the Chat Generator to it. + + :returns: A tuple containing: + + 1. The token budget for a single summary. This is `max_summary_tokens`, unless the generator already + configures a recognized output limit of its own, in which case the generator's setting wins. + 2. The kwargs to pass to the generator's `run`. This carries a `generation_kwargs` entry for a built-in + generator that has no limit configured, and is empty for every other generator, since the + `ChatGenerator` protocol does not standardize the setting. When it is empty, the budget reaches the + model only as prompt guidance and `_apply_summary` measures the result instead. + """ summary_tokens, generation_kwargs = _resolve_output_token_limit( chat_generator=self.chat_generator, default_limit=self.max_summary_tokens ) @@ -373,15 +412,6 @@ def _report_failure(self, error: Exception) -> None: error=error, ) - @staticmethod - def _reduced( - original: list[ChatMessage], working: list[ChatMessage], token_counter: TokenCounter - ) -> list[ChatMessage] | None: - """Return partial or complete progress only when it made the original conversation smaller.""" - if token_counter.count(messages=working) < token_counter.count(messages=original): - return working - return None - def warm_up(self) -> None: """Warm up the Chat Generator that writes summaries.""" if hasattr(self.chat_generator, "warm_up"): From b5ca520bfd349b78dc6df01f528847860b7838cd Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:59:04 +0200 Subject: [PATCH 17/38] adding more clarity --- haystack/hooks/compaction/summarization.py | 52 ++++++++++------------ 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index 7e2ed9053ee..da23d3b2d7f 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -197,10 +197,16 @@ def compact( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ - # How large each summary may be, and the run kwargs, if any, that hold the generator to it. - summary_tokens, run_kwargs = self._summary_limit() + # How large each summary may be. A recognized built-in generator is held to it by its own runtime setting; + # any other generator gets it as prompt guidance only, since the protocol guarantees nothing beyond `run`. + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + run_kwargs = {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. compacted = messages + summarized = False while True: # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( @@ -220,13 +226,14 @@ def compact( compacted = self._apply_summary( messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) + summarized = True except Exception as error: # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than - # the untouched input means real progress, whether or not the target was met. - return None if compacted is messages else compacted + # Every applied summary was measured as shrinking the conversation, so any summary at all is real progress, + # whether or not the target was met. Without one there is nothing to hand back. + return compacted if summarized else None async def compact_async( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter @@ -239,10 +246,16 @@ async def compact_async( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ - # How large each summary may be, and the run kwargs, if any, that hold the generator to it. - summary_tokens, run_kwargs = self._summary_limit() + # How large each summary may be. A recognized built-in generator is held to it by its own runtime setting; + # any other generator gets it as prompt guidance only, since the protocol guarantees nothing beyond `run`. + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + run_kwargs = {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. compacted = messages + summarized = False while True: # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( @@ -264,13 +277,14 @@ async def compact_async( compacted = self._apply_summary( messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) + summarized = True except Exception as error: # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than - # the untouched input means real progress, whether or not the target was met. - return None if compacted is messages else compacted + # Every applied summary was measured as shrinking the conversation, so any summary at all is real progress, + # whether or not the target was met. Without one there is nothing to hand back. + return compacted if summarized else None def _next_summary( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int @@ -343,24 +357,6 @@ def _next_summary( ) return oldest_steps, _CURRENT_TASK_STEPS - def _summary_limit(self) -> tuple[int, dict[str, Any]]: - """ - Work out how large one summary may be and how to hold the Chat Generator to it. - - :returns: A tuple containing: - - 1. The token budget for a single summary. This is `max_summary_tokens`, unless the generator already - configures a recognized output limit of its own, in which case the generator's setting wins. - 2. The kwargs to pass to the generator's `run`. This carries a `generation_kwargs` entry for a built-in - generator that has no limit configured, and is empty for every other generator, since the - `ChatGenerator` protocol does not standardize the setting. When it is empty, the budget reaches the - model only as prompt guidance and `_apply_summary` measures the result instead. - """ - summary_tokens, generation_kwargs = _resolve_output_token_limit( - chat_generator=self.chat_generator, default_limit=self.max_summary_tokens - ) - return summary_tokens, {"generation_kwargs": generation_kwargs} if generation_kwargs else {} - def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" transcript = _rendered_conversation(_messages_at(messages=messages, indices=indices)) From fbf90d6e89069c09a2da03cf015f09bd9c5fa8d9 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 09:25:07 +0200 Subject: [PATCH 18/38] Add better placeholders for images and files --- haystack/hooks/compaction/summarization.py | 40 +++++++++++++-- haystack/token_counters/utils.py | 25 ++++++--- test/hooks/compaction/test_summarization.py | 57 ++++++++++++++++++++- test/token_counters/test_utils.py | 11 ++++ 4 files changed, 120 insertions(+), 13 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index da23d3b2d7f..387146deb89 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -8,7 +8,8 @@ from haystack.components.generators.chat.types import ChatGenerator from haystack.components.generators.chat.utils import _resolve_output_token_limit from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict -from haystack.dataclasses import ChatMessage +from haystack.dataclasses import ChatMessage, FileContent, ImageContent +from haystack.dataclasses.chat_message import ChatMessageContentT from haystack.hooks.compaction.types import Compactor from haystack.hooks.compaction.utils import ( _COMPACTION_META_KEY, @@ -46,10 +47,34 @@ - Exact file paths, URLs, identifiers, and references to stored data. - Unresolved work and the immediate next step. +Images and files appear only as and placeholders; their contents are not available to you \ +and will be lost. Keep the names and details shown in the placeholder, along with whatever the conversation says \ +about them, so they can be supplied again if they are needed. + Fold any existing blocks into one summary. Record only what the conversation shows. Do not \ infer or add advice. Use plain prose or short bullets, and do not address the user.""" +def _identifying_details(metadata: dict[str, Any]) -> list[str]: + """Render an attachment's metadata as `key=value` pairs, such as the path a file was loaded from.""" + # For ease, we don't support nested keys, we are mostly interested in the top-level keys that identify the + # attachment, such as a file path or URL. + return [f"{key}={value}" for key, value in sorted(metadata.items()) if isinstance(value, (str, int, float, bool))] + + +def _attachment_placeholder(content: ChatMessageContentT) -> str: + """ + Render a placeholder for an attachment that cannot survive summarization, so the summary can preserve its identity. + """ + if isinstance(content, ImageContent): + # Images have no filename, so whatever identifies one lives in its `meta`. + return f"" + if isinstance(content, FileContent): + details = [content.filename or "unnamed", content.mime_type or "unknown type"] + return f"" + return f"<{type(content).__name__}>" + + def _is_summary(message: ChatMessage) -> bool: """Whether a message is a summary this strategy wrote.""" return _is_compaction_message(message=message, strategy=_STRATEGY) @@ -160,7 +185,7 @@ def __init__( *, min_keep_steps: int = 1, max_summary_tokens: int = 1024, - summary_instruction: str | None = None, + summary_instruction: str = _DEFAULT_SUMMARY_INSTRUCTION, raise_on_failure: bool = False, ) -> None: """ @@ -171,7 +196,10 @@ def __init__( :param min_keep_steps: The fewest complete recent Agent steps to keep, even when they exceed the target. :param max_summary_tokens: The output-token budget reserved for each summary. Known built-in generators receive the corresponding runtime generation setting unless one is already configured on the generator. - :param summary_instruction: An instruction replacing the built-in summary prompt. + :param summary_instruction: What the model is told to preserve when it writes a summary. The default asks for + the user's goal, decisions and their reasoning, completed work, exact identifiers, the names of attachments + that cannot survive summarization, and the next step. The token budget is appended to whatever is given + here, so a replacement does not need to mention it. :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is logged and any successful partial compaction is returned. :raises ValueError: If `min_keep_steps` is negative or `max_summary_tokens` is not positive. @@ -183,7 +211,7 @@ def __init__( self.chat_generator = chat_generator self.min_keep_steps = min_keep_steps self.max_summary_tokens = max_summary_tokens - self.summary_instruction = summary_instruction or _DEFAULT_SUMMARY_INSTRUCTION + self.summary_instruction = summary_instruction self.raise_on_failure = raise_on_failure def compact( @@ -359,7 +387,9 @@ def _next_summary( def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" - transcript = _rendered_conversation(_messages_at(messages=messages, indices=indices)) + transcript = _rendered_conversation( + _messages_at(messages=messages, indices=indices), placeholder=_attachment_placeholder + ) instruction = ( f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately " f"{summary_tokens} tokens. Prioritize completeness within that limit so the response is not cut off." diff --git a/haystack/token_counters/utils.py b/haystack/token_counters/utils.py index 9dbf9b77940..7f4d8f3ef71 100644 --- a/haystack/token_counters/utils.py +++ b/haystack/token_counters/utils.py @@ -3,11 +3,15 @@ # SPDX-License-Identifier: Apache-2.0 import json +from collections.abc import Callable from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent from haystack.dataclasses.chat_message import ChatMessageContentT, ToolCallResultContentT from haystack.tools import ToolsType, flatten_tools_or_toolsets +# Builds the stand-in for message content that has no text form, such as an image. +_PlaceholderFn = Callable[[ChatMessageContentT], str] + def _non_text_placeholder(content: ChatMessageContentT) -> str: """A short stand-in, such as ``, for message content that has no text form.""" @@ -18,26 +22,33 @@ def _non_text_placeholder(content: ChatMessageContentT) -> str: return f"<{type(content).__name__}>" -def _tool_result_text(result: ToolCallResultContentT) -> str: +def _tool_result_text(result: ToolCallResultContentT, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """A tool result as a single string, with placeholders standing in for any non-text parts.""" if isinstance(result, str): return result - return "".join(block.text if isinstance(block, TextContent) else _non_text_placeholder(block) for block in result) + return "".join(block.text if isinstance(block, TextContent) else placeholder(block) for block in result) -def _render_message(message: ChatMessage) -> str: +def _render_message(message: ChatMessage, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """ One message as one or more lines of plain text. Reasoning content is deliberately left out: providers discard it between turns, so it is not part of the context being measured. + + :param message: The message to render. + :param placeholder: Builds the stand-in for content that has no text form. The default is short and stable, which + is what a counter needs because the stand-in's own length is what gets measured. A caller rendering for a model + to read can pass one that describes the content instead. + :returns: The rendered message. """ role = message.role.value # A tool-result msg only carries tool_call_results, so it is rendered on its own and labelled with the tool that # produced it. if results := message.tool_call_results: return "\n".join( - f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] {_tool_result_text(result.result)}" + f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] " + f"{_tool_result_text(result.result, placeholder=placeholder)}" for result in results ) @@ -50,13 +61,13 @@ def _render_message(message: ChatMessage) -> str: # Images and files cost tokens too, so they need a stand-in rather than being skipped. non_text: list[ChatMessageContentT] = [*message.images, *message.files] for content in non_text: - lines.append(f"[{role}] {_non_text_placeholder(content)}") + lines.append(f"[{role}] {placeholder(content)}") return "\n".join(lines) if lines else f"[{role}] " -def _rendered_conversation(messages: list[ChatMessage]) -> str: +def _rendered_conversation(messages: list[ChatMessage], *, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """The whole conversation as one plain-text block, which is what a counter measures.""" - return "\n".join(_render_message(message) for message in messages) + return "\n".join(_render_message(message, placeholder=placeholder) for message in messages) def _rendered_tools(tools: ToolsType | None) -> str: diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py index c9d916fa56d..b2b5197125e 100644 --- a/test/hooks/compaction/test_summarization.py +++ b/test/hooks/compaction/test_summarization.py @@ -5,8 +5,9 @@ import pytest from haystack.components.generators.chat import MockChatGenerator -from haystack.dataclasses import ChatMessage, ChatRole +from haystack.dataclasses import ChatMessage, ChatRole, FileContent, ImageContent, TextContent, ToolCall from haystack.hooks.compaction import SummarizationCompactor +from haystack.hooks.compaction.summarization import _attachment_placeholder from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result @@ -77,6 +78,38 @@ def a_task_with_two_steps() -> list[ChatMessage]: ] +class TestAttachmentPlaceholder: + @pytest.mark.parametrize( + ("content", "expected"), + [ + pytest.param(ImageContent(base64_image="Zm9v", mime_type="image/png"), "", id="image"), + pytest.param( + ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"file_path": "/tmp/shot.png"}), + "", + id="image-named-by-meta", + ), + pytest.param( + FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="q3.pdf"), + "", + id="file", + ), + pytest.param( + FileContent(base64_data="Zm9v", mime_type="application/pdf", extra={"page": 4}), + "", + id="file-unnamed-with-extra", + ), + # A nested value could be arbitrarily large, so it is left out rather than bloating the prompt. + pytest.param( + ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"boxes": [[1, 2], [3, 4]]}), + "", + id="nested-metadata-left-out", + ), + ], + ) + def test_names_the_attachment(self, content, expected): + assert _attachment_placeholder(content) == expected + + class TestSummarizationCompactor: def test_summarizes_oldest_historical_turn(self): messages = two_turns_and_a_task() @@ -171,6 +204,28 @@ def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, e result = compacted or messages assert sum(message.is_from(role=ChatRole.ASSISTANT) for message in result) == expected + def test_attachments_are_named_in_the_transcript(self): + image = ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"file_path": "/tmp/shot.png"}) + pdf = FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="q3.pdf") + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user(content_parts=["review this " * 20, pdf]), + tool_call("c1"), + # An attachment a tool returned is nested inside the tool result rather than on the message. + ChatMessage.from_tool( + tool_result=[TextContent(text="captured " * 20), image], + origin=ToolCall(tool_name="browse", arguments={}, id="c1"), + ), + ChatMessage.from_user("current task"), + ] + generator, prompts = summarizer("summary") + SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + # The summary cannot reproduce either attachment, so the transcript has to name them well enough to ask again. + assert "" in prompts[0] + assert "" in prompts[0] + def test_custom_summary_instruction_replaces_the_default(self): generator, prompts = summarizer("summary") SummarizationCompactor(generator, summary_instruction="Only list file paths.", max_summary_tokens=1).compact( diff --git a/test/token_counters/test_utils.py b/test/token_counters/test_utils.py index b3101da360b..8fb06a896a2 100644 --- a/test/token_counters/test_utils.py +++ b/test/token_counters/test_utils.py @@ -75,6 +75,17 @@ def test_joins_messages_with_newlines(self): def test_empty_conversation(self): assert _rendered_conversation([]) == "" + def test_a_custom_placeholder_reaches_nested_tool_results(self): + messages = [ + ChatMessage.from_user(content_parts=["look:", IMAGE]), + ChatMessage.from_tool( + tool_result=[TextContent(text="screenshot: "), IMAGE], + origin=ToolCall(tool_name="browse", arguments={}, id="c1"), + ), + ] + rendered = _rendered_conversation(messages, placeholder=lambda content: "") + assert rendered == "[user] look:\n[user] \n[tool:browse] screenshot: " + @tool def search(query: Annotated[str, "the search query"]) -> str: From c8a911e195f4aa458fc1ead8f7c7f0750d95120d Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 13:25:07 +0200 Subject: [PATCH 19/38] Add support for keeping old user-assistant turns --- haystack/hooks/compaction/sliding_window.py | 105 +++++++++++++----- ...t-context-compaction-3258c08dec9d2b34.yaml | 12 +- test/hooks/compaction/test_sliding_window.py | 83 +++++++++++++- 3 files changed, 169 insertions(+), 31 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 9cba0dc924c..97c7e81240e 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -44,59 +44,114 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: return None +def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: + """ + Return spans for complete user turns in a bounded section of conversation history. + + Each turn begins with a real user message and continues up to, but does not include, the next real user message. + This groups a user's request with every assistant step and tool result produced in response to it. + + :param messages: The full conversation to analyze, ordered oldest to newest. + :param start: The inclusive index at which to begin looking for historical turns. + :param end: The exclusive index at which to stop. This is normally the current task's user-message index. + :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to + `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. + """ + # Compaction notes use the user role for provider compatibility, but they do not begin a new conversation turn. + # Ignoring marked messages here also lets a subsequent compaction fold an old note into its replacement. + user_indices = [ + index + for index in range(start, end) + if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta + ] + + # A real user message closes the preceding turn and starts the next one. The final historical turn extends to the + # supplied boundary, which is typically where the protected current task begins. + return [ + (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) + for position, index in enumerate(user_indices) + ] + + def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int ) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage], int]: - """Split messages into the protected task context, removable history, and retained Agent steps.""" + """Split messages into the protected prefix, removable history, and retained conversation window.""" # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) # Find the latest user message to use as the current task anchor. task_index = _latest_user_index(messages=messages) task = [messages[task_index]] if task_index is not None else [] - # Find complete Agent steps after the current task, or after the system messages when there is no task anchor. - steps = _agent_step_spans(messages=messages, start=(task_index + 1) if task_index is not None else system_end) + step_start = (task_index + 1) if task_index is not None else system_end + # Current-task steps can be removed individually. Earlier user/assistant exchanges are kept as complete turns so + # an assistant reply is never retained without the user message it answers. + steps = _agent_step_spans(messages=messages, start=step_start) + historical_end = task_index if task_index is not None else system_end + historical_turns = _historical_turn_spans(messages, system_end, historical_end) # Protect the Agent instructions and current task from removal. protected = [*messages[:system_end], *task] - # The remaining token budget to retain recent Agent steps. + # The remaining token budget after protecting the instructions and current task. available_tokens = target_tokens - token_counter.count(messages=protected) - # Work backwards through the steps, keeping as many as fit in the remaining budget. - kept_step_start = len(steps) - while kept_step_start > 0: - start, end = steps[kept_step_start - 1] - step_tokens = token_counter.count(messages=messages[start:end]) - # Stop at the first step that does not fit - if step_tokens > available_tokens: - break - available_tokens -= step_tokens - kept_step_start -= 1 + kept_turn_start = len(historical_turns) + all_step_tokens = token_counter.count(messages=[message for start, end in steps for message in messages[start:end]]) + if all_step_tokens <= available_tokens: + # Historical turns are considered only when the entire current task fits. This ensures that compaction removes + # every older turn before it starts trimming individual steps from the task the Agent is actively working on. + kept_step_start = 0 + available_tokens -= all_step_tokens + while kept_turn_start > 0: + start, end = historical_turns[kept_turn_start - 1] + turn = [message for message in messages[start:end] if _COMPACTION_META_KEY not in message.meta] + turn_tokens = token_counter.count(messages=turn) + if turn_tokens > available_tokens: + break + available_tokens -= turn_tokens + kept_turn_start -= 1 + else: + # Even after dropping every historical turn, the current task is too large. Work backwards through its Agent + # steps and retain the most recent complete suffix that fits. + kept_step_start = len(steps) + while kept_step_start > 0: + start, end = steps[kept_step_start - 1] + step_tokens = token_counter.count(messages=messages[start:end]) + if step_tokens > available_tokens: + break + available_tokens -= step_tokens + kept_step_start -= 1 # Enforce the minimum number of complete steps, even when they exceed the target token budget. kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) - kept_spans = steps[kept_step_start:] + kept_step_spans = steps[kept_step_start:] + kept_turn_spans = historical_turns[kept_turn_start:] # Record every protected or retained message index; equal ChatMessages can appear more than once in the list. kept_indices = {*range(system_end)} if task_index is not None: kept_indices.add(task_index) - for start, end in kept_spans: - kept_indices.update(range(start, end)) + for start, end in [*kept_turn_spans, *kept_step_spans]: + kept_indices.update(index for index in range(start, end) if _COMPACTION_META_KEY not in messages[index].meta) # Everything outside the protected context and retained steps can be removed. removable = [message for index, message in enumerate(messages) if index not in kept_indices] - # Flatten the retained spans back into the message list expected by the compactor. - kept_steps = [message for start, end in kept_spans for message in messages[start:end]] - return protected, removable, kept_steps, len(kept_spans) + kept_turns = [ + message + for start, end in kept_turn_spans + for message in messages[start:end] + if _COMPACTION_META_KEY not in message.meta + ] + kept_steps = [message for start, end in kept_step_spans for message in messages[start:end]] + return [*messages[:system_end], *kept_turns, *task], removable, kept_steps, len(steps) - kept_step_start @_experimental class SlidingWindowCompactor(Compactor): """ - Keeps the Agent's instructions, current task, and as many complete recent steps as the target allows. + Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows. - Leading system messages and the latest user message are protected. Recent history is retained in complete Agent - steps, where a step is an assistant message together with all immediately following tool results. An - `omission_note` is left in place of what was removed. + Leading system messages and the latest user message are protected. Earlier user/assistant turns are retained when + they fit, and the current task's history is retained in complete Agent steps, where a step is an assistant message + together with all immediately following tool results. An `omission_note` is left in place of what was removed. ```python from haystack.components.agents import Agent @@ -134,7 +189,7 @@ def compact( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter ) -> list[ChatMessage] | None: """ - Drop older history while preserving the task anchor and a recent window of complete Agent steps. + Drop older history while preserving the task anchor and a complete recent conversation window. :param messages: The conversation to compact, oldest to newest. :param target_tokens: The size the retained conversation should come in under. diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index 8e0a4810a64..c71c803f853 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -5,8 +5,9 @@ features: shortens the conversation when it reaches a configured fraction of the model's context window. The first built-in strategy, ``SlidingWindowCompactor``, preserves leading system messages, the latest user task, - and as many complete recent Agent steps as the target allows. It replaces removed history with a short omission - note. + and as much complete recent conversation as the target allows. It removes earlier user/assistant turns as whole + units first. Only when removing all earlier turns is insufficient does it remove individual Agent steps from the + current task. It replaces removed history with a short omission note. .. code-block:: python @@ -30,9 +31,10 @@ features: tool schemas. Leave headroom above ``compact_at`` for the next reply and its tool results. ``SlidingWindowCompactor`` treats an assistant message and its following tool results as one step, so a tool call is - never separated from its results. It may retain slightly more history than the requested target when preserving the - current task or the configured minimum number of steps. Compaction is lossy: removed messages cannot be recovered - or summarized by this strategy. Implement the ``Compactor`` protocol to provide a custom strategy. + never separated from its results. Historical turns are also kept whole, so an assistant reply is not retained + without the user message it answers. It may retain slightly more history than the requested target when preserving + the current task or the configured minimum number of steps. Compaction is lossy: removed messages cannot be + recovered or summarized by this strategy. Implement the ``Compactor`` protocol to provide a custom strategy. ``CompactionHook`` and ``SlidingWindowCompactor`` emit an ``ExperimentalWarning`` and may change without a deprecation cycle. diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index d395a32f795..635f4407d73 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -6,7 +6,7 @@ from haystack.dataclasses import ChatMessage, ChatRole from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import FakeCounter, count_markers, long_conversation, tool_call, tool_result @@ -18,6 +18,46 @@ COUNTER = FakeCounter() +class TestHistoricalTurnSpans: + def test_groups_each_user_message_with_its_assistant_steps_and_tool_results(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("first task"), + tool_call("c1"), + tool_result("first result", call_id="c1"), + ChatMessage.from_assistant("first answer"), + ChatMessage.from_user("second task"), + ChatMessage.from_assistant("second answer"), + ] + spans = _historical_turn_spans(messages=messages, start=1, end=len(messages)) + assert spans == [(1, 5), (5, 7)] + assert messages[slice(*spans[0])] == messages[1:5] + assert messages[slice(*spans[1])] == messages[5:7] + + def test_only_returns_turns_within_the_requested_bounds(self): + messages = [ + ChatMessage.from_user("outside"), + ChatMessage.from_assistant("outside answer"), + ChatMessage.from_user("inside"), + ChatMessage.from_assistant("inside answer"), + ChatMessage.from_user("current task"), + ] + assert _historical_turn_spans(messages=messages, start=2, end=4) == [(2, 4)] + + def test_compaction_note_does_not_start_a_new_turn(self): + note = ChatMessage.from_user( + "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ) + messages = [ + ChatMessage.from_user("task"), + ChatMessage.from_assistant("first step"), + note, + ChatMessage.from_assistant("second step"), + ChatMessage.from_user("next task"), + ] + assert _historical_turn_spans(messages=messages, start=0, end=len(messages)) == [(0, 4), (4, 5)] + + class TestSlidingWindowCompactor: def test_replaces_the_middle_with_an_omission_note(self): messages = long_conversation() @@ -48,6 +88,47 @@ def test_a_roomier_target_keeps_more(self): assert len(tight) == 4 assert len(roomy) == 8 + def test_keeps_complete_recent_user_assistant_turns_that_fit(self): + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user(text="old question"), + ChatMessage.from_assistant(text="old answer"), + ChatMessage.from_user(text="recent question"), + ChatMessage.from_assistant(text="recent answer"), + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Drops one historical turn + expected = [messages[0], *messages[3:]] + target_tokens = COUNTER.count(expected) + compacted = SlidingWindowCompactor(omission_note=None).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER + ) + assert compacted == expected + + def test_drops_historical_context_and_one_current_task_step_to_reach_target(self): + system_message = ChatMessage.from_system(text="rules") + historical_turn = [ + ChatMessage.from_user(text="old question"), + tool_call("old-call"), + tool_result(result="old result", call_id="old-call"), + ChatMessage.from_assistant(text="old final answer"), + ] + current_task = [ + ChatMessage.from_user(text="current task"), + tool_call("current-call-1"), + tool_result(result="large intermediate result " * 100, call_id="current-call-1"), + tool_call("current-call-2"), + tool_result(result="latest result", call_id="current-call-2"), + ] + messages = [system_message, *historical_turn, *current_task] + # The target is small enough that the historical context and one step in the current task must be removed. + target_tokens = 52 + compacted = SlidingWindowCompactor(omission_note=None).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER + ) + assert compacted == [system_message, current_task[0], *current_task[-2:]] + def test_returns_none_when_the_conversation_already_fits(self): assert ( SlidingWindowCompactor().compact(messages=long_conversation(), target_tokens=100_000, token_counter=COUNTER) From bee98f4644485d4b204732ed1d12918a1a589b9e Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 13:39:38 +0200 Subject: [PATCH 20/38] refactoring --- haystack/hooks/compaction/sliding_window.py | 179 ++++++++++++++----- test/hooks/compaction/test_sliding_window.py | 14 ++ 2 files changed, 148 insertions(+), 45 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 97c7e81240e..f3ed181bacb 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -73,10 +73,120 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> ] +def _messages_from_spans( + messages: list[ChatMessage], spans: list[tuple[int, int]], *, skip_compaction_notes: bool = False +) -> list[ChatMessage]: + """Flatten message spans, optionally excluding messages produced by an earlier compaction.""" + return [ + message + for start, end in spans + for message in messages[start:end] + if not skip_compaction_notes or _COMPACTION_META_KEY not in message.meta + ] + + +def _fitting_suffix_start( + messages: list[ChatMessage], + spans: list[tuple[int, int]], + available_tokens: int, + token_counter: TokenCounter, + *, + skip_compaction_notes: bool = False, +) -> int: + """ + Return the first span in the newest contiguous suffix that fits the token budget. + + Spans are measured from newest to oldest. Retention stops as soon as a span does not fit, ensuring that an older + span is never kept after a newer one has been removed. + + :param messages: The full conversation containing the messages referenced by `spans`. + :param spans: Ordered `(start_index, end_index)` pairs to consider for retention. Both indices refer to `messages`, + and `end_index` is exclusive. + :param available_tokens: The token budget available for retaining messages from `spans`. + :param token_counter: The `TokenCounter` used to measure each span. + :param skip_compaction_notes: Whether messages produced by an earlier compaction are excluded from token counting. + :returns: The index in `spans` at which the retained suffix begins. If no span fits, returns `len(spans)`; if every + span fits, returns `0`. + """ + kept_start = len(spans) + while kept_start > 0: + span_messages = _messages_from_spans( + messages=messages, spans=[spans[kept_start - 1]], skip_compaction_notes=skip_compaction_notes + ) + span_tokens = token_counter.count(messages=span_messages) + if span_tokens > available_tokens: + break + available_tokens -= span_tokens + kept_start -= 1 + return kept_start + + +def _retained_span_starts( + messages: list[ChatMessage], + protected: list[ChatMessage], + historical_turns: list[tuple[int, int]], + steps: list[tuple[int, int]], + target_tokens: int, + token_counter: TokenCounter, + min_keep_steps: int, +) -> tuple[int, int]: + """Return the first retained historical turn and current-task step.""" + available_tokens = target_tokens - token_counter.count(messages=protected) + all_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=steps, skip_compaction_notes=False) + ) + kept_turn_start = len(historical_turns) + if all_step_tokens <= available_tokens: + # Historical turns are considered only when the entire current task fits. This ensures that compaction removes + # every older turn before it starts trimming individual steps from the task the Agent is actively working on. + kept_step_start = 0 + kept_turn_start = _fitting_suffix_start( + messages=messages, + spans=historical_turns, + available_tokens=available_tokens - all_step_tokens, + token_counter=token_counter, + skip_compaction_notes=True, + ) + else: + # Even after dropping every historical turn, the current task is too large. Retain the most recent complete + # suffix of Agent steps that fits. + kept_step_start = _fitting_suffix_start( + messages=messages, + spans=steps, + available_tokens=available_tokens, + token_counter=token_counter, + skip_compaction_notes=False, + ) + kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) + return kept_turn_start, kept_step_start + + def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int ) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage], int]: - """Split messages into the protected prefix, removable history, and retained conversation window.""" + """ + Split a conversation into the messages kept before and after an omission note and the messages to remove. + + Leading system messages and the latest real user message are always retained. Historical user turns are retained + whole when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed + individually while preserving complete assistant/tool-result groups. + + :param messages: The full conversation to split, ordered oldest to newest. + :param target_tokens: The target token budget for the retained messages. + :param token_counter: The `TokenCounter` used to decide which historical turns and current-task steps fit. + :param min_keep_steps: The minimum number of recent current-task Agent steps to retain, even if they exceed the + target token budget. + :returns: A tuple containing: + + 1. Messages retained before the omission-note position. + 2. Every message selected for removal. + 3. Messages retained after the omission-note position. + 4. The number of retained current-task Agent steps. + + When only historical turns are removed, the first and third elements place the note immediately before the + current task. When current-task steps are also removed, they place it after the latest user message and before + the retained current-task steps. + """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) # Find the latest user message to use as the current task anchor. @@ -87,41 +197,19 @@ def _task_and_step_split( # an assistant reply is never retained without the user message it answers. steps = _agent_step_spans(messages=messages, start=step_start) historical_end = task_index if task_index is not None else system_end - historical_turns = _historical_turn_spans(messages, system_end, historical_end) + historical_turns = _historical_turn_spans(messages=messages, start=system_end, end=historical_end) # Protect the Agent instructions and current task from removal. protected = [*messages[:system_end], *task] - # The remaining token budget after protecting the instructions and current task. - available_tokens = target_tokens - token_counter.count(messages=protected) - kept_turn_start = len(historical_turns) - all_step_tokens = token_counter.count(messages=[message for start, end in steps for message in messages[start:end]]) - if all_step_tokens <= available_tokens: - # Historical turns are considered only when the entire current task fits. This ensures that compaction removes - # every older turn before it starts trimming individual steps from the task the Agent is actively working on. - kept_step_start = 0 - available_tokens -= all_step_tokens - while kept_turn_start > 0: - start, end = historical_turns[kept_turn_start - 1] - turn = [message for message in messages[start:end] if _COMPACTION_META_KEY not in message.meta] - turn_tokens = token_counter.count(messages=turn) - if turn_tokens > available_tokens: - break - available_tokens -= turn_tokens - kept_turn_start -= 1 - else: - # Even after dropping every historical turn, the current task is too large. Work backwards through its Agent - # steps and retain the most recent complete suffix that fits. - kept_step_start = len(steps) - while kept_step_start > 0: - start, end = steps[kept_step_start - 1] - step_tokens = token_counter.count(messages=messages[start:end]) - if step_tokens > available_tokens: - break - available_tokens -= step_tokens - kept_step_start -= 1 - - # Enforce the minimum number of complete steps, even when they exceed the target token budget. - kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) + kept_turn_start, kept_step_start = _retained_span_starts( + messages=messages, + protected=protected, + historical_turns=historical_turns, + steps=steps, + target_tokens=target_tokens, + token_counter=token_counter, + min_keep_steps=min_keep_steps, + ) kept_step_spans = steps[kept_step_start:] kept_turn_spans = historical_turns[kept_turn_start:] @@ -134,14 +222,15 @@ def _task_and_step_split( # Everything outside the protected context and retained steps can be removed. removable = [message for index, message in enumerate(messages) if index not in kept_indices] - kept_turns = [ - message - for start, end in kept_turn_spans - for message in messages[start:end] - if _COMPACTION_META_KEY not in message.meta - ] - kept_steps = [message for start, end in kept_step_spans for message in messages[start:end]] - return [*messages[:system_end], *kept_turns, *task], removable, kept_steps, len(steps) - kept_step_start + kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) + kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) + if kept_step_start == 0: + kept_before_note = [*messages[:system_end], *kept_turns] + kept_after_note = [*task, *kept_steps] + else: + kept_before_note = [*messages[:system_end], *task] + kept_after_note = kept_steps + return kept_before_note, removable, kept_after_note, len(steps) - kept_step_start @_experimental @@ -199,7 +288,7 @@ def compact( """ if token_counter.count(messages) <= target_tokens: return None - protected, removable, kept_steps, kept_step_count = _task_and_step_split( + kept_before_note, removable, kept_after_note, kept_step_count = _task_and_step_split( messages=messages, target_tokens=target_tokens, token_counter=token_counter, @@ -208,7 +297,7 @@ def compact( if not removable: return None if not self.omission_note: - return [*protected, *kept_steps] + return [*kept_before_note, *kept_after_note] # We prefer user over system since not all providers support multiple system messages note = ChatMessage.from_user( @@ -218,7 +307,7 @@ def compact( _COMPACTION_META_KEY: { "strategy": "sliding_window", "removed_messages": len(removable), - "kept_messages": len(protected) + len(kept_steps), + "kept_messages": len(kept_before_note) + len(kept_after_note), "kept_steps": kept_step_count, } }, @@ -226,7 +315,7 @@ def compact( # The note costs tokens of its own, so it is only worth leaving behind if what it stands in for is bigger. if token_counter.count([note]) >= token_counter.count(removable): return None - return [*protected, note, *kept_steps] + return [*kept_before_note, note, *kept_after_note] def to_dict(self) -> dict[str, Any]: """ diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 635f4407d73..4ba9627283e 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -76,6 +76,20 @@ def test_replaces_the_middle_with_an_omission_note(self): # A user message, not a system one, so providers that hoist system messages cannot move it out of position. assert compacted[2].is_from(role=ChatRole.USER) + def test_places_omission_note_before_current_task_when_only_historical_turns_are_removed(self): + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user(text="old question"), + ChatMessage.from_assistant(text="old answer " * 100), + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current answer"), + ] + compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=20, token_counter=COUNTER) + assert compacted is not None + assert compacted[0] == messages[0] + assert _COMPACTION_META_KEY in compacted[1].meta + assert compacted[2:] == messages[3:] + def test_a_roomier_target_keeps_more(self): messages = [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="task")] for index in range(4): From 7c7f87d80450c6e3bf178efdae8da086d3fbaed9 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 13:58:05 +0200 Subject: [PATCH 21/38] making the logic less insane --- haystack/hooks/compaction/sliding_window.py | 100 ++++++++++++-------- 1 file changed, 62 insertions(+), 38 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index f3ed181bacb..fe8636182dc 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -121,44 +121,55 @@ def _fitting_suffix_start( return kept_start -def _retained_span_starts( +def _get_turn_start( messages: list[ChatMessage], - protected: list[ChatMessage], + available_tokens: int, historical_turns: list[tuple[int, int]], - steps: list[tuple[int, int]], - target_tokens: int, + current_agent_steps: list[tuple[int, int]], token_counter: TokenCounter, - min_keep_steps: int, -) -> tuple[int, int]: - """Return the first retained historical turn and current-task step.""" - available_tokens = target_tokens - token_counter.count(messages=protected) - all_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=steps, skip_compaction_notes=False) +) -> int: + """Return the start of the retained historical turns, or `len(historical_turns)` if none fit.""" + all_current_agent_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) ) kept_turn_start = len(historical_turns) - if all_step_tokens <= available_tokens: + if all_current_agent_step_tokens <= available_tokens: # Historical turns are considered only when the entire current task fits. This ensures that compaction removes # every older turn before it starts trimming individual steps from the task the Agent is actively working on. - kept_step_start = 0 kept_turn_start = _fitting_suffix_start( messages=messages, spans=historical_turns, - available_tokens=available_tokens - all_step_tokens, + available_tokens=available_tokens - all_current_agent_step_tokens, token_counter=token_counter, skip_compaction_notes=True, ) - else: - # Even after dropping every historical turn, the current task is too large. Retain the most recent complete - # suffix of Agent steps that fits. - kept_step_start = _fitting_suffix_start( - messages=messages, - spans=steps, - available_tokens=available_tokens, - token_counter=token_counter, - skip_compaction_notes=False, - ) - kept_step_start = min(kept_step_start, max(len(steps) - min_keep_steps, 0)) - return kept_turn_start, kept_step_start + return kept_turn_start + + +def _get_step_start( + messages: list[ChatMessage], + available_tokens: int, + current_agent_steps: list[tuple[int, int]], + token_counter: TokenCounter, + min_keep_steps: int, +) -> int: + """Return the start of the retained current-task steps, or `len(current_agent_steps)` if none fit.""" + all_current_agent_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) + ) + if all_current_agent_step_tokens <= available_tokens: + # Since all current-task steps fit, we can retain all of them. + return 0 + # Even after dropping every historical turn, the current task is too large. Retain the most recent complete + # suffix of Agent steps that fits. + kept_step_start = _fitting_suffix_start( + messages=messages, + spans=current_agent_steps, + available_tokens=available_tokens, + token_counter=token_counter, + skip_compaction_notes=False, + ) + return min(kept_step_start, max(len(current_agent_steps) - min_keep_steps, 0)) def _task_and_step_split( @@ -189,30 +200,43 @@ def _task_and_step_split( """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) + # Find the latest user message to use as the current task anchor. task_index = _latest_user_index(messages=messages) task = [messages[task_index]] if task_index is not None else [] - step_start = (task_index + 1) if task_index is not None else system_end - # Current-task steps can be removed individually. Earlier user/assistant exchanges are kept as complete turns so - # an assistant reply is never retained without the user message it answers. - steps = _agent_step_spans(messages=messages, start=step_start) + + # Find the complete Agent steps that follow the current task anchor. + current_task_step_start = (task_index + 1) if task_index is not None else system_end + current_agent_steps = _agent_step_spans(messages=messages, start=current_task_step_start) + + # Find all complete historical turns (i.e. user-assistant) that precede the current task. historical_end = task_index if task_index is not None else system_end historical_turns = _historical_turn_spans(messages=messages, start=system_end, end=historical_end) - # Protect the Agent instructions and current task from removal. - protected = [*messages[:system_end], *task] - kept_turn_start, kept_step_start = _retained_span_starts( + # Calculate the token count of the protected context (leading system messages and current task) + protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + + # Get the start of where we should retain historical turns. + # If no turns are kept the value of this is `len(historical_turns)`, which is the same as `historical_end`. + kept_turn_start = _get_turn_start( messages=messages, - protected=protected, + available_tokens=target_tokens - protected_tokens, historical_turns=historical_turns, - steps=steps, - target_tokens=target_tokens, + current_agent_steps=current_agent_steps, token_counter=token_counter, - min_keep_steps=min_keep_steps, ) - kept_step_spans = steps[kept_step_start:] kept_turn_spans = historical_turns[kept_turn_start:] + # Get the start of where we should retain current-task steps + kept_step_start = _get_step_start( + messages=messages, + available_tokens=target_tokens - protected_tokens, + current_agent_steps=current_agent_steps, + token_counter=token_counter, + min_keep_steps=min_keep_steps, + ) + kept_step_spans = current_agent_steps[kept_step_start:] + # Record every protected or retained message index; equal ChatMessages can appear more than once in the list. kept_indices = {*range(system_end)} if task_index is not None: @@ -230,7 +254,7 @@ def _task_and_step_split( else: kept_before_note = [*messages[:system_end], *task] kept_after_note = kept_steps - return kept_before_note, removable, kept_after_note, len(steps) - kept_step_start + return kept_before_note, removable, kept_after_note, len(current_agent_steps) - kept_step_start @_experimental From 1bd22359476f7f4f16bd598df2741e1d2909e884 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:07:13 +0200 Subject: [PATCH 22/38] more logic refactoring --- haystack/hooks/compaction/sliding_window.py | 34 ++++++++++----------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index fe8636182dc..00640f7d226 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -74,7 +74,7 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> def _messages_from_spans( - messages: list[ChatMessage], spans: list[tuple[int, int]], *, skip_compaction_notes: bool = False + messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False ) -> list[ChatMessage]: """Flatten message spans, optionally excluding messages produced by an earlier compaction.""" return [ @@ -90,8 +90,7 @@ def _fitting_suffix_start( spans: list[tuple[int, int]], available_tokens: int, token_counter: TokenCounter, - *, - skip_compaction_notes: bool = False, + skip_compaction_notes: bool, ) -> int: """ Return the first span in the newest contiguous suffix that fits the token budget. @@ -124,14 +123,11 @@ def _fitting_suffix_start( def _get_turn_start( messages: list[ChatMessage], available_tokens: int, + all_current_agent_step_tokens: int, historical_turns: list[tuple[int, int]], - current_agent_steps: list[tuple[int, int]], token_counter: TokenCounter, ) -> int: """Return the start of the retained historical turns, or `len(historical_turns)` if none fit.""" - all_current_agent_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) - ) kept_turn_start = len(historical_turns) if all_current_agent_step_tokens <= available_tokens: # Historical turns are considered only when the entire current task fits. This ensures that compaction removes @@ -149,14 +145,12 @@ def _get_turn_start( def _get_step_start( messages: list[ChatMessage], available_tokens: int, + all_current_agent_step_tokens: int, current_agent_steps: list[tuple[int, int]], token_counter: TokenCounter, min_keep_steps: int, ) -> int: """Return the start of the retained current-task steps, or `len(current_agent_steps)` if none fit.""" - all_current_agent_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) - ) if all_current_agent_step_tokens <= available_tokens: # Since all current-task steps fit, we can retain all of them. return 0 @@ -216,38 +210,44 @@ def _task_and_step_split( # Calculate the token count of the protected context (leading system messages and current task) protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + # Calculate the size of the current task's Agent steps + all_current_agent_step_tokens = token_counter.count( + messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) + ) + # Get the start of where we should retain historical turns. # If no turns are kept the value of this is `len(historical_turns)`, which is the same as `historical_end`. kept_turn_start = _get_turn_start( messages=messages, available_tokens=target_tokens - protected_tokens, + all_current_agent_step_tokens=all_current_agent_step_tokens, historical_turns=historical_turns, - current_agent_steps=current_agent_steps, token_counter=token_counter, ) kept_turn_spans = historical_turns[kept_turn_start:] + kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) # Get the start of where we should retain current-task steps kept_step_start = _get_step_start( messages=messages, available_tokens=target_tokens - protected_tokens, + all_current_agent_step_tokens=all_current_agent_step_tokens, current_agent_steps=current_agent_steps, token_counter=token_counter, min_keep_steps=min_keep_steps, ) kept_step_spans = current_agent_steps[kept_step_start:] + kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) - # Record every protected or retained message index; equal ChatMessages can appear more than once in the list. + # Record every index that we are keeping kept_indices = {*range(system_end)} if task_index is not None: kept_indices.add(task_index) for start, end in [*kept_turn_spans, *kept_step_spans]: - kept_indices.update(index for index in range(start, end) if _COMPACTION_META_KEY not in messages[index].meta) - - # Everything outside the protected context and retained steps can be removed. + kept_indices.update(index for index in range(start, end)) + # Remove everything else that's not in the kept indices removable = [message for index, message in enumerate(messages) if index not in kept_indices] - kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) - kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) + if kept_step_start == 0: kept_before_note = [*messages[:system_end], *kept_turns] kept_after_note = [*task, *kept_steps] From f10a2bb85cb4597c5eb08cfb7bf320ea09d23675 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:29:03 +0200 Subject: [PATCH 23/38] simplifications --- haystack/hooks/compaction/sliding_window.py | 237 +++++++++---------- test/hooks/compaction/test_sliding_window.py | 25 +- 2 files changed, 135 insertions(+), 127 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 00640f7d226..d2b95090c90 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -73,124 +73,124 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> ] -def _messages_from_spans( +def _index_groups( messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False -) -> list[ChatMessage]: - """Flatten message spans, optionally excluding messages produced by an earlier compaction.""" +) -> list[list[int]]: + """ + Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. + """ return [ - message + [ + index + for index in range(start, end) + if not skip_compaction_notes or _COMPACTION_META_KEY not in messages[index].meta + ] for start, end in spans - for message in messages[start:end] - if not skip_compaction_notes or _COMPACTION_META_KEY not in message.meta ] -def _fitting_suffix_start( - messages: list[ChatMessage], - spans: list[tuple[int, int]], - available_tokens: int, - token_counter: TokenCounter, - skip_compaction_notes: bool, +def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages at the given indices, in conversation order.""" + return [messages[index] for index in indices] + + +def _flatten(groups: list[list[int]]) -> list[int]: + """Join index groups into a single ordered list of indices.""" + return [index for group in groups for index in group] + + +def _first_group_to_keep( + messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter ) -> int: """ - Return the first span in the newest contiguous suffix that fits the token budget. - - Spans are measured from newest to oldest. Retention stops as soon as a span does not fit, ensuring that an older - span is never kept after a newer one has been removed. - - :param messages: The full conversation containing the messages referenced by `spans`. - :param spans: Ordered `(start_index, end_index)` pairs to consider for retention. Both indices refer to `messages`, - and `end_index` is exclusive. - :param available_tokens: The token budget available for retaining messages from `spans`. - :param token_counter: The `TokenCounter` used to measure each span. - :param skip_compaction_notes: Whether messages produced by an earlier compaction are excluded from token counting. - :returns: The index in `spans` at which the retained suffix begins. If no span fits, returns `len(spans)`; if every - span fits, returns `0`. + Return the oldest group that the budget can still pay for, working backwards from the newest. + + Groups are added up from newest to oldest and counting stops at the first group that does not fit, so what is kept + is always a run of groups at the end of the list. An older group is never kept once a newer one has been dropped, + which would leave a hole in the conversation. + + :param messages: The full conversation containing the messages referenced by `groups`. + :param groups: Ordered index groups to choose from, oldest group first. + :param available_tokens: The token budget available for these groups. + :param token_counter: The `TokenCounter` used to measure each group. + :returns: The position in `groups` to start keeping from: `len(groups)` when nothing fits, `0` when it all fits. """ - kept_start = len(spans) - while kept_start > 0: - span_messages = _messages_from_spans( - messages=messages, spans=[spans[kept_start - 1]], skip_compaction_notes=skip_compaction_notes - ) - span_tokens = token_counter.count(messages=span_messages) - if span_tokens > available_tokens: + first_kept = len(groups) + while first_kept > 0: + group_tokens = token_counter.count(messages=_messages_at(messages=messages, indices=groups[first_kept - 1])) + if group_tokens > available_tokens: break - available_tokens -= span_tokens - kept_start -= 1 - return kept_start + available_tokens -= group_tokens + first_kept -= 1 + return first_kept -def _get_turn_start( +def _first_turn_and_step_to_keep( messages: list[ChatMessage], + turn_groups: list[list[int]], + step_groups: list[list[int]], available_tokens: int, - all_current_agent_step_tokens: int, - historical_turns: list[tuple[int, int]], token_counter: TokenCounter, -) -> int: - """Return the start of the retained historical turns, or `len(historical_turns)` if none fit.""" - kept_turn_start = len(historical_turns) - if all_current_agent_step_tokens <= available_tokens: - # Historical turns are considered only when the entire current task fits. This ensures that compaction removes - # every older turn before it starts trimming individual steps from the task the Agent is actively working on. - kept_turn_start = _fitting_suffix_start( - messages=messages, - spans=historical_turns, - available_tokens=available_tokens - all_current_agent_step_tokens, - token_counter=token_counter, - skip_compaction_notes=True, + min_keep_steps: int, +) -> tuple[int, int]: + """ + Return which historical turn and which Agent step of the current task to start keeping from. + + :param messages: The full conversation containing the messages referenced by both group lists. + :param turn_groups: Index groups for the complete historical turns preceding the current task, oldest first. + :param step_groups: Index groups for the current task's Agent steps, oldest first. + :param available_tokens: The token budget left once the protected context is paid for. + :param token_counter: The `TokenCounter` used to measure the groups. + :param min_keep_steps: The fewest recent Agent steps to keep, even when they exceed the budget. + :returns: The position in `turn_groups` and the position in `step_groups` to start keeping from. Either is the + length of its list when nothing from it is kept. + """ + current_task_tokens = token_counter.count( + messages=_messages_at(messages=messages, indices=_flatten(groups=step_groups)) + ) + if current_task_tokens > available_tokens: + # The current task alone overruns the budget, so every historical turn is dropped and + # the current task's own oldest steps trimmed until what remains fits. + first_kept_step = _first_group_to_keep( + messages=messages, groups=step_groups, available_tokens=available_tokens, token_counter=token_counter ) - return kept_turn_start + # The newest steps are kept regardless of the budget. + return len(turn_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) - -def _get_step_start( - messages: list[ChatMessage], - available_tokens: int, - all_current_agent_step_tokens: int, - current_agent_steps: list[tuple[int, int]], - token_counter: TokenCounter, - min_keep_steps: int, -) -> int: - """Return the start of the retained current-task steps, or `len(current_agent_steps)` if none fit.""" - if all_current_agent_step_tokens <= available_tokens: - # Since all current-task steps fit, we can retain all of them. - return 0 - # Even after dropping every historical turn, the current task is too large. Retain the most recent complete - # suffix of Agent steps that fits. - kept_step_start = _fitting_suffix_start( + # The whole current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. + first_kept_turn = _first_group_to_keep( messages=messages, - spans=current_agent_steps, - available_tokens=available_tokens, + groups=turn_groups, + available_tokens=available_tokens - current_task_tokens, token_counter=token_counter, - skip_compaction_notes=False, ) - return min(kept_step_start, max(len(current_agent_steps) - min_keep_steps, 0)) + return first_kept_turn, 0 def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int -) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage], int]: +) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage]]: """ Split a conversation into the messages kept before and after an omission note and the messages to remove. - Leading system messages and the latest real user message are always retained. Historical user turns are retained - whole when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed - individually while preserving complete assistant/tool-result groups. + Leading system messages and the latest real user message are always kept. Historical user turns are kept whole when + they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time + while keeping each assistant message together with its tool results. :param messages: The full conversation to split, ordered oldest to newest. - :param target_tokens: The target token budget for the retained messages. + :param target_tokens: The token budget for the messages that are kept. :param token_counter: The `TokenCounter` used to decide which historical turns and current-task steps fit. - :param min_keep_steps: The minimum number of recent current-task Agent steps to retain, even if they exceed the - target token budget. + :param min_keep_steps: The fewest recent current-task Agent steps to keep, even if they exceed the target token + budget. :returns: A tuple containing: - 1. Messages retained before the omission-note position. - 2. Every message selected for removal. - 3. Messages retained after the omission-note position. - 4. The number of retained current-task Agent steps. + 1. Messages kept before the omission-note position. + 2. Messages kept after the omission-note position. + 3. Every message selected for removal. - When only historical turns are removed, the first and third elements place the note immediately before the - current task. When current-task steps are also removed, they place it after the latest user message and before - the retained current-task steps. + When only historical turns are removed, the first two elements place the note immediately before the current + task. When current-task steps are also removed, they place it after the latest user message and before the + current-task steps that survived. """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) @@ -199,62 +199,48 @@ def _task_and_step_split( task_index = _latest_user_index(messages=messages) task = [messages[task_index]] if task_index is not None else [] - # Find the complete Agent steps that follow the current task anchor. - current_task_step_start = (task_index + 1) if task_index is not None else system_end - current_agent_steps = _agent_step_spans(messages=messages, start=current_task_step_start) + # Group the complete Agent steps that follow the current task anchor. + step_start_index = (task_index + 1) if task_index is not None else system_end + step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start_index)) - # Find all complete historical turns (i.e. user-assistant) that precede the current task. + # Group the complete historical turns (i.e. user-assistant) that precede the current task. An earlier compaction's + # note is left out of its turn, so keeping the turn folds that note into the note this compaction leaves behind. historical_end = task_index if task_index is not None else system_end - historical_turns = _historical_turn_spans(messages=messages, start=system_end, end=historical_end) - - # Calculate the token count of the protected context (leading system messages and current task) - protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) - - # Calculate the size of the current task's Agent steps - all_current_agent_step_tokens = token_counter.count( - messages=_messages_from_spans(messages=messages, spans=current_agent_steps, skip_compaction_notes=False) - ) - - # Get the start of where we should retain historical turns. - # If no turns are kept the value of this is `len(historical_turns)`, which is the same as `historical_end`. - kept_turn_start = _get_turn_start( + turn_groups = _index_groups( messages=messages, - available_tokens=target_tokens - protected_tokens, - all_current_agent_step_tokens=all_current_agent_step_tokens, - historical_turns=historical_turns, - token_counter=token_counter, + spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), + skip_compaction_notes=True, ) - kept_turn_spans = historical_turns[kept_turn_start:] - kept_turns = _messages_from_spans(messages=messages, spans=kept_turn_spans, skip_compaction_notes=True) - # Get the start of where we should retain current-task steps - kept_step_start = _get_step_start( + # The Agent instructions and the current task are never removed, so they are paid for out of the target first. + protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + first_kept_turn, first_kept_step = _first_turn_and_step_to_keep( messages=messages, + turn_groups=turn_groups, + step_groups=step_groups, available_tokens=target_tokens - protected_tokens, - all_current_agent_step_tokens=all_current_agent_step_tokens, - current_agent_steps=current_agent_steps, token_counter=token_counter, min_keep_steps=min_keep_steps, ) - kept_step_spans = current_agent_steps[kept_step_start:] - kept_steps = _messages_from_spans(messages=messages, spans=kept_step_spans, skip_compaction_notes=False) + kept_turn_indices = _flatten(groups=turn_groups[first_kept_turn:]) + kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) - # Record every index that we are keeping - kept_indices = {*range(system_end)} + # A message survives only by being protected or by falling in a group we are keeping; everything else goes. + kept_indices = {*range(system_end), *kept_turn_indices, *kept_step_indices} if task_index is not None: kept_indices.add(task_index) - for start, end in [*kept_turn_spans, *kept_step_spans]: - kept_indices.update(index for index in range(start, end)) - # Remove everything else that's not in the kept indices removable = [message for index, message in enumerate(messages) if index not in kept_indices] - if kept_step_start == 0: - kept_before_note = [*messages[:system_end], *kept_turns] - kept_after_note = [*task, *kept_steps] + if first_kept_step == 0: + # The current task is untouched, so the note stands in for the older turns and belongs in front of the task. + kept_before_note = [*messages[:system_end], *_messages_at(messages=messages, indices=kept_turn_indices)] + kept_after_note = [*task, *_messages_at(messages=messages, indices=kept_step_indices)] else: + # Steps were cut from the current task, which means every historical turn was already dropped, so the note + # goes between the task and the steps that survived it. kept_before_note = [*messages[:system_end], *task] - kept_after_note = kept_steps - return kept_before_note, removable, kept_after_note, len(current_agent_steps) - kept_step_start + kept_after_note = _messages_at(messages=messages, indices=kept_step_indices) + return kept_before_note, kept_after_note, removable @_experimental @@ -312,7 +298,7 @@ def compact( """ if token_counter.count(messages) <= target_tokens: return None - kept_before_note, removable, kept_after_note, kept_step_count = _task_and_step_split( + kept_before_note, kept_after_note, removable = _task_and_step_split( messages=messages, target_tokens=target_tokens, token_counter=token_counter, @@ -332,7 +318,6 @@ def compact( "strategy": "sliding_window", "removed_messages": len(removable), "kept_messages": len(kept_before_note) + len(kept_after_note), - "kept_steps": kept_step_count, } }, ) diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 4ba9627283e..32754f81ebe 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -70,7 +70,6 @@ def test_replaces_the_middle_with_an_omission_note(self): "strategy": "sliding_window", "removed_messages": 2, "kept_messages": 4, - "kept_steps": 1, } assert compacted[2].text == _DEFAULT_OMISSION_NOTE.replace("{num_removed}", "2") # A user message, not a system one, so providers that hoist system messages cannot move it out of position. @@ -143,6 +142,30 @@ def test_drops_historical_context_and_one_current_task_step_to_reach_target(self ) assert compacted == [system_message, current_task[0], *current_task[-2:]] + def test_folds_an_earlier_note_inside_a_retained_turn_and_counts_it_as_removed(self): + note = ChatMessage.from_user( + "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ) + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user(text="old question " * 200), + ChatMessage.from_assistant(text="old answer"), + ChatMessage.from_user(text="recent question"), + ChatMessage.from_assistant(text="recent answer"), + note, + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Room for everything but the oldest turn, so the turn holding the earlier note is retained around it. + compacted = SlidingWindowCompactor().compact( + messages=messages, target_tokens=COUNTER.count([messages[0], *messages[3:]]), token_counter=COUNTER + ) + assert compacted is not None + # The earlier note is replaced by the new one rather than surviving alongside it, and it counts as removed. + assert count_markers(messages=compacted) == 1 + assert compacted == [messages[0], *messages[3:5], compacted[3], *messages[6:]] + assert compacted[3].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 + def test_returns_none_when_the_conversation_already_fits(self): assert ( SlidingWindowCompactor().compact(messages=long_conversation(), target_tokens=100_000, token_counter=COUNTER) From 1bb59a8ab285dc38208bbdfbfbbf5c903393d2f1 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:40:17 +0200 Subject: [PATCH 24/38] more simplification --- haystack/hooks/compaction/sliding_window.py | 49 ++++++++------------ test/hooks/compaction/test_sliding_window.py | 23 --------- 2 files changed, 20 insertions(+), 52 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index d2b95090c90..f18d1fe07a4 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -169,9 +169,9 @@ def _first_turn_and_step_to_keep( def _task_and_step_split( messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, min_keep_steps: int -) -> tuple[list[ChatMessage], list[ChatMessage], list[ChatMessage]]: +) -> tuple[list[ChatMessage], int, list[ChatMessage]]: """ - Split a conversation into the messages kept before and after an omission note and the messages to remove. + Split a conversation into the messages to keep and the messages to remove. Leading system messages and the latest real user message are always kept. Historical user turns are kept whole when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time @@ -184,13 +184,10 @@ def _task_and_step_split( budget. :returns: A tuple containing: - 1. Messages kept before the omission-note position. - 2. Messages kept after the omission-note position. + 1. The messages to keep, ordered oldest to newest. + 2. The position in that list where an omission note belongs, immediately before the current task when only + historical turns were removed and immediately before the surviving steps when the task itself was trimmed. 3. Every message selected for removal. - - When only historical turns are removed, the first two elements place the note immediately before the current - task. When current-task steps are also removed, they place it after the latest user message and before the - current-task steps that survived. """ # Find the leading system messages that contain the Agent instructions. system_end = _leading_system_end(messages=messages) @@ -224,6 +221,9 @@ def _task_and_step_split( ) kept_turn_indices = _flatten(groups=turn_groups[first_kept_turn:]) kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) + kept_turns = _messages_at(messages=messages, indices=kept_turn_indices) + kept_steps = _messages_at(messages=messages, indices=kept_step_indices) + kept = [*messages[:system_end], *kept_turns, *task, *kept_steps] # A message survives only by being protected or by falling in a group we are keeping; everything else goes. kept_indices = {*range(system_end), *kept_turn_indices, *kept_step_indices} @@ -231,16 +231,10 @@ def _task_and_step_split( kept_indices.add(task_index) removable = [message for index, message in enumerate(messages) if index not in kept_indices] - if first_kept_step == 0: - # The current task is untouched, so the note stands in for the older turns and belongs in front of the task. - kept_before_note = [*messages[:system_end], *_messages_at(messages=messages, indices=kept_turn_indices)] - kept_after_note = [*task, *_messages_at(messages=messages, indices=kept_step_indices)] - else: - # Steps were cut from the current task, which means every historical turn was already dropped, so the note - # goes between the task and the steps that survived it. - kept_before_note = [*messages[:system_end], *task] - kept_after_note = _messages_at(messages=messages, indices=kept_step_indices) - return kept_before_note, kept_after_note, removable + # The note stands in for the newest thing that was dropped. That is the current task's own steps when those were + # cut, in which case every historical turn went too and `kept_turns` is empty; otherwise it is the older turns. + note_index = system_end + len(kept_turns) + (len(task) if first_kept_step > 0 else 0) + return kept, note_index, removable @_experimental @@ -291,14 +285,14 @@ def compact( Drop older history while preserving the task anchor and a complete recent conversation window. :param messages: The conversation to compact, oldest to newest. - :param target_tokens: The size the retained conversation should come in under. + :param target_tokens: The size the kept conversation should come in under. :param token_counter: The `TokenCounter` to measure messages with. - :returns: The protected context, an omission note if configured, and the retained steps; or None when there is - nothing worth removing. + :returns: The protected context, an omission note if configured, and the steps that survived; or None when + there is nothing to remove. """ - if token_counter.count(messages) <= target_tokens: + if token_counter.count(messages=messages) <= target_tokens: return None - kept_before_note, kept_after_note, removable = _task_and_step_split( + kept, note_index, removable = _task_and_step_split( messages=messages, target_tokens=target_tokens, token_counter=token_counter, @@ -307,7 +301,7 @@ def compact( if not removable: return None if not self.omission_note: - return [*kept_before_note, *kept_after_note] + return kept # We prefer user over system since not all providers support multiple system messages note = ChatMessage.from_user( @@ -317,14 +311,11 @@ def compact( _COMPACTION_META_KEY: { "strategy": "sliding_window", "removed_messages": len(removable), - "kept_messages": len(kept_before_note) + len(kept_after_note), + "kept_messages": len(kept), } }, ) - # The note costs tokens of its own, so it is only worth leaving behind if what it stands in for is bigger. - if token_counter.count([note]) >= token_counter.count(removable): - return None - return [*kept_before_note, note, *kept_after_note] + return [*kept[:note_index], note, *kept[note_index:]] def to_dict(self) -> dict[str, Any]: """ diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 32754f81ebe..bafacf91205 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -268,29 +268,6 @@ def test_repeated_compaction_folds_the_previous_note(self): assert second[0].text == "rules" assert second[1].text == "start" - @pytest.mark.parametrize( - ("removable", "worth_replacing"), - [ - pytest.param("a", False, id="cheaper-than-the-note"), - pytest.param("x" * 4000, True, id="dearer-than-the-note"), - ], - ) - def test_a_cut_is_made_only_when_the_note_costs_less_than_what_it_replaces(self, removable, worth_replacing): - # A note is not free, so what matters is the size of what goes rather than how many messages it is: one long - # tool result is worth replacing, one short message is not. - messages = [ - ChatMessage.from_system(text="rules"), - ChatMessage.from_user(text="task"), - ChatMessage.from_assistant(text=removable), - ChatMessage.from_assistant(text="latest step"), - ] - compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) - assert (compacted is not None) is worth_replacing - # Without a note there is nothing to pay for, so the same cut is always worth making. - assert SlidingWindowCompactor(omission_note=None).compact( - messages=messages, target_tokens=SMALLEST, token_counter=COUNTER - ) == [*messages[:2], messages[-1]] - def test_keeping_no_steps_still_preserves_the_current_task(self): messages = long_conversation() compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( From f60d07a9f0033e4058202b4007dc50494457c4e8 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Fri, 7 Aug 2026 14:52:28 +0200 Subject: [PATCH 25/38] update reno and add more dev comments --- haystack/hooks/compaction/sliding_window.py | 15 ++++++++------- ...agent-context-compaction-3258c08dec9d2b34.yaml | 8 +++++--- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index f18d1fe07a4..d3961abd32f 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -103,21 +103,22 @@ def _first_group_to_keep( messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter ) -> int: """ - Return the oldest group that the budget can still pay for, working backwards from the newest. + Return the position in `groups` to start keeping from. - Groups are added up from newest to oldest and counting stops at the first group that does not fit, so what is kept - is always a run of groups at the end of the list. An older group is never kept once a newer one has been dropped, - which would leave a hole in the conversation. + Groups are added up from newest to oldest and counting stops at the first group that does not fit. :param messages: The full conversation containing the messages referenced by `groups`. - :param groups: Ordered index groups to choose from, oldest group first. + :param groups: Ordered index groups to choose from, the oldest group first is at position 0. :param available_tokens: The token budget available for these groups. :param token_counter: The `TokenCounter` used to measure each group. :returns: The position in `groups` to start keeping from: `len(groups)` when nothing fits, `0` when it all fits. """ + # We count backwards so first_kept starts such that nothing would be kept first_kept = len(groups) while first_kept > 0: + # Calculate the tokens consumed by the group that would be kept next group_tokens = token_counter.count(messages=_messages_at(messages=messages, indices=groups[first_kept - 1])) + # If this group does not fit, we stop if group_tokens > available_tokens: break available_tokens -= group_tokens @@ -149,8 +150,8 @@ def _first_turn_and_step_to_keep( messages=_messages_at(messages=messages, indices=_flatten(groups=step_groups)) ) if current_task_tokens > available_tokens: - # The current task alone overruns the budget, so every historical turn is dropped and - # the current task's own oldest steps trimmed until what remains fits. + # The current task alone overruns the budget, so every historical turn is dropped and the current task's own + # oldest steps trimmed until what remains fits. first_kept_step = _first_group_to_keep( messages=messages, groups=step_groups, available_tokens=available_tokens, token_counter=token_counter ) diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index c71c803f853..1c7d53ec068 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -32,9 +32,11 @@ features: ``SlidingWindowCompactor`` treats an assistant message and its following tool results as one step, so a tool call is never separated from its results. Historical turns are also kept whole, so an assistant reply is not retained - without the user message it answers. It may retain slightly more history than the requested target when preserving - the current task or the configured minimum number of steps. Compaction is lossy: removed messages cannot be - recovered or summarized by this strategy. Implement the ``Compactor`` protocol to provide a custom strategy. + without the user message it answers. It can also land above the requested target rather than under it, because + leading system messages and the current task are never removed and ``min_keep_steps`` holds on to the newest Agent + steps whatever their size, so a long system prompt or one large tool result can leave the conversation well over + the target. Compaction is lossy: removed messages cannot be recovered or summarized by this strategy. Implement + the ``Compactor`` protocol to provide a custom strategy. ``CompactionHook`` and ``SlidingWindowCompactor`` emit an ``ExperimentalWarning`` and may change without a deprecation cycle. From c8690f738df2b045efb06f62504bd7fc77f1512d Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 10:50:26 +0200 Subject: [PATCH 26/38] PR comments --- haystack/hooks/compaction/sliding_window.py | 42 ++++---- ...t-context-compaction-3258c08dec9d2b34.yaml | 5 +- test/hooks/compaction/helpers.py | 7 +- test/hooks/compaction/test_hooks.py | 12 +-- test/hooks/compaction/test_sliding_window.py | 95 +++++++++---------- 5 files changed, 80 insertions(+), 81 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index d3961abd32f..750bb24de06 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -128,7 +128,7 @@ def _first_group_to_keep( def _first_turn_and_step_to_keep( messages: list[ChatMessage], - turn_groups: list[list[int]], + historical_groups: list[list[int]], step_groups: list[list[int]], available_tokens: int, token_counter: TokenCounter, @@ -138,12 +138,12 @@ def _first_turn_and_step_to_keep( Return which historical turn and which Agent step of the current task to start keeping from. :param messages: The full conversation containing the messages referenced by both group lists. - :param turn_groups: Index groups for the complete historical turns preceding the current task, oldest first. + :param historical_groups: Index groups for the complete historical turns preceding the current task, oldest first. :param step_groups: Index groups for the current task's Agent steps, oldest first. :param available_tokens: The token budget left once the protected context is paid for. :param token_counter: The `TokenCounter` used to measure the groups. :param min_keep_steps: The fewest recent Agent steps to keep, even when they exceed the budget. - :returns: The position in `turn_groups` and the position in `step_groups` to start keeping from. Either is the + :returns: The position in `historical_groups` and the position in `step_groups` to start keeping from. Either is the length of its list when nothing from it is kept. """ current_task_tokens = token_counter.count( @@ -156,12 +156,12 @@ def _first_turn_and_step_to_keep( messages=messages, groups=step_groups, available_tokens=available_tokens, token_counter=token_counter ) # The newest steps are kept regardless of the budget. - return len(turn_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) + return len(historical_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) # The whole current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. first_kept_turn = _first_group_to_keep( messages=messages, - groups=turn_groups, + groups=historical_groups, available_tokens=available_tokens - current_task_tokens, token_counter=token_counter, ) @@ -186,8 +186,9 @@ def _task_and_step_split( :returns: A tuple containing: 1. The messages to keep, ordered oldest to newest. - 2. The position in that list where an omission note belongs, immediately before the current task when only - historical turns were removed and immediately before the surviving steps when the task itself was trimmed. + 2. The position in that list where an omission note belongs, which is where the removed messages used to sit: + directly after the leading system messages when only historical turns were removed, and directly after the + user message anchoring the current task when the task's own steps were removed. 3. Every message selected for removal. """ # Find the leading system messages that contain the Agent instructions. @@ -204,7 +205,7 @@ def _task_and_step_split( # Group the complete historical turns (i.e. user-assistant) that precede the current task. An earlier compaction's # note is left out of its turn, so keeping the turn folds that note into the note this compaction leaves behind. historical_end = task_index if task_index is not None else system_end - turn_groups = _index_groups( + historical_groups = _index_groups( messages=messages, spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), skip_compaction_notes=True, @@ -214,13 +215,13 @@ def _task_and_step_split( protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) first_kept_turn, first_kept_step = _first_turn_and_step_to_keep( messages=messages, - turn_groups=turn_groups, + historical_groups=historical_groups, step_groups=step_groups, available_tokens=target_tokens - protected_tokens, token_counter=token_counter, min_keep_steps=min_keep_steps, ) - kept_turn_indices = _flatten(groups=turn_groups[first_kept_turn:]) + kept_turn_indices = _flatten(groups=historical_groups[first_kept_turn:]) kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) kept_turns = _messages_at(messages=messages, indices=kept_turn_indices) kept_steps = _messages_at(messages=messages, indices=kept_step_indices) @@ -232,9 +233,10 @@ def _task_and_step_split( kept_indices.add(task_index) removable = [message for index, message in enumerate(messages) if index not in kept_indices] - # The note stands in for the newest thing that was dropped. That is the current task's own steps when those were - # cut, in which case every historical turn went too and `kept_turns` is empty; otherwise it is the older turns. - note_index = system_end + len(kept_turns) + (len(task) if first_kept_step > 0 else 0) + # The note stands in for what was dropped, so it goes where the dropped messages used to sit. Either right after + # the leading system messages when the historical turns were trimmed, or right after the user message that anchors + # the current task when its own Agent steps were trimmed. Both positions are counted off the layout of `kept`. + note_index = system_end if first_kept_step == 0 else system_end + len(kept_turns) + len(task) return kept, note_index, removable @@ -243,9 +245,13 @@ class SlidingWindowCompactor(Compactor): """ Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows. - Leading system messages and the latest user message are protected. Earlier user/assistant turns are retained when - they fit, and the current task's history is retained in complete Agent steps, where a step is an assistant message - together with all immediately following tool results. An `omission_note` is left in place of what was removed. + Leading system messages and the latest user message are protected. Earlier user/assistant turns are kept when they + fit, and the current task's history is kept in complete Agent steps, where a step is an assistant message together + with all immediately following tool results. + + An `omission_note` is left where the removed messages used to sit: directly after the leading system messages when + only earlier turns were removed, and directly after the latest user message when the current task's own steps were + removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement. ```python from haystack.components.agents import Agent @@ -288,8 +294,8 @@ def compact( :param messages: The conversation to compact, oldest to newest. :param target_tokens: The size the kept conversation should come in under. :param token_counter: The `TokenCounter` to measure messages with. - :returns: The protected context, an omission note if configured, and the steps that survived; or None when - there is nothing to remove. + :returns: The conversation that survived, with an omission note if configured standing where the removed + messages used to sit; or None when there is nothing to remove. """ if token_counter.count(messages=messages) <= target_tokens: return None diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index 1c7d53ec068..6caa6298451 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -7,7 +7,10 @@ features: The first built-in strategy, ``SlidingWindowCompactor``, preserves leading system messages, the latest user task, and as much complete recent conversation as the target allows. It removes earlier user/assistant turns as whole units first. Only when removing all earlier turns is insufficient does it remove individual Agent steps from the - current task. It replaces removed history with a short omission note. + current task. It replaces removed history with a short omission note, left where the removed messages used to sit: + directly after the leading system messages when only earlier turns were removed, and directly after the latest user + message when the current task's own steps were removed. Only one note is ever present, because a later compaction + folds an earlier one into its replacement. .. code-block:: python diff --git a/test/hooks/compaction/helpers.py b/test/hooks/compaction/helpers.py index 637e6005d98..6966aade7f5 100644 --- a/test/hooks/compaction/helpers.py +++ b/test/hooks/compaction/helpers.py @@ -67,12 +67,11 @@ def make_state(messages: list[ChatMessage], **data: Any) -> State: return State(schema=_SCHEMA, data={**base, **data}) -def long_conversation() -> list[ChatMessage]: +def fresh_conversation_with_two_steps() -> list[ChatMessage]: """ - Six messages: a system prefix, a user turn, then two tool round-trips. + A system prefix and a first user task with two Agent steps behind it, so there are no earlier turns to remove. - The results are padded so that removing them saves more than an omission note costs, which is what a compactor - weighs before leaving one behind. + The tool results are padded so that dropping a step is a saving worth making. """ return [ ChatMessage.from_system("rules"), diff --git a/test/hooks/compaction/test_hooks.py b/test/hooks/compaction/test_hooks.py index 4a46630bc9a..1dbeb56be36 100644 --- a/test/hooks/compaction/test_hooks.py +++ b/test/hooks/compaction/test_hooks.py @@ -18,7 +18,7 @@ from test.hooks.compaction.helpers import ( FakeCounter, count_markers, - long_conversation, + fresh_conversation_with_two_steps, make_state, tool_call, tool_result, @@ -165,7 +165,7 @@ class TestCompactionHook: ) def test_trigger(self, context_tokens, should_compact): compactor = _RecordingCompactor() - _hook(compactor).run(make_state(messages=long_conversation(), context_tokens=context_tokens)) + _hook(compactor).run(make_state(messages=fresh_conversation_with_two_steps(), context_tokens=context_tokens)) assert compactor.calls == (["compact"] if should_compact else []) def test_fires_without_reported_usage_by_counting_locally(self): @@ -201,7 +201,7 @@ def test_subtracts_provider_overhead_from_the_target(self): # The reported count also covers tool schemas and template overhead, which a compactor cannot remove. Here that # overhead alone exceeds the target, so the compactor is told to cut the messages as far as it is allowed. compactor = _RecordingCompactor() - _hook(compactor).run(make_state(messages=long_conversation(), context_tokens=800)) + _hook(compactor).run(make_state(messages=fresh_conversation_with_two_steps(), context_tokens=800)) assert compactor.targets[0] == 0 def test_warns_when_the_token_counter_exceeds_the_context_estimate(self, caplog): @@ -213,7 +213,7 @@ def test_warns_when_the_token_counter_exceeds_the_context_estimate(self, caplog) def test_rewrites_messages_and_re_estimates_context_tokens(self): counter = FakeCounter() hook = _hook(token_counter=counter) - messages = long_conversation() + messages = fresh_conversation_with_two_steps() original_context_tokens = 800 estimated = _estimated_context_tokens( messages=messages, context_tokens=original_context_tokens, token_counter=counter @@ -244,7 +244,7 @@ def test_preserves_the_no_usage_sentinel_after_compaction(self): assert state.data["context_tokens"] == 0 def test_leaves_the_conversation_alone_when_the_compactor_declines(self): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() state = make_state(messages, context_tokens=800) _hook(_RecordingCompactor(result=None)).run(state=state) assert state.data["messages"] == messages @@ -312,7 +312,7 @@ class TestCompactionHookAsync: @pytest.mark.asyncio async def test_run_async_uses_the_async_compaction_path(self): compactor = _RecordingCompactor() - await _hook(compactor).run_async(make_state(long_conversation(), context_tokens=800)) + await _hook(compactor).run_async(make_state(fresh_conversation_with_two_steps(), context_tokens=800)) assert compactor.calls == ["compact_async"] @pytest.mark.asyncio diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index bafacf91205..9745b7b1aa9 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -8,7 +8,13 @@ from haystack.hooks.compaction import SlidingWindowCompactor from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans from haystack.hooks.compaction.utils import _COMPACTION_META_KEY -from test.hooks.compaction.helpers import FakeCounter, count_markers, long_conversation, tool_call, tool_result +from test.hooks.compaction.helpers import ( + FakeCounter, + count_markers, + fresh_conversation_with_two_steps, + tool_call, + tool_result, +) pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") @@ -59,67 +65,43 @@ def test_compaction_note_does_not_start_a_new_turn(self): class TestSlidingWindowCompactor: - def test_replaces_the_middle_with_an_omission_note(self): - messages = long_conversation() + def test_replaces_oldest_agent_step(self): + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) assert compacted is not None - # The instructions and task survive, the note stands in for the older step, and the latest step stays intact. - assert compacted[:2] == messages[:2] - assert compacted[3:] == messages[4:] + # Only the oldest agent step is removed which is now where the omission note is. + assert compacted == [*messages[:2], compacted[2], *messages[4:]] + + # Test omission note assert compacted[2].meta[_COMPACTION_META_KEY] == { "strategy": "sliding_window", "removed_messages": 2, "kept_messages": 4, } assert compacted[2].text == _DEFAULT_OMISSION_NOTE.replace("{num_removed}", "2") - # A user message, not a system one, so providers that hoist system messages cannot move it out of position. assert compacted[2].is_from(role=ChatRole.USER) - def test_places_omission_note_before_current_task_when_only_historical_turns_are_removed(self): + def test_replaces_oldest_historical_turn(self): messages = [ ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="old question"), ChatMessage.from_assistant(text="old answer " * 100), - ChatMessage.from_user(text="current task"), - ChatMessage.from_assistant(text="current answer"), - ] - compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=20, token_counter=COUNTER) - assert compacted is not None - assert compacted[0] == messages[0] - assert _COMPACTION_META_KEY in compacted[1].meta - assert compacted[2:] == messages[3:] - - def test_a_roomier_target_keeps_more(self): - messages = [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="task")] - for index in range(4): - messages.extend([tool_call(f"c{index}"), tool_result(result="x" * 400, call_id=f"c{index}")]) - compactor = SlidingWindowCompactor(omission_note=None) - tight = compactor.compact(messages=messages, target_tokens=60, token_counter=COUNTER) - roomy = compactor.compact(messages=messages, target_tokens=350, token_counter=COUNTER) - assert tight is not None - assert roomy is not None - assert len(tight) == 4 - assert len(roomy) == 8 - - def test_keeps_complete_recent_user_assistant_turns_that_fit(self): - messages = [ - ChatMessage.from_system(text="rules"), - ChatMessage.from_user(text="old question"), - ChatMessage.from_assistant(text="old answer"), ChatMessage.from_user(text="recent question"), ChatMessage.from_assistant(text="recent answer"), ChatMessage.from_user(text="current task"), - ChatMessage.from_assistant(text="current step"), + ChatMessage.from_assistant(text="current answer"), ] - # Drops one historical turn - expected = [messages[0], *messages[3:]] - target_tokens = COUNTER.count(expected) - compacted = SlidingWindowCompactor(omission_note=None).compact( + # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn. + target_tokens = 30 + compacted = SlidingWindowCompactor().compact( messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) - assert compacted == expected + assert compacted is not None + # The oldest historical turn is removed and replaced with an omission note. + assert compacted == [messages[0], compacted[1], *messages[3:]] + assert _COMPACTION_META_KEY in compacted[1].meta - def test_drops_historical_context_and_one_current_task_step_to_reach_target(self): + def test_drops_all_historical_turns_and_oldest_agent_step(self): system_message = ChatMessage.from_system(text="rules") historical_turn = [ ChatMessage.from_user(text="old question"), @@ -147,33 +129,40 @@ def test_folds_an_earlier_note_inside_a_retained_turn_and_counts_it_as_removed(s "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} ) messages = [ + # System ChatMessage.from_system(text="rules"), + # Historical ChatMessage.from_user(text="old question " * 200), ChatMessage.from_assistant(text="old answer"), ChatMessage.from_user(text="recent question"), ChatMessage.from_assistant(text="recent answer"), note, + # Current Task ChatMessage.from_user(text="current task"), ChatMessage.from_assistant(text="current step"), ] - # Room for everything but the oldest turn, so the turn holding the earlier note is retained around it. + # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn, so + # the turn holding the earlier note survives around it. + target_tokens = 30 compacted = SlidingWindowCompactor().compact( - messages=messages, target_tokens=COUNTER.count([messages[0], *messages[3:]]), token_counter=COUNTER + messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) assert compacted is not None + assert compacted == [messages[0], compacted[1], *messages[3:5], *messages[6:]] # The earlier note is replaced by the new one rather than surviving alongside it, and it counts as removed. assert count_markers(messages=compacted) == 1 - assert compacted == [messages[0], *messages[3:5], compacted[3], *messages[6:]] - assert compacted[3].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 + assert compacted[1].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 def test_returns_none_when_the_conversation_already_fits(self): assert ( - SlidingWindowCompactor().compact(messages=long_conversation(), target_tokens=100_000, token_counter=COUNTER) + SlidingWindowCompactor().compact( + messages=fresh_conversation_with_two_steps(), target_tokens=100_000, token_counter=COUNTER + ) is None ) def test_omission_note_can_be_turned_off(self): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(omission_note=None).compact( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) @@ -192,7 +181,7 @@ def test_omission_note_can_be_turned_off(self): ) def test_omission_note_can_be_customized(self, note, expected): compacted = SlidingWindowCompactor(omission_note=note).compact( - messages=long_conversation(), target_tokens=SMALLEST, token_counter=COUNTER + messages=fresh_conversation_with_two_steps(), target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None assert compacted[2].text == expected @@ -232,7 +221,7 @@ def test_keeps_a_parallel_tool_call_together_with_all_results(self): @pytest.mark.parametrize(("min_keep_steps", "expected"), [(0, 0), (1, 1), (2, 2), (20, 2)]) def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, expected): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(min_keep_steps=min_keep_steps, omission_note=None).compact( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) @@ -257,7 +246,9 @@ def test_preserves_the_latest_user_task_when_compacting_input_history(self): def test_repeated_compaction_folds_the_previous_note(self): compactor = SlidingWindowCompactor() - first = compactor.compact(messages=long_conversation(), target_tokens=SMALLEST, token_counter=COUNTER) + first = compactor.compact( + messages=fresh_conversation_with_two_steps(), target_tokens=SMALLEST, token_counter=COUNTER + ) assert first is not None # Simulate two more turns arriving on top of the already-compacted conversation. grown = [*first, tool_call("c3"), tool_result(result="third result", call_id="c3")] @@ -269,7 +260,7 @@ def test_repeated_compaction_folds_the_previous_note(self): assert second[1].text == "start" def test_keeping_no_steps_still_preserves_the_current_task(self): - messages = long_conversation() + messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) @@ -292,7 +283,7 @@ class TestSlidingWindowCompactorAsync: async def test_compact_async_matches_compact(self): # `SlidingWindowCompactor` does no I/O, so it relies on the protocol's default `compact_async`. compactor = SlidingWindowCompactor() - messages = long_conversation() + messages = fresh_conversation_with_two_steps() assert await compactor.compact_async( messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) == compactor.compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) From 71f19e4c37ae0683b31c2ab9c7d2efb371dbb1ed Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 11:35:47 +0200 Subject: [PATCH 27/38] refactoring tests and improve previous compaction note detection --- haystack/hooks/compaction/sliding_window.py | 23 ++- test/hooks/compaction/test_sliding_window.py | 198 +++++++++++++------ 2 files changed, 159 insertions(+), 62 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 750bb24de06..37227923091 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -11,6 +11,9 @@ from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental +# Recorded as the strategy on every message this compactor produces, so a later run can recognize its own notes. +_STRATEGY = "sliding_window" + # Placeholder a custom omission note may include to have the number of removed messages substituted in. _NUM_REMOVED_PLACEHOLDER = "{num_removed}" @@ -44,6 +47,18 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: return None +def _is_compaction_note(message: ChatMessage) -> bool: + """ + Whether a message is an omission note this strategy left in place of removed history. + + Every compactor marks what it produces with the same meta key, including the tool results + `ToolResultPruningCompactor` rewrites into a placeholder. Matching on the role and the strategy keeps those out: + they are still part of the conversation and have to travel with the turn they belong to. + """ + marker = message.meta.get(_COMPACTION_META_KEY) + return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY + + def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: """ Return spans for complete user turns in a bounded section of conversation history. @@ -80,11 +95,7 @@ def _index_groups( Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. """ return [ - [ - index - for index in range(start, end) - if not skip_compaction_notes or _COMPACTION_META_KEY not in messages[index].meta - ] + [index for index in range(start, end) if not (skip_compaction_notes and _is_compaction_note(messages[index]))] for start, end in spans ] @@ -316,7 +327,7 @@ def compact( self.omission_note.replace(_NUM_REMOVED_PLACEHOLDER, str(len(removable))), meta={ _COMPACTION_META_KEY: { - "strategy": "sliding_window", + "strategy": _STRATEGY, "removed_messages": len(removable), "kept_messages": len(kept), } diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 9745b7b1aa9..332332b71b6 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -4,9 +4,9 @@ import pytest -from haystack.dataclasses import ChatMessage, ChatRole +from haystack.dataclasses import ChatMessage, ChatRole, ToolCall from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans, _is_compaction_note from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import ( FakeCounter, @@ -24,6 +24,51 @@ COUNTER = FakeCounter() +class TestIsCompactionNote: + @pytest.mark.parametrize( + ("message", "expected"), + [ + pytest.param( + ChatMessage.from_user( + text="Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), + True, + id="note-this-strategy-left", + ), + # A pruned result carries the same meta key but is still part of the conversation, so it is not a note. + pytest.param( + ChatMessage.from_tool( + tool_result="[Tool result removed to free up context.]", + origin=ToolCall(tool_name="search", arguments={}, id="c1"), + meta={_COMPACTION_META_KEY: {"strategy": "tool_result_pruning", "original_tokens": 180}}, + ), + False, + id="tool-result-another-strategy-pruned", + ), + # Another strategy's note is not this one's to fold away, so it is left where it is. + pytest.param( + ChatMessage.from_user( + text="A summary of what came before.", meta={_COMPACTION_META_KEY: {"strategy": "summarization"}} + ), + False, + id="note-another-strategy-left", + ), + pytest.param( + ChatMessage.from_system(text="rules", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}}), + False, + id="system-message", + ), + pytest.param( + ChatMessage.from_user(text="odd", meta={_COMPACTION_META_KEY: "sliding_window"}), + False, + id="marker-that-is-not-a-mapping", + ), + ], + ) + def test_only_matches_sliding_window_omission_message(self, message, expected): + assert _is_compaction_note(message=message) is expected + + class TestHistoricalTurnSpans: def test_groups_each_user_message_with_its_assistant_steps_and_tool_results(self): messages = [ @@ -51,20 +96,36 @@ def test_only_returns_turns_within_the_requested_bounds(self): assert _historical_turn_spans(messages=messages, start=2, end=4) == [(2, 4)] def test_compaction_note_does_not_start_a_new_turn(self): - note = ChatMessage.from_user( - "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} - ) messages = [ + # Historical turns + ChatMessage.from_user( + "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), ChatMessage.from_user("task"), ChatMessage.from_assistant("first step"), - note, - ChatMessage.from_assistant("second step"), ChatMessage.from_user("next task"), + ChatMessage.from_assistant("second step"), ] - assert _historical_turn_spans(messages=messages, start=0, end=len(messages)) == [(0, 4), (4, 5)] + # The note is skipped which is why the first span starts at 1 + assert _historical_turn_spans(messages=messages, start=0, end=len(messages)) == [(1, 3), (3, 5)] class TestSlidingWindowCompactor: + def test_replaces_all_historical_turns(self): + messages = [ + # System + ChatMessage.from_system(text="rules"), + # Historical turn + ChatMessage.from_user(text="old question " * 100), + ChatMessage.from_assistant(text="old answer"), + # Current task + ChatMessage.from_user(text="current task"), + ] + compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + assert compacted == [messages[0], messages[-1]] + def test_replaces_oldest_agent_step(self): messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) @@ -124,42 +185,87 @@ def test_drops_all_historical_turns_and_oldest_agent_step(self): ) assert compacted == [system_message, current_task[0], *current_task[-2:]] - def test_folds_an_earlier_note_inside_a_retained_turn_and_counts_it_as_removed(self): - note = ChatMessage.from_user( - "Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} - ) + def test_replaces_earlier_note_and_oldest_historical_turn(self): messages = [ # System ChatMessage.from_system(text="rules"), + # The note an earlier compaction left, which sits at the top of the historical turns. + ChatMessage.from_user( + text="Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), # Historical ChatMessage.from_user(text="old question " * 200), ChatMessage.from_assistant(text="old answer"), ChatMessage.from_user(text="recent question"), ChatMessage.from_assistant(text="recent answer"), - note, # Current Task ChatMessage.from_user(text="current task"), ChatMessage.from_assistant(text="current step"), ] - # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn, so - # the turn holding the earlier note survives around it. + # Enough for the instructions, the recent turn, and the task with its step, but not the padded oldest turn. target_tokens = 30 compacted = SlidingWindowCompactor().compact( messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) assert compacted is not None - assert compacted == [messages[0], compacted[1], *messages[3:5], *messages[6:]] - # The earlier note is replaced by the new one rather than surviving alongside it, and it counts as removed. + # The earlier note goes with the turn it stood in front of, and the new note takes its place. + assert compacted == [messages[0], compacted[1], *messages[4:]] assert count_markers(messages=compacted) == 1 + # The earlier note is counted among the removed, alongside the two messages of the oldest turn. assert compacted[1].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 - def test_returns_none_when_the_conversation_already_fits(self): - assert ( - SlidingWindowCompactor().compact( - messages=fresh_conversation_with_two_steps(), target_tokens=100_000, token_counter=COUNTER - ) - is None + def test_replaces_earlier_note_and_oldest_agent_step(self): + messages = [ + # System + ChatMessage.from_system(text="rules"), + # Current Task + ChatMessage.from_user(text="current task"), + # The note an earlier compaction left, which sits right after the task when its own steps were trimmed. + ChatMessage.from_user( + text="Earlier messages were removed.", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}} + ), + tool_call("c1"), + tool_result(result="first result", call_id="c1"), + tool_call("c2"), + tool_result(result="second result", call_id="c2"), + ] + compacted = SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) + assert compacted is not None + # The earlier note goes with the step it stood in front of, and the new note takes its place. + assert compacted == [*messages[:2], compacted[2], *messages[5:]] + assert count_markers(messages=compacted) == 1 + # The earlier note is counted among the removed, alongside the two messages of the oldest step. + assert compacted[2].meta[_COMPACTION_META_KEY]["removed_messages"] == 3 + + def test_keeps_a_pruned_tool_result_inside_a_kept_turn(self): + messages = [ + # System + ChatMessage.from_system(text="rules"), + # Historical, dropped to make room + ChatMessage.from_user(text="ancient question " * 200), + ChatMessage.from_assistant(text="ancient answer"), + # Historical, kept + ChatMessage.from_user(text="old question"), + tool_call("old"), + # A result `ToolResultPruningCompactor` already pruned, which carries the same meta key as an omission note + # but is part of the conversation rather than standing in for removed history. + ChatMessage.from_tool( + tool_result="[Tool result removed to free up context.]", + origin=ToolCall(tool_name="search", arguments={}, id="old"), + meta={_COMPACTION_META_KEY: {"strategy": "tool_result_pruning", "original_tokens": 180}}, + ), + # Current Task + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Enough for the instructions, the kept turn, and the task with its step, but not the padded oldest turn. + target_tokens = 45 + compacted = SlidingWindowCompactor(omission_note=None).compact( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) + assert compacted is not None + # The pruned result is not an omission note, so it stays with the turn and its tool call keeps its answer. + assert compacted == [messages[0], *messages[3:]] def test_omission_note_can_be_turned_off(self): messages = fresh_conversation_with_two_steps() @@ -187,18 +293,25 @@ def test_omission_note_can_be_customized(self, note, expected): assert compacted[2].text == expected @pytest.mark.parametrize( - "messages", + ("messages", "target_tokens"), [ - pytest.param([], id="empty"), - pytest.param([ChatMessage.from_system(text="a"), ChatMessage.from_system(text="b")], id="only-system"), + # The conversation is already under the target, so there is nothing to do. + pytest.param(fresh_conversation_with_two_steps(), 100_000, id="conversation-already-fits"), + # Over the target, but everything that is left is protected, so there is nothing the compactor may remove. pytest.param( - [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="hi")], id="nothing-outside-window" + [ChatMessage.from_system(text="a"), ChatMessage.from_system(text="b")], SMALLEST, id="only-system" + ), + pytest.param( + [ChatMessage.from_system(text="rules"), ChatMessage.from_user(text="hi")], + SMALLEST, + id="only-system-and-task", ), ], ) - def test_returns_none_when_there_is_nothing_to_remove(self, messages): + def test_returns_none_when_there_is_nothing_to_remove(self, messages, target_tokens): assert ( - SlidingWindowCompactor().compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) is None + SlidingWindowCompactor().compact(messages=messages, target_tokens=target_tokens, token_counter=COUNTER) + is None ) def test_keeps_a_parallel_tool_call_together_with_all_results(self): @@ -232,33 +345,6 @@ def test_rejects_negative_min_keep_steps(self): with pytest.raises(ValueError, match="`min_keep_steps` must be at least 0"): SlidingWindowCompactor(min_keep_steps=-1) - def test_preserves_the_latest_user_task_when_compacting_input_history(self): - messages = [ - ChatMessage.from_system(text="rules"), - ChatMessage.from_user(text="old question " * 100), - ChatMessage.from_assistant(text="old answer"), - ChatMessage.from_user(text="current task"), - ] - compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( - messages=messages, target_tokens=SMALLEST, token_counter=COUNTER - ) - assert compacted == [messages[0], messages[-1]] - - def test_repeated_compaction_folds_the_previous_note(self): - compactor = SlidingWindowCompactor() - first = compactor.compact( - messages=fresh_conversation_with_two_steps(), target_tokens=SMALLEST, token_counter=COUNTER - ) - assert first is not None - # Simulate two more turns arriving on top of the already-compacted conversation. - grown = [*first, tool_call("c3"), tool_result(result="third result", call_id="c3")] - second = compactor.compact(messages=grown, target_tokens=SMALLEST, token_counter=COUNTER) - assert second is not None - # The original task stays anchored and the earlier compaction note is folded into the new one. - assert count_markers(messages=second) == 1 - assert second[0].text == "rules" - assert second[1].text == "start" - def test_keeping_no_steps_still_preserves_the_current_task(self): messages = fresh_conversation_with_two_steps() compacted = SlidingWindowCompactor(min_keep_steps=0, omission_note=None).compact( From bda7829529d7551a818f640e383efc4359cb375b Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 11:45:55 +0200 Subject: [PATCH 28/38] PR comments --- haystack/hooks/compaction/sliding_window.py | 91 +++++++++++++-------- 1 file changed, 56 insertions(+), 35 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 37227923091..f1d0cfbbd7f 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -101,15 +101,51 @@ def _index_groups( def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages at the given indices, in conversation order.""" + """Return the messages at the given indices, in the order the indices are given.""" return [messages[index] for index in indices] +def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages the given indices leave out, in conversation order.""" + left_out = set(indices) + return [message for index, message in enumerate(messages) if index not in left_out] + + def _flatten(groups: list[list[int]]) -> list[int]: """Join index groups into a single ordered list of indices.""" return [index for group in groups for index in group] +def _removable_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> tuple[list[list[int]], list[list[int]]]: + """ + Group the two stretches of conversation compaction is allowed to remove. + + :param messages: The full conversation, ordered oldest to newest. + :param system_end: The end of the leading system-message block. + :param task_index: The index of the user message anchoring the current task, or None when there is none. + :returns: Index groups for the complete historical turns preceding the current task, and index groups for the + current task's own Agent steps. Both are ordered oldest group first. A group is the unit of removal: it is + kept or removed entire, which is what keeps a tool call with its results and an assistant reply with the user + message it answers. + """ + # Steps belong to the current task, so they start after its anchor, or after the instructions when the + # conversation has no user message to anchor on. + step_start = (task_index + 1) if task_index is not None else system_end + step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start)) + + # An earlier compaction's note is left out of its turn, so keeping the turn folds that note into the note this + # compaction leaves behind. + historical_end = task_index if task_index is not None else system_end + historical_groups = _index_groups( + messages=messages, + spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), + skip_compaction_notes=True, + ) + return historical_groups, step_groups + + def _first_group_to_keep( messages: list[ChatMessage], groups: list[list[int]], available_tokens: int, token_counter: TokenCounter ) -> int: @@ -202,53 +238,38 @@ def _task_and_step_split( user message anchoring the current task when the task's own steps were removed. 3. Every message selected for removal. """ - # Find the leading system messages that contain the Agent instructions. + # The landmarks the split is built around: the Agent instructions, and the user message anchoring the current task. system_end = _leading_system_end(messages=messages) - - # Find the latest user message to use as the current task anchor. task_index = _latest_user_index(messages=messages) - task = [messages[task_index]] if task_index is not None else [] + task_indices = [task_index] if task_index is not None else [] - # Group the complete Agent steps that follow the current task anchor. - step_start_index = (task_index + 1) if task_index is not None else system_end - step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start_index)) + # The two stretches compaction may remove. A group is the unit of removal, so a turn or a step is never split. + historical_groups, step_groups = _removable_groups(messages=messages, system_end=system_end, task_index=task_index) - # Group the complete historical turns (i.e. user-assistant) that precede the current task. An earlier compaction's - # note is left out of its turn, so keeping the turn folds that note into the note this compaction leaves behind. - historical_end = task_index if task_index is not None else system_end - historical_groups = _index_groups( - messages=messages, - spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), - skip_compaction_notes=True, - ) - - # The Agent instructions and the current task are never removed, so they are paid for out of the target first. - protected_tokens = token_counter.count(messages=[*messages[:system_end], *task]) + # The instructions and the current task are never removed, so they come off the budget first. + protected = _messages_at(messages=messages, indices=[*range(system_end), *task_indices]) first_kept_turn, first_kept_step = _first_turn_and_step_to_keep( messages=messages, historical_groups=historical_groups, step_groups=step_groups, - available_tokens=target_tokens - protected_tokens, + available_tokens=target_tokens - token_counter.count(messages=protected), token_counter=token_counter, min_keep_steps=min_keep_steps, ) + + # What survives, laid out in conversation order: the instructions, the turns that fit, the task, then its steps. kept_turn_indices = _flatten(groups=historical_groups[first_kept_turn:]) kept_step_indices = _flatten(groups=step_groups[first_kept_step:]) - kept_turns = _messages_at(messages=messages, indices=kept_turn_indices) - kept_steps = _messages_at(messages=messages, indices=kept_step_indices) - kept = [*messages[:system_end], *kept_turns, *task, *kept_steps] - - # A message survives only by being protected or by falling in a group we are keeping; everything else goes. - kept_indices = {*range(system_end), *kept_turn_indices, *kept_step_indices} - if task_index is not None: - kept_indices.add(task_index) - removable = [message for index, message in enumerate(messages) if index not in kept_indices] - - # The note stands in for what was dropped, so it goes where the dropped messages used to sit. Either right after - # the leading system messages when the historical turns were trimmed, or right after the user message that anchors - # the current task when its own Agent steps were trimmed. Both positions are counted off the layout of `kept`. - note_index = system_end if first_kept_step == 0 else system_end + len(kept_turns) + len(task) - return kept, note_index, removable + kept_indices = [*range(system_end), *kept_turn_indices, *task_indices, *kept_step_indices] + + # The note stands in for what was dropped, so it goes where those messages used to sit. Either right after the + # instructions when the historical turns were trimmed, or right after the task anchor when its own steps were. + note_index = system_end if first_kept_step == 0 else system_end + len(kept_turn_indices) + len(task_indices) + return ( + _messages_at(messages=messages, indices=kept_indices), + note_index, + _messages_except(messages=messages, indices=kept_indices), + ) @_experimental From b4bda9ca314e44813267f11dece2d38559bb1af1 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 11:52:07 +0200 Subject: [PATCH 29/38] update docs pages --- .../agents-1/compaction.mdx | 4 +-- .../agents-1/compaction/compaction-hook.mdx | 2 +- .../compaction/sliding-window-compactor.mdx | 27 +++++++++---------- 3 files changed, 16 insertions(+), 17 deletions(-) diff --git a/docs-website/docs/pipeline-components/agents-1/compaction.mdx b/docs-website/docs/pipeline-components/agents-1/compaction.mdx index b8562acc820..792a8a231d8 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction.mdx @@ -72,7 +72,7 @@ Compactors receive the current messages, a target token count, and the same toke | Compactor | Strategy | Trade-off | | --- | --- | --- | -| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions, latest user task, and recent complete Agent steps while removing older history. | Fast and local, but discarded information is not summarized. | +| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions and latest user task, keeps earlier user/assistant turns whole while they fit, and trims the current task's own Agent steps only when that is not enough. | Fast and local, but discarded information is not summarized. | | [`ToolResultPruningCompactor`](compaction/tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while preserving tool-call/result structure. | Retains the shape of the run and recent results, but removes the content of pruned results. | ## Combining compaction strategies @@ -108,7 +108,7 @@ agent = Agent( ) ``` -If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes older complete Agent steps. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. +If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes earlier turns and then, if needed, the current task's oldest Agent steps. A result the pruning compactor already replaced with a placeholder stays with the turn it belongs to, so its tool call keeps an answer. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. ### Creating a custom compactor diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx index 0b8a27a66a4..db5394c1d74 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx @@ -90,7 +90,7 @@ The compactor controls what information survives: | Compactor | Strategy | | --- | --- | -| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and the most recent complete Agent steps, dropping older history. | +| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and as much complete recent conversation as fits, dropping earlier turns whole before it trims the task's own steps. | | [`ToolResultPruningCompactor`](tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while keeping recent results intact. | You can also implement the `Compactor` protocol for a custom strategy. See [Context Compaction](../compaction.mdx#creating-a-custom-compactor) for its requirements. diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx index 7de7773be7e..d27a358a733 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx @@ -2,12 +2,12 @@ title: "SlidingWindowCompactor" id: sliding-window-compactor slug: "/sliding-window-compactor" -description: "Use SlidingWindowCompactor to remove older Agent history while preserving the current task and recent complete steps." +description: "Use SlidingWindowCompactor to remove older Agent history while preserving the current task and as much recent complete conversation as fits." --- # SlidingWindowCompactor -`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as many recent complete Agent steps as the token target allows. +`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as much complete recent conversation as the token target allows. It removes earlier user/assistant turns as whole units first, and only trims the current task's own Agent steps when removing every earlier turn is not enough. :::warning[Experimental] @@ -51,17 +51,19 @@ compaction_hook = CompactionHook( ## How the sliding window is selected -The compactor divides a conversation into protected context, removable history, and recent Agent steps: +The compactor divides a conversation into protected context, earlier turns, and the current task's Agent steps: 1. It preserves all leading system messages as the Agent's instructions. 2. It preserves the latest user message as the current task. -3. It groups each assistant message and all immediately following tool-result messages into one complete Agent step. -4. Working backwards from the latest step, it keeps as many complete steps as fit within the target. -5. It replaces the removable middle history with an omission note, unless the note is disabled. +3. It groups the history before that task into complete user turns, each running from one user message up to the next. +4. It groups each assistant message and all immediately following tool-result messages into one complete Agent step. +5. Working backwards from the newest, it keeps as many whole earlier turns as fit within the target. +6. Only when the current task alone still exceeds the target does it begin removing that task's own steps, oldest first. +7. It replaces what it removed with an omission note, unless the note is disabled. -Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. +Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. Keeping whole turns likewise ensures an assistant reply is never retained without the user message it answers. -The target is a goal rather than a guarantee. Protected messages and the configured minimum number of recent steps take precedence when they already exceed the available token budget. +The target is a goal rather than a guarantee, and the conversation can end up above it rather than below. Leading system messages and the current task are never removed, and `min_keep_steps` holds on to the newest Agent steps whatever their size, so a long system prompt or a single large tool result can leave the conversation well over the target. ## Configuration @@ -76,16 +78,13 @@ The target is a goal rather than a guarantee. Protected messages and the configu An omission note tells the model that earlier context is missing. Without one, the shortened conversation can appear complete and the model may repeat work or behave as though it still has the removed information. -The compactor only inserts the note when it costs fewer tokens than the messages it replaces. Otherwise, it returns `None` and leaves the conversation unchanged. Repeated compactions fold the previous note into the newly removed block, leaving a single current note. +The note is left where the removed messages used to sit: directly after the leading system messages when only earlier turns were removed, and directly after the latest user message when the current task's own steps were removed. Repeated compactions fold an earlier note into the new one, so the conversation carries at most one. -Compaction metadata is stored on the note, including the strategy name and the numbers of removed messages, retained messages, and retained steps. +Compaction metadata is stored on the note, including the strategy name and the numbers of removed and retained messages. ## When the conversation is unchanged The compactor returns `None` without changing the conversation when: - The conversation already fits within `target_tokens`. -- There is no removable history outside the protected messages and retained steps. -- The configured omission note would cost at least as many tokens as the history it replaces. - -Setting `omission_note=None` removes the final condition because no replacement note needs to fit. +- There is no removable history outside the protected messages and the history it retained. From b535ef7a8b6a11420f1c69205ff1aff61dac09e9 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 12:54:35 +0200 Subject: [PATCH 30/38] fix docs --- .../agents-1/compaction.mdx | 4 ++-- .../agents-1/compaction/compaction-hook.mdx | 2 +- .../compaction/sliding-window-compactor.mdx | 12 +++++----- haystack/hooks/compaction/sliding_window.py | 12 +++++----- ...t-context-compaction-3258c08dec9d2b34.yaml | 24 +++++++++---------- 5 files changed, 27 insertions(+), 27 deletions(-) diff --git a/docs-website/docs/pipeline-components/agents-1/compaction.mdx b/docs-website/docs/pipeline-components/agents-1/compaction.mdx index 792a8a231d8..8a76a1185d6 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction.mdx @@ -72,7 +72,7 @@ Compactors receive the current messages, a target token count, and the same toke | Compactor | Strategy | Trade-off | | --- | --- | --- | -| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions and latest user task, keeps earlier user/assistant turns whole while they fit, and trims the current task's own Agent steps only when that is not enough. | Fast and local, but discarded information is not summarized. | +| [`SlidingWindowCompactor`](compaction/sliding-window-compactor.mdx) | Preserves the Agent's instructions and latest user task, keeps complete historical turns while they fit, and trims the current task's own Agent steps only when that is not enough. | Fast and local, but discarded information is not summarized. | | [`ToolResultPruningCompactor`](compaction/tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while preserving tool-call/result structure. | Retains the shape of the run and recent results, but removes the content of pruned results. | ## Combining compaction strategies @@ -108,7 +108,7 @@ agent = Agent( ) ``` -If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes earlier turns and then, if needed, the current task's oldest Agent steps. A result the pruning compactor already replaced with a placeholder stays with the turn it belongs to, so its tool call keeps an answer. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. +If pruning brings the updated context below `compact_at`, the sliding-window hook does nothing. If pruning returns `None` because no eligible results remain, or it shortens the context without getting below the trigger, the sliding window removes historical turns and then, if needed, the current task's oldest Agent steps. A result the pruning compactor already replaced with a placeholder stays with the historical turn it belongs to, so its tool call keeps an answer. Configure both hooks for the same model context window and compatible token counters so they make decisions from comparable estimates. ### Creating a custom compactor diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx index db5394c1d74..8837f6d2b51 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/compaction-hook.mdx @@ -90,7 +90,7 @@ The compactor controls what information survives: | Compactor | Strategy | | --- | --- | -| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and as much complete recent conversation as fits, dropping earlier turns whole before it trims the task's own steps. | +| [`SlidingWindowCompactor`](sliding-window-compactor.mdx) | Keeps the current task and as much complete recent conversation as fits, removing complete historical turns before it trims the task's own steps. | | [`ToolResultPruningCompactor`](tool-result-pruning-compactor.mdx) | Replaces older, large tool results with short placeholders while keeping recent results intact. | You can also implement the `Compactor` protocol for a custom strategy. See [Context Compaction](../compaction.mdx#creating-a-custom-compactor) for its requirements. diff --git a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx index d27a358a733..e66e5c915fe 100644 --- a/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx +++ b/docs-website/docs/pipeline-components/agents-1/compaction/sliding-window-compactor.mdx @@ -7,7 +7,7 @@ description: "Use SlidingWindowCompactor to remove older Agent history while pre # SlidingWindowCompactor -`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as much complete recent conversation as the token target allows. It removes earlier user/assistant turns as whole units first, and only trims the current task's own Agent steps when removing every earlier turn is not enough. +`SlidingWindowCompactor` removes older conversation history while preserving the Agent's instructions, current task, and as much complete recent conversation as the token target allows. It removes complete historical turns first, and only trims the current task's own Agent steps when removing every historical turn is not enough. :::warning[Experimental] @@ -51,17 +51,17 @@ compaction_hook = CompactionHook( ## How the sliding window is selected -The compactor divides a conversation into protected context, earlier turns, and the current task's Agent steps: +The compactor divides a conversation into protected context, historical turns, and the current task's Agent steps: 1. It preserves all leading system messages as the Agent's instructions. 2. It preserves the latest user message as the current task. -3. It groups the history before that task into complete user turns, each running from one user message up to the next. +3. It groups the history before that task into complete historical turns, each running from one user message up to the next. 4. It groups each assistant message and all immediately following tool-result messages into one complete Agent step. -5. Working backwards from the newest, it keeps as many whole earlier turns as fit within the target. +5. Working backwards from the newest, it keeps as many complete historical turns as fit within the target. 6. Only when the current task alone still exceeds the target does it begin removing that task's own steps, oldest first. 7. It replaces what it removed with an omission note, unless the note is disabled. -Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. Keeping whole turns likewise ensures an assistant reply is never retained without the user message it answers. +Keeping complete steps ensures that an assistant tool call is not separated from its results, including batches of parallel tool calls. Incomplete tool-call exchanges are rejected by chat-completion providers. Historical turns are kept or removed in full for the same reason: an assistant reply is never retained without the user message it answers. The target is a goal rather than a guarantee, and the conversation can end up above it rather than below. Leading system messages and the current task are never removed, and `min_keep_steps` holds on to the newest Agent steps whatever their size, so a long system prompt or a single large tool result can leave the conversation well over the target. @@ -78,7 +78,7 @@ The target is a goal rather than a guarantee, and the conversation can end up ab An omission note tells the model that earlier context is missing. Without one, the shortened conversation can appear complete and the model may repeat work or behave as though it still has the removed information. -The note is left where the removed messages used to sit: directly after the leading system messages when only earlier turns were removed, and directly after the latest user message when the current task's own steps were removed. Repeated compactions fold an earlier note into the new one, so the conversation carries at most one. +The note is left where the removed messages used to sit: directly after the leading system messages when only historical turns were removed, and directly after the latest user message when the current task's own steps were removed. Repeated compactions fold an earlier note into the new one, so the conversation carries at most one. Compaction metadata is stored on the note, including the strategy name and the numbers of removed and retained messages. diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index f1d0cfbbd7f..bd780e83d10 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -205,7 +205,7 @@ def _first_turn_and_step_to_keep( # The newest steps are kept regardless of the budget. return len(historical_groups), min(first_kept_step, max(len(step_groups) - min_keep_steps, 0)) - # The whole current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. + # The entire current task fits, so every step stays and the rest of the budget goes on the newest turns that fit. first_kept_turn = _first_group_to_keep( messages=messages, groups=historical_groups, @@ -221,7 +221,7 @@ def _task_and_step_split( """ Split a conversation into the messages to keep and the messages to remove. - Leading system messages and the latest real user message are always kept. Historical user turns are kept whole when + Leading system messages and the latest real user message are always kept. Historical turns are kept in full when they fit. If the current task itself exceeds the available budget, its oldest Agent steps are removed one at a time while keeping each assistant message together with its tool results. @@ -277,13 +277,13 @@ class SlidingWindowCompactor(Compactor): """ Keeps the Agent's instructions, current task, and as much complete recent conversation as the target allows. - Leading system messages and the latest user message are protected. Earlier user/assistant turns are kept when they - fit, and the current task's history is kept in complete Agent steps, where a step is an assistant message together + Leading system messages and the latest user message are protected. Historical turns are kept in full when they fit, + and the current task's history is kept in complete Agent steps, where a step is an assistant message together with all immediately following tool results. An `omission_note` is left where the removed messages used to sit: directly after the leading system messages when - only earlier turns were removed, and directly after the latest user message when the current task's own steps were - removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement. + only historical turns were removed, and directly after the latest user message when the current task's own steps + were removed. Only one note is ever present, since a later compaction folds an earlier note into its replacement. ```python from haystack.components.agents import Agent diff --git a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml index 6caa6298451..4f6955996ca 100644 --- a/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml +++ b/releasenotes/notes/add-agent-context-compaction-3258c08dec9d2b34.yaml @@ -5,12 +5,12 @@ features: shortens the conversation when it reaches a configured fraction of the model's context window. The first built-in strategy, ``SlidingWindowCompactor``, preserves leading system messages, the latest user task, - and as much complete recent conversation as the target allows. It removes earlier user/assistant turns as whole - units first. Only when removing all earlier turns is insufficient does it remove individual Agent steps from the - current task. It replaces removed history with a short omission note, left where the removed messages used to sit: - directly after the leading system messages when only earlier turns were removed, and directly after the latest user - message when the current task's own steps were removed. Only one note is ever present, because a later compaction - folds an earlier one into its replacement. + and as much complete recent conversation as the target allows. It removes complete historical turns first, and only + when removing every historical turn is insufficient does it remove individual Agent steps from the current task. It + replaces removed history with a short omission note, left where the removed messages used to sit: directly after + the leading system messages when only historical turns were removed, and directly after the latest user message + when the current task's own steps were removed. Only one note is ever present, because a later compaction folds an + earlier one into its replacement. .. code-block:: python @@ -34,12 +34,12 @@ features: tool schemas. Leave headroom above ``compact_at`` for the next reply and its tool results. ``SlidingWindowCompactor`` treats an assistant message and its following tool results as one step, so a tool call is - never separated from its results. Historical turns are also kept whole, so an assistant reply is not retained - without the user message it answers. It can also land above the requested target rather than under it, because - leading system messages and the current task are never removed and ``min_keep_steps`` holds on to the newest Agent - steps whatever their size, so a long system prompt or one large tool result can leave the conversation well over - the target. Compaction is lossy: removed messages cannot be recovered or summarized by this strategy. Implement - the ``Compactor`` protocol to provide a custom strategy. + never separated from its results. Historical turns are likewise kept or removed in full, so an assistant reply is + not retained without the user message it answers. It can also land above the requested target rather than under it, + because leading system messages and the current task are never removed and ``min_keep_steps`` holds on to the + newest Agent steps whatever their size, so a long system prompt or one large tool result can leave the conversation + well over the target. Compaction is lossy: removed messages cannot be recovered or summarized by this strategy. + Implement the ``Compactor`` protocol to provide a custom strategy. ``CompactionHook`` and ``SlidingWindowCompactor`` emit an ``ExperimentalWarning`` and may change without a deprecation cycle. From 1a971d5e705c09fe4f2b81c9564ed2091a963391 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 12:59:39 +0200 Subject: [PATCH 31/38] changes --- haystack/hooks/compaction/sliding_window.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index bd780e83d10..c916e5958ef 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -48,13 +48,7 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: def _is_compaction_note(message: ChatMessage) -> bool: - """ - Whether a message is an omission note this strategy left in place of removed history. - - Every compactor marks what it produces with the same meta key, including the tool results - `ToolResultPruningCompactor` rewrites into a placeholder. Matching on the role and the strategy keeps those out: - they are still part of the conversation and have to travel with the turn they belong to. - """ + """Whether a message is an omission note this strategy left in place of removed history.""" marker = message.meta.get(_COMPACTION_META_KEY) return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY @@ -72,8 +66,8 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. """ - # Compaction notes use the user role for provider compatibility, but they do not begin a new conversation turn. - # Ignoring marked messages here also lets a subsequent compaction fold an old note into its replacement. + # Reject any user-role message an earlier compaction produced, whichever strategy made it: none of them are user + # requests, so none of them begin a turn. Leaving them out also lets this compaction fold an old note away. user_indices = [ index for index in range(start, end) From 36fa8069a9036c7ce9881f6223fc967d48a2de12 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:22:31 +0200 Subject: [PATCH 32/38] First pass at adding summarization compactor --- haystack/components/generators/chat/utils.py | 46 ++ haystack/hooks/compaction/__init__.py | 2 + haystack/hooks/compaction/sliding_window.py | 100 +-- haystack/hooks/compaction/summarization.py | 588 ++++++++++++++++++ haystack/hooks/compaction/utils.py | 62 ++ pydoc/hooks_api.yml | 2 +- ...marization-compactor-91b6be6855f478df.yaml | 6 + test/components/generators/chat/test_utils.py | 46 ++ test/hooks/compaction/test_sliding_window.py | 4 +- test/hooks/compaction/test_summarization.py | 249 ++++++++ 10 files changed, 1017 insertions(+), 88 deletions(-) create mode 100644 haystack/components/generators/chat/utils.py create mode 100644 haystack/hooks/compaction/summarization.py create mode 100644 releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml create mode 100644 test/components/generators/chat/test_utils.py create mode 100644 test/hooks/compaction/test_summarization.py diff --git a/haystack/components/generators/chat/utils.py b/haystack/components/generators/chat/utils.py new file mode 100644 index 00000000000..9a1a6d5b824 --- /dev/null +++ b/haystack/components/generators/chat/utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.components.generators.chat.types import ChatGenerator + +_CHAT_COMPLETIONS_GENERATORS = {"OpenAIChatGenerator", "AzureOpenAIChatGenerator"} +_RESPONSES_GENERATORS = {"OpenAIResponsesChatGenerator", "AzureOpenAIResponsesChatGenerator"} + + +def _generator_output_token_limit_key(chat_generator: ChatGenerator) -> str | None: + """Return the output-token parameter used by a known built-in Chat Generator.""" + class_names = {cls.__name__ for cls in type(chat_generator).__mro__} + if class_names & _RESPONSES_GENERATORS: + return "max_output_tokens" + if class_names & _CHAT_COMPLETIONS_GENERATORS: + return "max_completion_tokens" + return None + + +def _resolve_output_token_limit(chat_generator: ChatGenerator, default_limit: int) -> tuple[int, dict[str, Any] | None]: + """ + Resolve an effective output-token limit and runtime kwargs for a Chat Generator. + + A recognized limit configured directly on a built-in generator wins and is not repeated at runtime. When the + built-in generator has no configured limit, the default is returned as its provider-specific runtime setting. + Unknown generators receive no runtime setting because the ChatGenerator protocol does not standardize the key. + + :param chat_generator: The generator whose output should be limited. + :param default_limit: The positive fallback output-token limit. + :returns: The effective limit and provider-specific runtime generation kwargs, or None for no runtime kwargs. + """ + limit_key = _generator_output_token_limit_key(chat_generator=chat_generator) + if limit_key is None: + return default_limit, None + + configured = getattr(chat_generator, "generation_kwargs", None) + if isinstance(configured, dict) and limit_key in configured: + value = configured[limit_key] + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value, None + # The generator owns this setting. Do not silently replace an invalid value; let it report the problem. + return default_limit, None + return default_limit, {limit_key: default_limit} diff --git a/haystack/hooks/compaction/__init__.py b/haystack/hooks/compaction/__init__.py index 878515a6d1f..fb0ae0a1630 100644 --- a/haystack/hooks/compaction/__init__.py +++ b/haystack/hooks/compaction/__init__.py @@ -10,6 +10,7 @@ _import_structure = { "hooks": ["CompactionHook"], "sliding_window": ["SlidingWindowCompactor"], + "summarization": ["SummarizationCompactor"], "tool_result_pruning": ["ToolResultPruningCompactor"], "types": ["Compactor"], } @@ -17,6 +18,7 @@ if TYPE_CHECKING: from .hooks import CompactionHook as CompactionHook from .sliding_window import SlidingWindowCompactor as SlidingWindowCompactor + from .summarization import SummarizationCompactor as SummarizationCompactor from .tool_result_pruning import ToolResultPruningCompactor as ToolResultPruningCompactor from .types import Compactor as Compactor else: diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index c916e5958ef..b3b90c44604 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -7,7 +7,15 @@ from haystack.core.serialization import default_to_dict from haystack.dataclasses import ChatMessage, ChatRole from haystack.hooks.compaction.types import Compactor -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _agent_step_spans +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _current_step_groups, + _historical_turn_groups, + _is_compaction_message, + _latest_user_index, + _leading_system_end, + _messages_except, +) from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental @@ -23,75 +31,9 @@ ) -def _leading_system_end(messages: list[ChatMessage]) -> int: - """Return the end of the leading system-message block.""" - for index, message in enumerate(messages): - # Find the first non-leading system message or a system message produced by compaction - if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: - return index - return len(messages) - - -def _latest_user_index(messages: list[ChatMessage]) -> int | None: - """ - Return the latest user message not produced by compaction. - - :param messages: The conversation to analyze, oldest to newest. - """ - # We loop backwards to find the latest user message - for index in reversed(range(len(messages))): - message = messages[index] - # Find the latest user message that was not produced by a previous compaction - if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: - return index - return None - - def _is_compaction_note(message: ChatMessage) -> bool: """Whether a message is an omission note this strategy left in place of removed history.""" - marker = message.meta.get(_COMPACTION_META_KEY) - return message.is_from(role=ChatRole.USER) and isinstance(marker, dict) and marker.get("strategy") == _STRATEGY - - -def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: - """ - Return spans for complete user turns in a bounded section of conversation history. - - Each turn begins with a real user message and continues up to, but does not include, the next real user message. - This groups a user's request with every assistant step and tool result produced in response to it. - - :param messages: The full conversation to analyze, ordered oldest to newest. - :param start: The inclusive index at which to begin looking for historical turns. - :param end: The exclusive index at which to stop. This is normally the current task's user-message index. - :returns: Ordered `(start_index, end_index)` pairs for each complete historical turn. Both indices refer to - `messages`, and `end_index` is exclusive, so a returned pair can be used directly as `messages[start:end]`. - """ - # Reject any user-role message an earlier compaction produced, whichever strategy made it: none of them are user - # requests, so none of them begin a turn. Leaving them out also lets this compaction fold an old note away. - user_indices = [ - index - for index in range(start, end) - if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta - ] - - # A real user message closes the preceding turn and starts the next one. The final historical turn extends to the - # supplied boundary, which is typically where the protected current task begins. - return [ - (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) - for position, index in enumerate(user_indices) - ] - - -def _index_groups( - messages: list[ChatMessage], spans: list[tuple[int, int]], skip_compaction_notes: bool = False -) -> list[list[int]]: - """ - Expand each span into the message indices it covers, optionally dropping messages an earlier compaction produced. - """ - return [ - [index for index in range(start, end) if not (skip_compaction_notes and _is_compaction_note(messages[index]))] - for start, end in spans - ] + return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER) def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: @@ -99,12 +41,6 @@ def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMe return [messages[index] for index in indices] -def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages the given indices leave out, in conversation order.""" - left_out = set(indices) - return [message for index, message in enumerate(messages) if index not in left_out] - - def _flatten(groups: list[list[int]]) -> list[int]: """Join index groups into a single ordered list of indices.""" return [index for group in groups for index in group] @@ -124,19 +60,13 @@ def _removable_groups( kept or removed entire, which is what keeps a tool call with its results and an assistant reply with the user message it answers. """ - # Steps belong to the current task, so they start after its anchor, or after the instructions when the - # conversation has no user message to anchor on. - step_start = (task_index + 1) if task_index is not None else system_end - step_groups = _index_groups(messages=messages, spans=_agent_step_spans(messages=messages, start=step_start)) - # An earlier compaction's note is left out of its turn, so keeping the turn folds that note into the note this # compaction leaves behind. - historical_end = task_index if task_index is not None else system_end - historical_groups = _index_groups( - messages=messages, - spans=_historical_turn_spans(messages=messages, start=system_end, end=historical_end), - skip_compaction_notes=True, - ) + historical_groups = [ + [index for index in group if not _is_compaction_note(message=messages[index])] + for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + ] + step_groups = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) return historical_groups, step_groups diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py new file mode 100644 index 00000000000..b466a7ee798 --- /dev/null +++ b/haystack/hooks/compaction/summarization.py @@ -0,0 +1,588 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Awaitable, Callable +from typing import Any + +from haystack import logging +from haystack.components.generators.chat.types import ChatGenerator +from haystack.components.generators.chat.utils import _resolve_output_token_limit +from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict +from haystack.dataclasses import ChatMessage +from haystack.hooks.compaction.types import Compactor +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _current_step_groups, + _historical_turn_groups, + _is_compaction_message, + _latest_user_index, + _leading_system_end, + _messages_except, +) +from haystack.token_counters import TokenCounter +from haystack.token_counters.utils import _render_message +from haystack.utils.async_utils import _execute_component_async +from haystack.utils.deserialization import deserialize_component_inplace +from haystack.utils.experimental import _experimental + +logger = logging.getLogger(__name__) + +_STRATEGY = "summarization" + +_DEFAULT_SUMMARY_INSTRUCTION = """You are compacting part of a conversation between a user and an AI agent so the \ +agent can keep working with fewer tokens. Write a self-contained summary that preserves: +- The user's goal, requirements, constraints, and preferences. +- Decisions and the reasoning behind them. +- Work already completed and important tool results. +- Exact file paths, URLs, identifiers, and references to stored data. +- Unresolved work and the immediate next step. + +Fold any existing blocks into one summary. Record only what the conversation shows. Do not \ +infer or add advice. Use plain prose or short bullets, and do not address the user.""" + + +def _indices(messages: list[ChatMessage], start: int, end: int, *, summaries: bool) -> list[int]: + """Return summary or non-summary indices in a bounded part of a conversation.""" + return [ + index + for index in range(start, end) + if _is_compaction_message(message=messages[index], strategy=_STRATEGY) is summaries + ] + + +def _raw_historical_turn_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> list[list[int]]: + """Return shared historical-turn groups with this strategy's summaries filtered out.""" + return [ + [index for index in group if not _is_compaction_message(message=messages[index], strategy=_STRATEGY)] + for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + ] + + +def _groups_to_summarize( + messages: list[ChatMessage], + groups: list[list[int]], + target_tokens: int, + summary_budget: int, + token_counter: TokenCounter, +) -> list[int]: + """Select the fewest oldest groups that should make room for a summary of the configured size.""" + selected: list[int] = [] + for group in groups: + selected.extend(group) + if ( + token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + summary_budget + <= target_tokens + ): + break + return selected + + +def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMessage: + """Build a marked summary message.""" + body = f"\n{text.strip()}\n" + meta = {_COMPACTION_META_KEY: {"strategy": _STRATEGY, "summarized_messages": summarized_messages, "source": source}} + return ChatMessage.from_user(text=body, meta=meta) + + +def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: + """Replace possibly non-contiguous selected messages with one summary at their first position.""" + selected = set(indices) + insertion_index = min(indices) + compacted: list[ChatMessage] = [] + for index, message in enumerate(messages): + if index == insertion_index: + compacted.append(summary) + if index not in selected: + compacted.append(message) + return compacted + + +@_experimental +class SummarizationCompactor(Compactor): + """ + Condenses old historical turns first, then old steps from the Agent's current task. + + Leading system messages and the latest real user message are always kept. Historical turns are summarized in full, + oldest first. Summaries normally accumulate so they are not repeatedly rewritten; if every historical turn has + already been summarized and more space is needed, those historical summaries are folded into one before any + current-task steps are summarized. An assistant message and all immediately following tool results form one step, + so tool calls are never separated from their results. + + Each summary is requested within `max_summary_tokens`. Built-in OpenAI Chat Completions and Responses generators + receive the corresponding runtime output limit unless they already configure one in their `generation_kwargs`, in + which case the generator's setting wins. Other generators receive the same limit as prompt guidance and the actual + result is measured before it is accepted. + + ```python + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIResponsesChatGenerator + from haystack.hooks.compaction import CompactionHook, SummarizationCompactor + + summary_generator = OpenAIResponsesChatGenerator(model="gpt-5.4-nano") + hook = CompactionHook( + compactor=SummarizationCompactor(chat_generator=summary_generator), + context_window=400_000, + compact_at=0.7, + compact_to=0.4, + ) + agent = Agent(chat_generator=agent_generator, tools=[web_search], hooks={"before_llm": [hook]}) + ``` + """ + + def __init__( + self, + chat_generator: ChatGenerator, + *, + min_keep_steps: int = 1, + max_summary_tokens: int = 1024, + summary_instruction: str | None = None, + raise_on_failure: bool = False, + ) -> None: + """ + Initialize the compactor. + + :param chat_generator: The Chat Generator used to write summaries. Its configured output-token limit takes + precedence over `max_summary_tokens` when the generator exposes a recognized setting. + :param min_keep_steps: The fewest complete recent Agent steps to keep, even when they exceed the target. + :param max_summary_tokens: The output-token budget reserved for each summary. Known built-in generators receive + the corresponding runtime generation setting unless one is already configured on the generator. + :param summary_instruction: An instruction replacing the built-in summary prompt. + :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is + logged and any successful partial compaction is returned. + """ + if min_keep_steps < 0: + raise ValueError(f"`min_keep_steps` must be at least 0, got {min_keep_steps}.") + if max_summary_tokens < 1: + raise ValueError(f"`max_summary_tokens` must be a positive number of tokens, got {max_summary_tokens}.") + self.chat_generator = chat_generator + self.min_keep_steps = min_keep_steps + self.max_summary_tokens = max_summary_tokens + self.summary_instruction = summary_instruction or _DEFAULT_SUMMARY_INSTRUCTION + self.raise_on_failure = raise_on_failure + + def compact( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Return a progressively summarized conversation, or None when no useful reduction is possible. + + :param messages: The conversation to compact, ordered oldest to newest. + :param target_tokens: The token budget the compacted messages should aim to fit. + :param token_counter: The counter used both to plan compaction and verify generated summaries. + :returns: A smaller replacement conversation, or None when nothing was reduced. + """ + if token_counter.count(messages=messages) <= target_tokens: + return None + budget, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + + def generate(prompt: list[ChatMessage]) -> dict[str, Any]: + kwargs: dict[str, Any] = {"messages": prompt} + if generation_kwargs is not None: + kwargs["generation_kwargs"] = generation_kwargs + return self.chat_generator.run(**kwargs) + + return self._compact( + original=messages, + target_tokens=target_tokens, + token_counter=token_counter, + budget=budget, + generate=generate, + ) + + async def compact_async( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Asynchronously return a progressively summarized conversation. + + :param messages: The conversation to compact, ordered oldest to newest. + :param target_tokens: The token budget the compacted messages should aim to fit. + :param token_counter: The counter used both to plan compaction and verify generated summaries. + :returns: A smaller replacement conversation, or None when nothing was reduced. + """ + if token_counter.count(messages=messages) <= target_tokens: + return None + budget, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + + async def generate(prompt: list[ChatMessage]) -> dict[str, Any]: + kwargs: dict[str, Any] = {"messages": prompt} + if generation_kwargs is not None: + kwargs["generation_kwargs"] = generation_kwargs + return await _execute_component_async(component_instance=self.chat_generator, **kwargs) + + return await self._compact_async( + original=messages, + target_tokens=target_tokens, + token_counter=token_counter, + budget=budget, + generate=generate, + ) + + def _prompt(self, messages: list[ChatMessage], budget: int) -> list[ChatMessage]: + """Build the bounded summarization instruction and rendered source transcript.""" + transcript = "\n".join(_render_message(message=message) for message in messages) + instruction = ( + f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately {budget} tokens. " + "Prioritize completeness within that limit so the response is not cut off." + ) + return [ + ChatMessage.from_system(text=instruction), + ChatMessage.from_user(text=f"\n{transcript}\n"), + ] + + def _apply_result( + self, + messages: list[ChatMessage], + indices: list[int], + source: str, + result: dict[str, Any], + token_counter: TokenCounter, + ) -> list[ChatMessage]: + """Validate a generator reply and replace its source messages when the result is smaller.""" + replies = result.get("replies") or [] + text = replies[-1].text if replies else None + if not text or not text.strip(): + raise RuntimeError("The Chat Generator returned no text to use as a conversation summary.") + summary = _summary_message(text=text, summarized_messages=len(indices), source=source) + compacted = _replace_indices(messages=messages, indices=indices, summary=summary) + before = token_counter.count(messages=messages) + after = token_counter.count(messages=compacted) + if after >= before: + raise RuntimeError( + f"The generated summary did not reduce the conversation size ({before} tokens before and {after} " + "tokens after)." + ) + return compacted + + def _attempt( + self, + messages: list[ChatMessage], + indices: list[int], + source: str, + budget: int, + token_counter: TokenCounter, + generate: Callable[[list[ChatMessage]], dict[str, Any]], + ) -> list[ChatMessage] | None: + """Attempt one synchronous summary, applying the configured failure policy.""" + try: + result = generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) + return self._apply_result( + messages=messages, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + if self.raise_on_failure: + raise + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) + return None + + async def _attempt_async( + self, + messages: list[ChatMessage], + indices: list[int], + source: str, + budget: int, + token_counter: TokenCounter, + generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], + ) -> list[ChatMessage] | None: + """Attempt one asynchronous summary, applying the configured failure policy.""" + try: + result = await generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) + return self._apply_result( + messages=messages, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + if self.raise_on_failure: + raise + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) + return None + + def _compact( + self, + original: list[ChatMessage], + target_tokens: int, + token_counter: TokenCounter, + budget: int, + generate: Callable[[list[ChatMessage]], dict[str, Any]], + ) -> list[ChatMessage] | None: + """Run synchronous historical, consolidation, and current-step compaction tiers in order.""" + working = list(original) + + # First replace the fewest oldest raw historical turns expected to reach the target. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + groups = [ + group + for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) + if group + ] + if not groups: + break + selected = _groups_to_summarize( + messages=working, + groups=groups, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = self._attempt( + messages=working, + indices=selected, + source="historical_turns", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # If all historical turns are summaries and the target is still unmet, fold them into one summary. + if token_counter.count(messages=working) > target_tokens: + summaries = self._summary_indices(messages=working, source="historical_summaries") + if len(summaries) > 1: + compacted = self._attempt( + messages=working, + indices=summaries, + source="historical_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + if not eligible: + break + + # Fold accumulated current-task summaries before consuming more raw agent steps. + summaries = self._summary_indices(messages=working, source="current_task_summaries") + if len(summaries) > 1: + compacted = self._attempt( + messages=working, + indices=summaries, + source="current_task_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + selected = _groups_to_summarize( + messages=working, + groups=eligible, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = self._attempt( + messages=working, + indices=selected, + source="current_task_steps", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + break + working = compacted + return self._result(original=original, working=working, token_counter=token_counter) + + async def _compact_async( + self, + original: list[ChatMessage], + target_tokens: int, + token_counter: TokenCounter, + budget: int, + generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], + ) -> list[ChatMessage] | None: + """Run asynchronous historical, consolidation, and current-step compaction tiers in order.""" + working = list(original) + + # First replace the fewest oldest raw historical turns expected to reach the target. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + groups = [ + group + for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) + if group + ] + if not groups: + break + selected = _groups_to_summarize( + messages=working, + groups=groups, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = await self._attempt_async( + messages=working, + indices=selected, + source="historical_turns", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # If all historical turns are summaries and the target is still unmet, fold them into one summary. + if token_counter.count(messages=working) > target_tokens: + summaries = self._summary_indices(messages=working, source="historical_summaries") + if len(summaries) > 1: + compacted = await self._attempt_async( + messages=working, + indices=summaries, + source="historical_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. + while token_counter.count(messages=working) > target_tokens: + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + if not eligible: + break + + # Fold accumulated current-task summaries before consuming more raw agent steps. + summaries = self._summary_indices(messages=working, source="current_task_summaries") + if len(summaries) > 1: + compacted = await self._attempt_async( + messages=working, + indices=summaries, + source="current_task_summaries", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + return self._result(original=original, working=working, token_counter=token_counter) + working = compacted + + # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. + system_end = _leading_system_end(messages=working) + task_index = _latest_user_index(messages=working) + step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) + eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] + selected = _groups_to_summarize( + messages=working, + groups=eligible, + target_tokens=target_tokens, + summary_budget=budget, + token_counter=token_counter, + ) + compacted = await self._attempt_async( + messages=working, + indices=selected, + source="current_task_steps", + budget=budget, + token_counter=token_counter, + generate=generate, + ) + if compacted is None: + break + working = compacted + return self._result(original=original, working=working, token_counter=token_counter) + + def _summary_indices(self, messages: list[ChatMessage], source: str) -> list[int]: + """Return historical or current-task summary indices based on their conversation position.""" + system_end = _leading_system_end(messages=messages) + task_index = _latest_user_index(messages=messages) + end = task_index if task_index is not None else len(messages) + if source == "historical_summaries": + return _indices(messages=messages, start=system_end, end=end, summaries=True) + start = task_index + 1 if task_index is not None else system_end + return _indices(messages=messages, start=start, end=len(messages), summaries=True) + + @staticmethod + def _result( + original: list[ChatMessage], working: list[ChatMessage], token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """Return partial or complete progress only when it reduced the original conversation.""" + if token_counter.count(messages=working) < token_counter.count(messages=original): + return working + return None + + def warm_up(self) -> None: + """Warm up the Chat Generator that writes summaries.""" + if hasattr(self.chat_generator, "warm_up"): + self.chat_generator.warm_up() + + async def warm_up_async(self) -> None: + """Warm up the Chat Generator on the serving event loop.""" + warm_up_async = getattr(self.chat_generator, "warm_up_async", None) + if warm_up_async is not None: + await warm_up_async() + elif hasattr(self.chat_generator, "warm_up"): + self.chat_generator.warm_up() + + def close(self) -> None: + """Release the Chat Generator's resources.""" + if hasattr(self.chat_generator, "close"): + self.chat_generator.close() + + async def close_async(self) -> None: + """Release the Chat Generator's resources.""" + close_async = getattr(self.chat_generator, "close_async", None) + if close_async is not None: + await close_async() + elif hasattr(self.chat_generator, "close"): + self.chat_generator.close() + + def to_dict(self) -> dict[str, Any]: + """Serialize the compactor and its Chat Generator.""" + return default_to_dict( + self, + chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"), + min_keep_steps=self.min_keep_steps, + max_summary_tokens=self.max_summary_tokens, + summary_instruction=self.summary_instruction, + raise_on_failure=self.raise_on_failure, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "SummarizationCompactor": + """Deserialize the compactor and reconstruct its Chat Generator.""" + init_params = data.get("init_parameters", {}) + if init_params.get("chat_generator") is not None: + deserialize_component_inplace(data=init_params, key="chat_generator") + return default_from_dict(cls=cls, data=data) diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 6a6d1402c52..99b0fe79d8e 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -10,6 +10,36 @@ _COMPACTION_META_KEY = "context_compaction" +def _leading_system_end(messages: list[ChatMessage]) -> int: + """Return the end of the leading system-message block, excluding system messages created by compaction.""" + for index, message in enumerate(messages): + if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta: + return index + return len(messages) + + +def _latest_user_index(messages: list[ChatMessage]) -> int | None: + """Return the latest user message not produced by compaction.""" + for index in reversed(range(len(messages))): + message = messages[index] + if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta: + return index + return None + + +def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages the given indices leave out, in conversation order.""" + left_out = set(indices) + return [message for index, message in enumerate(messages) if index not in left_out] + + +def _is_compaction_message(message: ChatMessage, strategy: str, role: ChatRole | None = None) -> bool: + """Return whether a message was produced by a compaction strategy and optionally has the requested role.""" + marker = message.meta.get(_COMPACTION_META_KEY) + has_role = role is None or message.is_from(role=role) + return has_role and isinstance(marker, dict) and marker.get("strategy") == strategy + + def _last_assistant_index(messages: list[ChatMessage]) -> int: """Return the index of the last assistant message, or -1 if none exists.""" for index in reversed(range(len(messages))): @@ -47,6 +77,38 @@ def _agent_step_spans(messages: list[ChatMessage], start: int) -> list[tuple[int return spans +def _current_step_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """Return message-index groups for complete Agent steps belonging to the current task.""" + step_start = task_index + 1 if task_index is not None else system_end + return [list(range(start, end)) for start, end in _agent_step_spans(messages=messages, start=step_start)] + + +def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]: + """ + Return spans for complete real-user turns in a bounded section of conversation history. + + Each turn begins with a user message not created by compaction and continues up to the next such message. + """ + user_indices = [ + index + for index in range(start, end) + if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta + ] + return [ + (index, user_indices[position + 1] if position + 1 < len(user_indices) else end) + for position, index in enumerate(user_indices) + ] + + +def _historical_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """Return message-index groups for complete historical turns preceding the current task.""" + historical_end = task_index if task_index is not None else system_end + return [ + list(range(start, end)) + for start, end in _historical_turn_spans(messages=messages, start=system_end, end=historical_end) + ] + + def _estimated_context_tokens( messages: list[ChatMessage], context_tokens: int, token_counter: TokenCounter, tools: ToolsType | None = None ) -> int: diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml index 4fed6fb1452..fef10aad17e 100644 --- a/pydoc/hooks_api.yml +++ b/pydoc/hooks_api.yml @@ -1,7 +1,7 @@ loaders: - search_path: [../haystack/hooks] modules: ["protocol", "from_function", "compaction/hooks", "compaction/sliding_window", - "compaction/tool_result_pruning", + "compaction/summarization", "compaction/tool_result_pruning", "compaction/types/protocol", "human_in_the_loop/dataclasses", "human_in_the_loop/hooks", "human_in_the_loop/policies", "human_in_the_loop/strategies", "human_in_the_loop/user_interfaces", "tool_result_offloading/hooks", "tool_result_offloading/policies", "tool_result_offloading/stores", diff --git a/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml b/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml new file mode 100644 index 00000000000..d6125f531b5 --- /dev/null +++ b/releasenotes/notes/add-summarization-compactor-91b6be6855f478df.yaml @@ -0,0 +1,6 @@ +--- +features: + - | + Added the experimental ``SummarizationCompactor``. It uses a Chat Generator to condense the Agent's oldest + historical turns, accumulated summaries, and old current-task steps as needed to reach the context target while + preserving leading system messages, the latest user task, and complete recent tool-calling steps. diff --git a/test/components/generators/chat/test_utils.py b/test/components/generators/chat/test_utils.py new file mode 100644 index 00000000000..6d64d54dd26 --- /dev/null +++ b/test/components/generators/chat/test_utils.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from haystack.components.generators.chat import MockChatGenerator +from haystack.components.generators.chat.utils import _generator_output_token_limit_key, _resolve_output_token_limit + + +class OpenAIChatGenerator(MockChatGenerator): + def __init__(self, generation_kwargs=None): + super().__init__("response") + self.generation_kwargs = generation_kwargs or {} + + +class OpenAIResponsesChatGenerator(MockChatGenerator): + def __init__(self, generation_kwargs=None): + super().__init__("response") + self.generation_kwargs = generation_kwargs or {} + + +class CustomGenerator(MockChatGenerator): + generation_kwargs = {"max_tokens": 7} + + +def test_resolves_chat_completions_limit(): + generator = OpenAIChatGenerator() + assert _generator_output_token_limit_key(generator) == "max_completion_tokens" + assert _resolve_output_token_limit(generator, 100) == (100, {"max_completion_tokens": 100}) + + +def test_resolves_responses_limit(): + generator = OpenAIResponsesChatGenerator() + assert _generator_output_token_limit_key(generator) == "max_output_tokens" + assert _resolve_output_token_limit(generator, 100) == (100, {"max_output_tokens": 100}) + + +def test_configured_generator_limit_wins_without_mutation(): + generator = OpenAIChatGenerator({"temperature": 0, "max_completion_tokens": 23}) + original = dict(generator.generation_kwargs) + assert _resolve_output_token_limit(generator, 100) == (23, None) + assert generator.generation_kwargs == original + + +def test_unknown_generator_receives_no_guessed_runtime_setting(): + assert _generator_output_token_limit_key(CustomGenerator()) is None + assert _resolve_output_token_limit(CustomGenerator(), 100) == (100, None) diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 332332b71b6..d847c50fe55 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -6,8 +6,8 @@ from haystack.dataclasses import ChatMessage, ChatRole, ToolCall from haystack.hooks.compaction import SlidingWindowCompactor -from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _historical_turn_spans, _is_compaction_note -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from haystack.hooks.compaction.sliding_window import _DEFAULT_OMISSION_NOTE, _is_compaction_note +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _historical_turn_spans from test.hooks.compaction.helpers import ( FakeCounter, count_markers, diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py new file mode 100644 index 00000000000..62064ffcd6d --- /dev/null +++ b/test/hooks/compaction/test_summarization.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +import pytest + +from haystack.components.generators.chat import MockChatGenerator +from haystack.dataclasses import ChatMessage, ChatRole +from haystack.hooks.compaction import SummarizationCompactor +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result + +pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") +COUNTER = FakeCounter(chars_per_token=1) + + +def recording_generator(responses: list[str | Exception]) -> tuple[MockChatGenerator, list[dict[str, Any]]]: + """Build a MockChatGenerator whose response function records prompts and can raise queued failures.""" + queued = list(responses) + calls: list[dict[str, Any]] = [] + + def respond(messages: list[ChatMessage]) -> str: + calls.append({"messages": messages}) + response = queued.pop(0) + if isinstance(response, Exception): + raise response + return response + + return MockChatGenerator(response_fn=respond), calls + + +def summary(text: str, source: str) -> ChatMessage: + return ChatMessage.from_user( + f"\n{text}\n", + meta={_COMPACTION_META_KEY: {"strategy": "summarization", "source": source}}, + ) + + +def transcript(call: dict[str, Any]) -> str: + return call["messages"][-1].text + + +class TestSummarizationCompactor: + def test_summarizes_the_minimum_number_of_oldest_historical_turns(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + generator, calls = recording_generator(["short historical summary"]) + retained_without_oldest = [messages[0], *messages[3:]] + target = COUNTER.count(retained_without_oldest) + 100 + compacted = SummarizationCompactor(generator, max_summary_tokens=100).compact( + messages=messages, target_tokens=target, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 1 + assert "oldest question" in transcript(calls[0]) + assert "recent question" not in transcript(calls[0]) + assert compacted[0] == messages[0] + assert compacted[2:] == messages[3:] + assert messages == [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + + def test_reaches_current_steps_in_the_same_call_after_historical_context(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question " * 40), + ChatMessage.from_assistant("old answer " * 40), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 40, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + generator, calls = recording_generator(["history", "old step"]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert "old question" in transcript(calls[0]) + assert "old result" in transcript(calls[1]) + assert compacted[-2:] == messages[-2:] + assert [m.meta[_COMPACTION_META_KEY]["source"] for m in compacted if _COMPACTION_META_KEY in m.meta] == [ + "historical_turns", + "current_task_steps", + ] + + def test_consolidates_historical_summaries_before_current_steps(self): + messages = [ + ChatMessage.from_system("rules"), + summary("first history " * 20, "historical_turns"), + summary("second history " * 20, "historical_turns"), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + generator, calls = recording_generator(["combined history", "old step"]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert "first history" in transcript(calls[0]) + assert "old result" not in transcript(calls[0]) + assert "old result" in transcript(calls[1]) + + def test_consolidates_current_summaries_before_more_raw_steps(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + summary("first step summary " * 20, "current_task_steps"), + summary("second step summary " * 20, "current_task_steps"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] + generator, calls = recording_generator(["combined steps", "old step"]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert "first step summary" in transcript(calls[0]) + assert "old result" in transcript(calls[1]) + assert compacted[-2:] == messages[-2:] + + def test_returns_partial_progress_when_a_later_tier_fails_by_default(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question " * 30), + ChatMessage.from_assistant("old answer " * 30), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("old step " * 30), + ChatMessage.from_assistant("new step"), + ] + generator, calls = recording_generator(["history", RuntimeError("provider unavailable")]) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert len(calls) == 2 + assert any( + message.meta.get(_COMPACTION_META_KEY, {}).get("source") == "historical_turns" for message in compacted + ) + assert messages[-2:] == compacted[-2:] + + def test_raises_when_configured_and_summary_does_not_shrink_context(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old"), + ChatMessage.from_assistant("answer"), + ChatMessage.from_user("current"), + ] + generator, _ = recording_generator(["much longer summary " * 100]) + with pytest.raises(RuntimeError, match="did not reduce"): + SummarizationCompactor(generator, max_summary_tokens=1, raise_on_failure=True).compact(messages, 1, COUNTER) + + def test_returns_none_without_calling_generator_when_context_fits(self): + generator, calls = recording_generator(["unused"]) + messages = [ChatMessage.from_system("rules"), ChatMessage.from_user("task")] + assert ( + SummarizationCompactor(generator).compact(messages=messages, target_tokens=10_000, token_counter=COUNTER) + is None + ) + assert calls == [] + + def test_summary_is_a_user_message_with_compaction_metadata(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old " * 100), + ChatMessage.from_assistant("answer " * 100), + ChatMessage.from_user("task"), + ] + compacted = SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + assert compacted is not None + generated = compacted[1] + assert generated.is_from(ChatRole.USER) + assert generated.meta[_COMPACTION_META_KEY] == { + "strategy": "summarization", + "summarized_messages": 2, + "source": "historical_turns", + } + + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"min_keep_steps": -1}, "`min_keep_steps` must be at least 0"), + ({"max_summary_tokens": 0}, "`max_summary_tokens` must be a positive"), + ], + ) + def test_rejects_invalid_settings(self, kwargs, match): + with pytest.raises(ValueError, match=match): + SummarizationCompactor(MockChatGenerator("summary"), **kwargs) + + def test_serde_round_trip(self): + compactor = SummarizationCompactor( + MockChatGenerator("summary"), + min_keep_steps=2, + max_summary_tokens=321, + summary_instruction="custom", + raise_on_failure=True, + ) + restored = SummarizationCompactor.from_dict(compactor.to_dict()) + assert isinstance(restored.chat_generator, MockChatGenerator) + assert restored.min_keep_steps == 2 + assert restored.max_summary_tokens == 321 + assert restored.summary_instruction == "custom" + assert restored.raise_on_failure is True + + +@pytest.mark.asyncio +async def test_async_compaction_uses_async_generator(): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old " * 100), + ChatMessage.from_assistant("answer " * 100), + ChatMessage.from_user("task"), + ] + generator, calls = recording_generator(["async summary"]) + compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + assert compacted is not None + assert len(calls) == 1 From ef54bcf0bd4c4d87b1d7a808320a0af137279b40 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:45:56 +0200 Subject: [PATCH 33/38] refactor to make it more understandable --- haystack/hooks/compaction/sliding_window.py | 6 +- haystack/hooks/compaction/summarization.py | 510 +++++++------------- haystack/hooks/compaction/utils.py | 5 + test/hooks/compaction/test_summarization.py | 254 +++++----- 4 files changed, 325 insertions(+), 450 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index b3b90c44604..59a5a3cd6d4 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -14,6 +14,7 @@ _is_compaction_message, _latest_user_index, _leading_system_end, + _messages_at, _messages_except, ) from haystack.token_counters import TokenCounter @@ -36,11 +37,6 @@ def _is_compaction_note(message: ChatMessage) -> bool: return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER) -def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: - """Return the messages at the given indices, in the order the indices are given.""" - return [messages[index] for index in indices] - - def _flatten(groups: list[list[int]]) -> list[int]: """Join index groups into a single ordered list of indices.""" return [index for group in groups for index in group] diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index b466a7ee798..2fa2add59a9 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -2,7 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Awaitable, Callable from typing import Any from haystack import logging @@ -18,18 +17,27 @@ _is_compaction_message, _latest_user_index, _leading_system_end, + _messages_at, _messages_except, ) from haystack.token_counters import TokenCounter -from haystack.token_counters.utils import _render_message +from haystack.token_counters.utils import _rendered_conversation from haystack.utils.async_utils import _execute_component_async from haystack.utils.deserialization import deserialize_component_inplace from haystack.utils.experimental import _experimental logger = logging.getLogger(__name__) +# Recorded as the strategy on every summary this compactor produces, so a later run can recognize its own summaries. _STRATEGY = "summarization" +# Recorded as the `source` on a summary, naming the stretch of conversation it stands in for. Compaction gives these up +# in order, so the Agent's current task is the last thing to go. +_HISTORICAL_TURNS = "historical_turns" +_HISTORICAL_SUMMARIES = "historical_summaries" +_CURRENT_TASK_SUMMARIES = "current_task_summaries" +_CURRENT_TASK_STEPS = "current_task_steps" + _DEFAULT_SUMMARY_INSTRUCTION = """You are compacting part of a conversation between a user and an AI agent so the \ agent can keep working with fewer tokens. Write a self-contained summary that preserves: - The user's goal, requirements, constraints, and preferences. @@ -42,53 +50,61 @@ infer or add advice. Use plain prose or short bullets, and do not address the user.""" -def _indices(messages: list[ChatMessage], start: int, end: int, *, summaries: bool) -> list[int]: - """Return summary or non-summary indices in a bounded part of a conversation.""" - return [ - index - for index in range(start, end) - if _is_compaction_message(message=messages[index], strategy=_STRATEGY) is summaries - ] +def _is_summary(message: ChatMessage) -> bool: + """Whether a message is a summary this strategy wrote.""" + return _is_compaction_message(message=message, strategy=_STRATEGY) + + +def _summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: + """Return the positions of this strategy's summaries within a bounded part of a conversation.""" + return [index for index in range(start, end) if _is_summary(message=messages[index])] -def _raw_historical_turn_groups( - messages: list[ChatMessage], system_end: int, task_index: int | None -) -> list[list[int]]: - """Return shared historical-turn groups with this strategy's summaries filtered out.""" - return [ - [index for index in group if not _is_compaction_message(message=messages[index], strategy=_STRATEGY)] +def _summarizable_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: + """ + Return the historical turns that still hold raw conversation, oldest turn first. + + Summaries this strategy already wrote are left out of their turn, so summarizing the turn folds that summary into + the summary this run produces. A turn that is nothing but summaries has nothing left to give up and is dropped. + """ + groups = [ + [index for index in group if not _is_summary(message=messages[index])] for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) ] + return [group for group in groups if group] def _groups_to_summarize( messages: list[ChatMessage], groups: list[list[int]], target_tokens: int, - summary_budget: int, + summary_tokens: int, token_counter: TokenCounter, ) -> list[int]: - """Select the fewest oldest groups that should make room for a summary of the configured size.""" + """ + Return the fewest oldest groups whose removal makes room for a summary of the expected size. + + Groups are taken oldest first and counting stops as soon as what remains, plus the summary that replaces them, + fits the target. When even taking all of them is not enough, all of them are returned. + """ selected: list[int] = [] for group in groups: selected.extend(group) - if ( - token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + summary_budget - <= target_tokens - ): + remaining = token_counter.count(messages=_messages_except(messages=messages, indices=selected)) + if remaining + summary_tokens <= target_tokens: break return selected def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMessage: - """Build a marked summary message.""" + """Build the marked user message that stands in for the messages the summary replaced.""" body = f"\n{text.strip()}\n" meta = {_COMPACTION_META_KEY: {"strategy": _STRATEGY, "summarized_messages": summarized_messages, "source": source}} return ChatMessage.from_user(text=body, meta=meta) def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: - """Replace possibly non-contiguous selected messages with one summary at their first position.""" + """Replace the selected messages, which need not be contiguous, with one summary at the oldest one's position.""" selected = set(indices) insertion_index = min(indices) compacted: list[ChatMessage] = [] @@ -152,6 +168,7 @@ def __init__( :param summary_instruction: An instruction replacing the built-in summary prompt. :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is logged and any successful partial compaction is returned. + :raises ValueError: If `min_keep_steps` is negative or `max_summary_tokens` is not positive. """ if min_keep_steps < 0: raise ValueError(f"`min_keep_steps` must be at least 0, got {min_keep_steps}.") @@ -174,25 +191,28 @@ def compact( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ - if token_counter.count(messages=messages) <= target_tokens: - return None - budget, generation_kwargs = _resolve_output_token_limit( - chat_generator=self.chat_generator, default_limit=self.max_summary_tokens - ) - - def generate(prompt: list[ChatMessage]) -> dict[str, Any]: - kwargs: dict[str, Any] = {"messages": prompt} - if generation_kwargs is not None: - kwargs["generation_kwargs"] = generation_kwargs - return self.chat_generator.run(**kwargs) - - return self._compact( - original=messages, - target_tokens=target_tokens, - token_counter=token_counter, - budget=budget, - generate=generate, - ) + summary_tokens, run_kwargs = self._summary_limit() + working = list(messages) + while True: + plan = self._next_summary( + messages=working, + target_tokens=target_tokens, + token_counter=token_counter, + summary_tokens=summary_tokens, + ) + if plan is None: + break + indices, source = plan + prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + try: + result = self.chat_generator.run(messages=prompt, **run_kwargs) + working = self._apply_summary( + messages=working, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + self._report_failure(error=error) + break + return self._reduced(original=messages, working=working, token_counter=token_counter) async def compact_async( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter @@ -205,47 +225,129 @@ async def compact_async( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ + summary_tokens, run_kwargs = self._summary_limit() + working = list(messages) + while True: + plan = self._next_summary( + messages=working, + target_tokens=target_tokens, + token_counter=token_counter, + summary_tokens=summary_tokens, + ) + if plan is None: + break + indices, source = plan + prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + try: + result = await _execute_component_async( + component_instance=self.chat_generator, messages=prompt, **run_kwargs + ) + working = self._apply_summary( + messages=working, indices=indices, source=source, result=result, token_counter=token_counter + ) + except Exception as error: + self._report_failure(error=error) + break + return self._reduced(original=messages, working=working, token_counter=token_counter) + + def _next_summary( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int + ) -> tuple[list[int], str] | None: + """ + Choose the next stretch of conversation to replace with a summary. + + Four tiers are tried in order, so the oldest and least useful context goes first and the Agent's current task + is given up last: + + 1. `_HISTORICAL_TURNS`: the fewest oldest raw turns that make room for a summary. + 2. `_HISTORICAL_SUMMARIES`: nothing raw is left in history, so fold its summaries into one. + 3. `_CURRENT_TASK_SUMMARIES`: fold the summaries earlier steps left behind before giving up more steps. + 4. `_CURRENT_TASK_STEPS`: the fewest oldest steps of the current task, keeping `min_keep_steps` of the newest. + + :param messages: The conversation as it stands, ordered oldest to newest. + :param target_tokens: The token budget the conversation should come in under. + :param token_counter: The counter used to measure candidate selections. + :param summary_tokens: The size a summary is expected to take, reserved when choosing how much to replace. + :returns: The message indices to summarize and the `source` to record on the resulting summary, or None when + the conversation already fits or nothing is left that may be given up. + """ if token_counter.count(messages=messages) <= target_tokens: return None - budget, generation_kwargs = _resolve_output_token_limit( - chat_generator=self.chat_generator, default_limit=self.max_summary_tokens - ) - async def generate(prompt: list[ChatMessage]) -> dict[str, Any]: - kwargs: dict[str, Any] = {"messages": prompt} - if generation_kwargs is not None: - kwargs["generation_kwargs"] = generation_kwargs - return await _execute_component_async(component_instance=self.chat_generator, **kwargs) + # The landmarks everything is measured against: the Agent's instructions, and the user message anchoring the + # current task. History runs from the instructions up to that anchor, the current task from the anchor on. + system_end = _leading_system_end(messages=messages) + task_index = _latest_user_index(messages=messages) + history_end = task_index if task_index is not None else system_end + task_start = task_index + 1 if task_index is not None else system_end + + turns = _summarizable_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + if turns: + oldest_turns = _groups_to_summarize( + messages=messages, + groups=turns, + target_tokens=target_tokens, + summary_tokens=summary_tokens, + token_counter=token_counter, + ) + return oldest_turns, _HISTORICAL_TURNS + + history_summaries = _summary_indices(messages=messages, start=system_end, end=history_end) + if len(history_summaries) > 1: + return history_summaries, _HISTORICAL_SUMMARIES + + # Only steps older than the `min_keep_steps` most recent ones may be given up. + steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) + eligible = steps[: max(len(steps) - self.min_keep_steps, 0)] + if not eligible: + return None + + task_summaries = _summary_indices(messages=messages, start=task_start, end=len(messages)) + if len(task_summaries) > 1: + return task_summaries, _CURRENT_TASK_SUMMARIES - return await self._compact_async( - original=messages, + oldest_steps = _groups_to_summarize( + messages=messages, + groups=eligible, target_tokens=target_tokens, + summary_tokens=summary_tokens, token_counter=token_counter, - budget=budget, - generate=generate, ) + return oldest_steps, _CURRENT_TASK_STEPS - def _prompt(self, messages: list[ChatMessage], budget: int) -> list[ChatMessage]: - """Build the bounded summarization instruction and rendered source transcript.""" - transcript = "\n".join(_render_message(message=message) for message in messages) + def _summary_limit(self) -> tuple[int, dict[str, Any]]: + """Return the token budget for one summary and the run kwargs, if any, that ask the generator to honor it.""" + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + return summary_tokens, {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + + def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: + """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" + transcript = _rendered_conversation(_messages_at(messages=messages, indices=indices)) instruction = ( - f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately {budget} tokens. " - "Prioritize completeness within that limit so the response is not cut off." + f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately " + f"{summary_tokens} tokens. Prioritize completeness within that limit so the response is not cut off." ) return [ ChatMessage.from_system(text=instruction), ChatMessage.from_user(text=f"\n{transcript}\n"), ] - def _apply_result( - self, + @staticmethod + def _apply_summary( messages: list[ChatMessage], indices: list[int], source: str, result: dict[str, Any], token_counter: TokenCounter, ) -> list[ChatMessage]: - """Validate a generator reply and replace its source messages when the result is smaller.""" + """ + Swap the selected messages for the generated summary. + + :raises RuntimeError: If the generator returned no usable text, or if the swap did not make the conversation + smaller, in which case keeping the raw messages is the better outcome. + """ replies = result.get("replies") or [] text = replies[-1].text if replies else None if not text or not text.strip(): @@ -261,283 +363,21 @@ def _apply_result( ) return compacted - def _attempt( - self, - messages: list[ChatMessage], - indices: list[int], - source: str, - budget: int, - token_counter: TokenCounter, - generate: Callable[[list[ChatMessage]], dict[str, Any]], - ) -> list[ChatMessage] | None: - """Attempt one synchronous summary, applying the configured failure policy.""" - try: - result = generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) - return self._apply_result( - messages=messages, indices=indices, source=source, result=result, token_counter=token_counter - ) - except Exception as error: - if self.raise_on_failure: - raise - logger.warning( - "Summarizing the conversation for context compaction failed; keeping the last successful result. " - "Error: {error}", - error=error, - ) - return None - - async def _attempt_async( - self, - messages: list[ChatMessage], - indices: list[int], - source: str, - budget: int, - token_counter: TokenCounter, - generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], - ) -> list[ChatMessage] | None: - """Attempt one asynchronous summary, applying the configured failure policy.""" - try: - result = await generate(self._prompt(messages=[messages[index] for index in indices], budget=budget)) - return self._apply_result( - messages=messages, indices=indices, source=source, result=result, token_counter=token_counter - ) - except Exception as error: - if self.raise_on_failure: - raise - logger.warning( - "Summarizing the conversation for context compaction failed; keeping the last successful result. " - "Error: {error}", - error=error, - ) - return None - - def _compact( - self, - original: list[ChatMessage], - target_tokens: int, - token_counter: TokenCounter, - budget: int, - generate: Callable[[list[ChatMessage]], dict[str, Any]], - ) -> list[ChatMessage] | None: - """Run synchronous historical, consolidation, and current-step compaction tiers in order.""" - working = list(original) - - # First replace the fewest oldest raw historical turns expected to reach the target. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - groups = [ - group - for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) - if group - ] - if not groups: - break - selected = _groups_to_summarize( - messages=working, - groups=groups, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = self._attempt( - messages=working, - indices=selected, - source="historical_turns", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # If all historical turns are summaries and the target is still unmet, fold them into one summary. - if token_counter.count(messages=working) > target_tokens: - summaries = self._summary_indices(messages=working, source="historical_summaries") - if len(summaries) > 1: - compacted = self._attempt( - messages=working, - indices=summaries, - source="historical_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - if not eligible: - break - - # Fold accumulated current-task summaries before consuming more raw agent steps. - summaries = self._summary_indices(messages=working, source="current_task_summaries") - if len(summaries) > 1: - compacted = self._attempt( - messages=working, - indices=summaries, - source="current_task_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - selected = _groups_to_summarize( - messages=working, - groups=eligible, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = self._attempt( - messages=working, - indices=selected, - source="current_task_steps", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - break - working = compacted - return self._result(original=original, working=working, token_counter=token_counter) - - async def _compact_async( - self, - original: list[ChatMessage], - target_tokens: int, - token_counter: TokenCounter, - budget: int, - generate: Callable[[list[ChatMessage]], Awaitable[dict[str, Any]]], - ) -> list[ChatMessage] | None: - """Run asynchronous historical, consolidation, and current-step compaction tiers in order.""" - working = list(original) - - # First replace the fewest oldest raw historical turns expected to reach the target. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - groups = [ - group - for group in _raw_historical_turn_groups(messages=working, system_end=system_end, task_index=task_index) - if group - ] - if not groups: - break - selected = _groups_to_summarize( - messages=working, - groups=groups, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = await self._attempt_async( - messages=working, - indices=selected, - source="historical_turns", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # If all historical turns are summaries and the target is still unmet, fold them into one summary. - if token_counter.count(messages=working) > target_tokens: - summaries = self._summary_indices(messages=working, source="historical_summaries") - if len(summaries) > 1: - compacted = await self._attempt_async( - messages=working, - indices=summaries, - source="historical_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Finally summarize the oldest eligible agent steps while preserving the configured recent steps. - while token_counter.count(messages=working) > target_tokens: - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - if not eligible: - break - - # Fold accumulated current-task summaries before consuming more raw agent steps. - summaries = self._summary_indices(messages=working, source="current_task_summaries") - if len(summaries) > 1: - compacted = await self._attempt_async( - messages=working, - indices=summaries, - source="current_task_summaries", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - return self._result(original=original, working=working, token_counter=token_counter) - working = compacted - - # Recompute positions after consolidation, then select the minimum useful prefix of raw steps. - system_end = _leading_system_end(messages=working) - task_index = _latest_user_index(messages=working) - step_groups = _current_step_groups(messages=working, system_end=system_end, task_index=task_index) - eligible = step_groups[: max(len(step_groups) - self.min_keep_steps, 0)] - selected = _groups_to_summarize( - messages=working, - groups=eligible, - target_tokens=target_tokens, - summary_budget=budget, - token_counter=token_counter, - ) - compacted = await self._attempt_async( - messages=working, - indices=selected, - source="current_task_steps", - budget=budget, - token_counter=token_counter, - generate=generate, - ) - if compacted is None: - break - working = compacted - return self._result(original=original, working=working, token_counter=token_counter) - - def _summary_indices(self, messages: list[ChatMessage], source: str) -> list[int]: - """Return historical or current-task summary indices based on their conversation position.""" - system_end = _leading_system_end(messages=messages) - task_index = _latest_user_index(messages=messages) - end = task_index if task_index is not None else len(messages) - if source == "historical_summaries": - return _indices(messages=messages, start=system_end, end=end, summaries=True) - start = task_index + 1 if task_index is not None else system_end - return _indices(messages=messages, start=start, end=len(messages), summaries=True) + def _report_failure(self, error: Exception) -> None: + """Re-raise a failed summarization or log it, so whatever compacted successfully so far is still returned.""" + if self.raise_on_failure: + raise error + logger.warning( + "Summarizing the conversation for context compaction failed; keeping the last successful result. " + "Error: {error}", + error=error, + ) @staticmethod - def _result( + def _reduced( original: list[ChatMessage], working: list[ChatMessage], token_counter: TokenCounter ) -> list[ChatMessage] | None: - """Return partial or complete progress only when it reduced the original conversation.""" + """Return partial or complete progress only when it made the original conversation smaller.""" if token_counter.count(messages=working) < token_counter.count(messages=original): return working return None diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 99b0fe79d8e..590ddb84b2d 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -27,6 +27,11 @@ def _latest_user_index(messages: list[ChatMessage]) -> int | None: return None +def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: + """Return the messages at the given indices, in the order the indices are given.""" + return [messages[index] for index in indices] + + def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]: """Return the messages the given indices leave out, in conversation order.""" left_out = set(indices) diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py index 62064ffcd6d..c9d916fa56d 100644 --- a/test/hooks/compaction/test_summarization.py +++ b/test/hooks/compaction/test_summarization.py @@ -2,8 +2,6 @@ # # SPDX-License-Identifier: Apache-2.0 -from typing import Any - import pytest from haystack.components.generators.chat import MockChatGenerator @@ -13,140 +11,175 @@ from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") + +# A target of one token forces every tier to run, isolating the structural rules from sizing. +SMALLEST = 1 +# One character per token, so the padded messages below are obviously the expensive ones. COUNTER = FakeCounter(chars_per_token=1) -def recording_generator(responses: list[str | Exception]) -> tuple[MockChatGenerator, list[dict[str, Any]]]: - """Build a MockChatGenerator whose response function records prompts and can raise queued failures.""" +def summarizer(*responses: str | Exception) -> tuple[MockChatGenerator, list[str]]: + """ + A Chat Generator returning the given summaries in order, recording the prompt it received for each. + + An `Exception` among the responses is raised instead of answering, so a test can fail one summarization step. + """ queued = list(responses) - calls: list[dict[str, Any]] = [] + prompts: list[str] = [] def respond(messages: list[ChatMessage]) -> str: - calls.append({"messages": messages}) + prompts.append("\n".join(message.text or "" for message in messages)) response = queued.pop(0) if isinstance(response, Exception): raise response return response - return MockChatGenerator(response_fn=respond), calls + return MockChatGenerator(response_fn=respond), prompts def summary(text: str, source: str) -> ChatMessage: + """A summary an earlier compaction left behind, marked the way this compactor marks its own.""" return ChatMessage.from_user( f"\n{text}\n", meta={_COMPACTION_META_KEY: {"strategy": "summarization", "source": source}}, ) -def transcript(call: dict[str, Any]) -> str: - return call["messages"][-1].text +def sources(messages: list[ChatMessage]) -> list[str]: + """Which stretch of conversation each summary in `messages` stands in for, oldest first.""" + return [ + message.meta[_COMPACTION_META_KEY]["source"] for message in messages if _COMPACTION_META_KEY in message.meta + ] + + +def two_turns_and_a_task() -> list[ChatMessage]: + """A padded oldest turn, a short recent turn, and the current task with one step behind it.""" + return [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("oldest question " * 30), + ChatMessage.from_assistant("oldest answer " * 30), + ChatMessage.from_user("recent question"), + ChatMessage.from_assistant("recent answer"), + ChatMessage.from_user("current task"), + ChatMessage.from_assistant("current step"), + ] + + +def a_task_with_two_steps() -> list[ChatMessage]: + """The current task with a padded oldest step and a cheap newest one, and no history in front of it.""" + return [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + tool_call("old"), + tool_result("old result " * 30, call_id="old"), + tool_call("new"), + tool_result("new result", call_id="new"), + ] class TestSummarizationCompactor: - def test_summarizes_the_minimum_number_of_oldest_historical_turns(self): - messages = [ - ChatMessage.from_system("rules"), - ChatMessage.from_user("oldest question " * 30), - ChatMessage.from_assistant("oldest answer " * 30), - ChatMessage.from_user("recent question"), - ChatMessage.from_assistant("recent answer"), - ChatMessage.from_user("current task"), - ChatMessage.from_assistant("current step"), - ] - generator, calls = recording_generator(["short historical summary"]) - retained_without_oldest = [messages[0], *messages[3:]] - target = COUNTER.count(retained_without_oldest) + 100 + def test_summarizes_oldest_historical_turn(self): + messages = two_turns_and_a_task() + generator, prompts = summarizer("short historical summary") + # Room for everything but the padded oldest turn, plus the summary standing in for it. + target_tokens = COUNTER.count([messages[0], *messages[3:]]) + 100 + compacted = SummarizationCompactor(generator, max_summary_tokens=100).compact( - messages=messages, target_tokens=target, token_counter=COUNTER + messages=messages, target_tokens=target_tokens, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 1 - assert "oldest question" in transcript(calls[0]) - assert "recent question" not in transcript(calls[0]) - assert compacted[0] == messages[0] - assert compacted[2:] == messages[3:] - assert messages == [ - ChatMessage.from_system("rules"), - ChatMessage.from_user("oldest question " * 30), - ChatMessage.from_assistant("oldest answer " * 30), - ChatMessage.from_user("recent question"), - ChatMessage.from_assistant("recent answer"), - ChatMessage.from_user("current task"), - ChatMessage.from_assistant("current step"), - ] + # Only the oldest turn was summarized, so the recent turn never reached the generator. + assert len(prompts) == 1 + assert "oldest question" in prompts[0] + assert "recent question" not in prompts[0] + assert compacted == [messages[0], compacted[1], *messages[3:]] + + def test_leaves_the_input_conversation_untouched(self): + messages = two_turns_and_a_task() + SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + assert messages == two_turns_and_a_task() - def test_reaches_current_steps_in_the_same_call_after_historical_context(self): + def test_summarizes_historical_turns_then_current_steps(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old question " * 40), ChatMessage.from_assistant("old answer " * 40), - ChatMessage.from_user("current task"), - tool_call("old"), - tool_result("old result " * 40, call_id="old"), - tool_call("new"), - tool_result("new result", call_id="new"), + *a_task_with_two_steps()[1:], ] - generator, calls = recording_generator(["history", "old step"]) + generator, prompts = summarizer("history", "old step") + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert "old question" in transcript(calls[0]) - assert "old result" in transcript(calls[1]) + assert "old question" in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["historical_turns", "current_task_steps"] + # The newest step is never given up. assert compacted[-2:] == messages[-2:] - assert [m.meta[_COMPACTION_META_KEY]["source"] for m in compacted if _COMPACTION_META_KEY in m.meta] == [ - "historical_turns", - "current_task_steps", - ] - def test_consolidates_historical_summaries_before_current_steps(self): + def test_folds_historical_summaries_before_current_steps(self): messages = [ ChatMessage.from_system("rules"), summary("first history " * 20, "historical_turns"), summary("second history " * 20, "historical_turns"), - ChatMessage.from_user("current task"), - tool_call("old"), - tool_result("old result " * 30, call_id="old"), - tool_call("new"), - tool_result("new result", call_id="new"), + *a_task_with_two_steps()[1:], ] - generator, calls = recording_generator(["combined history", "old step"]) + generator, prompts = summarizer("combined history", "old step") + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert "first history" in transcript(calls[0]) - assert "old result" not in transcript(calls[0]) - assert "old result" in transcript(calls[1]) + # The two historical summaries are folded into one before any current-task step is touched. + assert "first history" in prompts[0] and "old result" not in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["historical_summaries", "current_task_steps"] - def test_consolidates_current_summaries_before_more_raw_steps(self): + def test_folds_current_task_summaries_before_more_steps(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("current task"), summary("first step summary " * 20, "current_task_steps"), summary("second step summary " * 20, "current_task_steps"), - tool_call("old"), - tool_result("old result " * 30, call_id="old"), - tool_call("new"), - tool_result("new result", call_id="new"), + *a_task_with_two_steps()[2:], ] - generator, calls = recording_generator(["combined steps", "old step"]) + generator, prompts = summarizer("combined steps", "old step") + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert "first step summary" in transcript(calls[0]) - assert "old result" in transcript(calls[1]) + assert "first step summary" in prompts[0] + assert "old result" in prompts[1] + assert sources(compacted) == ["current_task_summaries", "current_task_steps"] assert compacted[-2:] == messages[-2:] - def test_returns_partial_progress_when_a_later_tier_fails_by_default(self): + @pytest.mark.parametrize(("min_keep_steps", "expected"), [(0, 0), (1, 1), (2, 2), (20, 2)]) + def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, expected): + messages = a_task_with_two_steps() + generator, _ = summarizer("step summary") + compacted = SummarizationCompactor(generator, min_keep_steps=min_keep_steps, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + result = compacted or messages + assert sum(message.is_from(role=ChatRole.ASSISTANT) for message in result) == expected + + def test_custom_summary_instruction_replaces_the_default(self): + generator, prompts = summarizer("summary") + SummarizationCompactor(generator, summary_instruction="Only list file paths.", max_summary_tokens=1).compact( + messages=two_turns_and_a_task(), target_tokens=SMALLEST, token_counter=COUNTER + ) + assert "Only list file paths." in prompts[0] + assert "You are compacting part of a conversation" not in prompts[0] + + def test_keeps_partial_progress_when_a_summary_fails(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old question " * 30), @@ -155,39 +188,41 @@ def test_returns_partial_progress_when_a_later_tier_fails_by_default(self): ChatMessage.from_assistant("old step " * 30), ChatMessage.from_assistant("new step"), ] - generator, calls = recording_generator(["history", RuntimeError("provider unavailable")]) + generator, prompts = summarizer("history", RuntimeError("provider unavailable")) + compacted = SummarizationCompactor(generator, max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - assert len(calls) == 2 - assert any( - message.meta.get(_COMPACTION_META_KEY, {}).get("source") == "historical_turns" for message in compacted - ) - assert messages[-2:] == compacted[-2:] + assert len(prompts) == 2 + # The history was summarized before the step summary failed, and that progress is kept. + assert sources(compacted) == ["historical_turns"] + assert compacted[-2:] == messages[-2:] - def test_raises_when_configured_and_summary_does_not_shrink_context(self): + def test_raises_when_a_summary_does_not_shrink_the_conversation(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old"), ChatMessage.from_assistant("answer"), ChatMessage.from_user("current"), ] - generator, _ = recording_generator(["much longer summary " * 100]) + compactor = SummarizationCompactor( + MockChatGenerator("much longer summary " * 100), max_summary_tokens=1, raise_on_failure=True + ) with pytest.raises(RuntimeError, match="did not reduce"): - SummarizationCompactor(generator, max_summary_tokens=1, raise_on_failure=True).compact(messages, 1, COUNTER) + compactor.compact(messages=messages, target_tokens=SMALLEST, token_counter=COUNTER) - def test_returns_none_without_calling_generator_when_context_fits(self): - generator, calls = recording_generator(["unused"]) + def test_returns_none_when_the_conversation_fits(self): + generator, prompts = summarizer("unused") messages = [ChatMessage.from_system("rules"), ChatMessage.from_user("task")] - assert ( - SummarizationCompactor(generator).compact(messages=messages, target_tokens=10_000, token_counter=COUNTER) - is None + compacted = SummarizationCompactor(generator).compact( + messages=messages, target_tokens=10_000, token_counter=COUNTER ) - assert calls == [] + assert compacted is None + assert prompts == [] - def test_summary_is_a_user_message_with_compaction_metadata(self): + def test_summary_is_a_marked_user_message(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old " * 100), @@ -195,12 +230,11 @@ def test_summary_is_a_user_message_with_compaction_metadata(self): ChatMessage.from_user("task"), ] compacted = SummarizationCompactor(MockChatGenerator("summary"), max_summary_tokens=1).compact( - messages=messages, target_tokens=1, token_counter=COUNTER + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER ) assert compacted is not None - generated = compacted[1] - assert generated.is_from(ChatRole.USER) - assert generated.meta[_COMPACTION_META_KEY] == { + assert compacted[1].is_from(role=ChatRole.USER) + assert compacted[1].meta[_COMPACTION_META_KEY] == { "strategy": "summarization", "summarized_messages": 2, "source": "historical_turns", @@ -233,17 +267,17 @@ def test_serde_round_trip(self): assert restored.raise_on_failure is True -@pytest.mark.asyncio -async def test_async_compaction_uses_async_generator(): - messages = [ - ChatMessage.from_system("rules"), - ChatMessage.from_user("old " * 100), - ChatMessage.from_assistant("answer " * 100), - ChatMessage.from_user("task"), - ] - generator, calls = recording_generator(["async summary"]) - compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( - messages=messages, target_tokens=1, token_counter=COUNTER - ) - assert compacted is not None - assert len(calls) == 1 +class TestSummarizationCompactorAsync: + @pytest.mark.asyncio + async def test_compact_async_matches_compact(self): + messages = two_turns_and_a_task() + generator, prompts = summarizer("async summary") + + compacted = await SummarizationCompactor(generator, max_summary_tokens=1).compact_async( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + + assert len(prompts) == 1 + assert compacted == SummarizationCompactor(MockChatGenerator("async summary"), max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) From 6f3d88c221a16e50915a78f6341c809d26306d0c Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:57:22 +0200 Subject: [PATCH 34/38] more refactoring --- haystack/hooks/compaction/summarization.py | 106 +++++++++++++-------- 1 file changed, 68 insertions(+), 38 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index 2fa2add59a9..7e2ed9053ee 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -55,18 +55,22 @@ def _is_summary(message: ChatMessage) -> bool: return _is_compaction_message(message=message, strategy=_STRATEGY) -def _summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: - """Return the positions of this strategy's summaries within a bounded part of a conversation.""" +def _previous_summary_indices(messages: list[ChatMessage], start: int, end: int) -> list[int]: + """Return the positions of the summaries an earlier compaction left in a bounded part of a conversation.""" return [index for index in range(start, end) if _is_summary(message=messages[index])] -def _summarizable_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]: +def _raw_historical_turn_groups( + messages: list[ChatMessage], system_end: int, task_index: int | None +) -> list[list[int]]: """ - Return the historical turns that still hold raw conversation, oldest turn first. + Return the historical turns that still hold raw, never-summarized conversation, oldest turn first. - Summaries this strategy already wrote are left out of their turn, so summarizing the turn folds that summary into - the summary this run produces. A turn that is nothing but summaries has nothing left to give up and is dropped. + Summaries an earlier compaction wrote are excluded, so summarizing a turn leaves them in place for + `_HISTORICAL_SUMMARIES` to fold later. The list is empty when there are no historical turns, or when every one of + them is already nothing but summaries. """ + # Strip the previous summaries out of each turn, then drop the turns that strip away to nothing. groups = [ [index for index in group if not _is_summary(message=messages[index])] for group in _historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) @@ -106,9 +110,11 @@ def _summary_message(text: str, summarized_messages: int, source: str) -> ChatMe def _replace_indices(messages: list[ChatMessage], indices: list[int], summary: ChatMessage) -> list[ChatMessage]: """Replace the selected messages, which need not be contiguous, with one summary at the oldest one's position.""" selected = set(indices) + # The summary stands in for everything it replaced, so it takes the position of the oldest message it covers. insertion_index = min(indices) compacted: list[ChatMessage] = [] for index, message in enumerate(messages): + # Emit the summary before the message it displaces, so the surrounding conversation keeps its order. if index == insertion_index: compacted.append(summary) if index not in selected: @@ -191,11 +197,14 @@ def compact( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ + # How large each summary may be, and the run kwargs, if any, that hold the generator to it. summary_tokens, run_kwargs = self._summary_limit() - working = list(messages) + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. + compacted = messages while True: + # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( - messages=working, + messages=compacted, target_tokens=target_tokens, token_counter=token_counter, summary_tokens=summary_tokens, @@ -203,16 +212,21 @@ def compact( if plan is None: break indices, source = plan - prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + prompt = self._prompt(messages=compacted, indices=indices, summary_tokens=summary_tokens) try: + # Summarize that stretch and swap it in, so the next round plans against the smaller conversation. + # A generator error or a summary that does not shrink raises out of here. result = self.chat_generator.run(messages=prompt, **run_kwargs) - working = self._apply_summary( - messages=working, indices=indices, source=source, result=result, token_counter=token_counter + compacted = self._apply_summary( + messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) except Exception as error: + # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - return self._reduced(original=messages, working=working, token_counter=token_counter) + # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than + # the untouched input means real progress, whether or not the target was met. + return None if compacted is messages else compacted async def compact_async( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter @@ -225,11 +239,14 @@ async def compact_async( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ + # How large each summary may be, and the run kwargs, if any, that hold the generator to it. summary_tokens, run_kwargs = self._summary_limit() - working = list(messages) + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. + compacted = messages while True: + # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( - messages=working, + messages=compacted, target_tokens=target_tokens, token_counter=token_counter, summary_tokens=summary_tokens, @@ -237,18 +254,23 @@ async def compact_async( if plan is None: break indices, source = plan - prompt = self._prompt(messages=working, indices=indices, summary_tokens=summary_tokens) + prompt = self._prompt(messages=compacted, indices=indices, summary_tokens=summary_tokens) try: + # Summarize that stretch and swap it in, so the next round plans against the smaller conversation. + # Only the generator call is awaited; planning and swapping are pure. result = await _execute_component_async( component_instance=self.chat_generator, messages=prompt, **run_kwargs ) - working = self._apply_summary( - messages=working, indices=indices, source=source, result=result, token_counter=token_counter + compacted = self._apply_summary( + messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) except Exception as error: + # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - return self._reduced(original=messages, working=working, token_counter=token_counter) + # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than + # the untouched input means real progress, whether or not the target was met. + return None if compacted is messages else compacted def _next_summary( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int @@ -271,6 +293,7 @@ def _next_summary( :returns: The message indices to summarize and the `source` to record on the resulting summary, or None when the conversation already fits or nothing is left that may be given up. """ + # Nothing to give up once the conversation fits. if token_counter.count(messages=messages) <= target_tokens: return None @@ -281,34 +304,39 @@ def _next_summary( history_end = task_index if task_index is not None else system_end task_start = task_index + 1 if task_index is not None else system_end - turns = _summarizable_turn_groups(messages=messages, system_end=system_end, task_index=task_index) - if turns: + # Tier 1. Raw history is the cheapest context to lose, so take the oldest turns that still hold any. + historical_turns = _raw_historical_turn_groups(messages=messages, system_end=system_end, task_index=task_index) + if historical_turns: oldest_turns = _groups_to_summarize( messages=messages, - groups=turns, + groups=historical_turns, target_tokens=target_tokens, summary_tokens=summary_tokens, token_counter=token_counter, ) return oldest_turns, _HISTORICAL_TURNS - history_summaries = _summary_indices(messages=messages, start=system_end, end=history_end) + # Tier 2. History is nothing but summaries now, so the only room left there is in folding them into one. They + # are left to accumulate until this point so that they are not rewritten on every compaction. + history_summaries = _previous_summary_indices(messages=messages, start=system_end, end=history_end) if len(history_summaries) > 1: return history_summaries, _HISTORICAL_SUMMARIES - # Only steps older than the `min_keep_steps` most recent ones may be given up. - steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) - eligible = steps[: max(len(steps) - self.min_keep_steps, 0)] - if not eligible: + # History is exhausted, so the current task has to pay. Its `min_keep_steps` newest steps are off limits. + agent_steps = _current_step_groups(messages=messages, system_end=system_end, task_index=task_index) + eligible_steps = agent_steps[: max(len(agent_steps) - self.min_keep_steps, 0)] + if not eligible_steps: return None - task_summaries = _summary_indices(messages=messages, start=task_start, end=len(messages)) + # Tier 3. Fold the summaries earlier steps left behind before spending another raw step on the same space. + task_summaries = _previous_summary_indices(messages=messages, start=task_start, end=len(messages)) if len(task_summaries) > 1: return task_summaries, _CURRENT_TASK_SUMMARIES + # Tier 4. Last resort: give up the oldest steps of the task the Agent is working on right now. oldest_steps = _groups_to_summarize( messages=messages, - groups=eligible, + groups=eligible_steps, target_tokens=target_tokens, summary_tokens=summary_tokens, token_counter=token_counter, @@ -316,7 +344,18 @@ def _next_summary( return oldest_steps, _CURRENT_TASK_STEPS def _summary_limit(self) -> tuple[int, dict[str, Any]]: - """Return the token budget for one summary and the run kwargs, if any, that ask the generator to honor it.""" + """ + Work out how large one summary may be and how to hold the Chat Generator to it. + + :returns: A tuple containing: + + 1. The token budget for a single summary. This is `max_summary_tokens`, unless the generator already + configures a recognized output limit of its own, in which case the generator's setting wins. + 2. The kwargs to pass to the generator's `run`. This carries a `generation_kwargs` entry for a built-in + generator that has no limit configured, and is empty for every other generator, since the + `ChatGenerator` protocol does not standardize the setting. When it is empty, the budget reaches the + model only as prompt guidance and `_apply_summary` measures the result instead. + """ summary_tokens, generation_kwargs = _resolve_output_token_limit( chat_generator=self.chat_generator, default_limit=self.max_summary_tokens ) @@ -373,15 +412,6 @@ def _report_failure(self, error: Exception) -> None: error=error, ) - @staticmethod - def _reduced( - original: list[ChatMessage], working: list[ChatMessage], token_counter: TokenCounter - ) -> list[ChatMessage] | None: - """Return partial or complete progress only when it made the original conversation smaller.""" - if token_counter.count(messages=working) < token_counter.count(messages=original): - return working - return None - def warm_up(self) -> None: """Warm up the Chat Generator that writes summaries.""" if hasattr(self.chat_generator, "warm_up"): From 47b2af232697c0a2084e11152ef46fb39c1927cc Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Mon, 10 Aug 2026 14:59:04 +0200 Subject: [PATCH 35/38] adding more clarity --- haystack/hooks/compaction/summarization.py | 52 ++++++++++------------ 1 file changed, 24 insertions(+), 28 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index 7e2ed9053ee..da23d3b2d7f 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -197,10 +197,16 @@ def compact( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ - # How large each summary may be, and the run kwargs, if any, that hold the generator to it. - summary_tokens, run_kwargs = self._summary_limit() + # How large each summary may be. A recognized built-in generator is held to it by its own runtime setting; + # any other generator gets it as prompt guidance only, since the protocol guarantees nothing beyond `run`. + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + run_kwargs = {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. compacted = messages + summarized = False while True: # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( @@ -220,13 +226,14 @@ def compact( compacted = self._apply_summary( messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) + summarized = True except Exception as error: # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than - # the untouched input means real progress, whether or not the target was met. - return None if compacted is messages else compacted + # Every applied summary was measured as shrinking the conversation, so any summary at all is real progress, + # whether or not the target was met. Without one there is nothing to hand back. + return compacted if summarized else None async def compact_async( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter @@ -239,10 +246,16 @@ async def compact_async( :param token_counter: The counter used both to plan compaction and verify generated summaries. :returns: A smaller replacement conversation, or None when nothing was reduced. """ - # How large each summary may be, and the run kwargs, if any, that hold the generator to it. - summary_tokens, run_kwargs = self._summary_limit() + # How large each summary may be. A recognized built-in generator is held to it by its own runtime setting; + # any other generator gets it as prompt guidance only, since the protocol guarantees nothing beyond `run`. + summary_tokens, generation_kwargs = _resolve_output_token_limit( + chat_generator=self.chat_generator, default_limit=self.max_summary_tokens + ) + run_kwargs = {"generation_kwargs": generation_kwargs} if generation_kwargs else {} + # Rebound only when a summary is applied, and never mutated, so `messages` is left as the caller passed it. compacted = messages + summarized = False while True: # Ask which stretch of the conversation to give up next. None means the target is met or nothing is left. plan = self._next_summary( @@ -264,13 +277,14 @@ async def compact_async( compacted = self._apply_summary( messages=compacted, indices=indices, source=source, result=result, token_counter=token_counter ) + summarized = True except Exception as error: # Stop at the last summary that worked, unless `raise_on_failure` says to propagate. self._report_failure(error=error) break - # Every applied summary was measured as shrinking the conversation, so reaching here with anything other than - # the untouched input means real progress, whether or not the target was met. - return None if compacted is messages else compacted + # Every applied summary was measured as shrinking the conversation, so any summary at all is real progress, + # whether or not the target was met. Without one there is nothing to hand back. + return compacted if summarized else None def _next_summary( self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter, summary_tokens: int @@ -343,24 +357,6 @@ def _next_summary( ) return oldest_steps, _CURRENT_TASK_STEPS - def _summary_limit(self) -> tuple[int, dict[str, Any]]: - """ - Work out how large one summary may be and how to hold the Chat Generator to it. - - :returns: A tuple containing: - - 1. The token budget for a single summary. This is `max_summary_tokens`, unless the generator already - configures a recognized output limit of its own, in which case the generator's setting wins. - 2. The kwargs to pass to the generator's `run`. This carries a `generation_kwargs` entry for a built-in - generator that has no limit configured, and is empty for every other generator, since the - `ChatGenerator` protocol does not standardize the setting. When it is empty, the budget reaches the - model only as prompt guidance and `_apply_summary` measures the result instead. - """ - summary_tokens, generation_kwargs = _resolve_output_token_limit( - chat_generator=self.chat_generator, default_limit=self.max_summary_tokens - ) - return summary_tokens, {"generation_kwargs": generation_kwargs} if generation_kwargs else {} - def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" transcript = _rendered_conversation(_messages_at(messages=messages, indices=indices)) From 695a56b25566ddb89bbf33581c60472fdc1d14b5 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 09:25:07 +0200 Subject: [PATCH 36/38] Add better placeholders for images and files --- haystack/hooks/compaction/summarization.py | 40 +++++++++++++-- haystack/token_counters/utils.py | 25 ++++++--- test/hooks/compaction/test_summarization.py | 57 ++++++++++++++++++++- test/token_counters/test_utils.py | 11 ++++ 4 files changed, 120 insertions(+), 13 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index da23d3b2d7f..387146deb89 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -8,7 +8,8 @@ from haystack.components.generators.chat.types import ChatGenerator from haystack.components.generators.chat.utils import _resolve_output_token_limit from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict -from haystack.dataclasses import ChatMessage +from haystack.dataclasses import ChatMessage, FileContent, ImageContent +from haystack.dataclasses.chat_message import ChatMessageContentT from haystack.hooks.compaction.types import Compactor from haystack.hooks.compaction.utils import ( _COMPACTION_META_KEY, @@ -46,10 +47,34 @@ - Exact file paths, URLs, identifiers, and references to stored data. - Unresolved work and the immediate next step. +Images and files appear only as and placeholders; their contents are not available to you \ +and will be lost. Keep the names and details shown in the placeholder, along with whatever the conversation says \ +about them, so they can be supplied again if they are needed. + Fold any existing blocks into one summary. Record only what the conversation shows. Do not \ infer or add advice. Use plain prose or short bullets, and do not address the user.""" +def _identifying_details(metadata: dict[str, Any]) -> list[str]: + """Render an attachment's metadata as `key=value` pairs, such as the path a file was loaded from.""" + # For ease, we don't support nested keys, we are mostly interested in the top-level keys that identify the + # attachment, such as a file path or URL. + return [f"{key}={value}" for key, value in sorted(metadata.items()) if isinstance(value, (str, int, float, bool))] + + +def _attachment_placeholder(content: ChatMessageContentT) -> str: + """ + Render a placeholder for an attachment that cannot survive summarization, so the summary can preserve its identity. + """ + if isinstance(content, ImageContent): + # Images have no filename, so whatever identifies one lives in its `meta`. + return f"" + if isinstance(content, FileContent): + details = [content.filename or "unnamed", content.mime_type or "unknown type"] + return f"" + return f"<{type(content).__name__}>" + + def _is_summary(message: ChatMessage) -> bool: """Whether a message is a summary this strategy wrote.""" return _is_compaction_message(message=message, strategy=_STRATEGY) @@ -160,7 +185,7 @@ def __init__( *, min_keep_steps: int = 1, max_summary_tokens: int = 1024, - summary_instruction: str | None = None, + summary_instruction: str = _DEFAULT_SUMMARY_INSTRUCTION, raise_on_failure: bool = False, ) -> None: """ @@ -171,7 +196,10 @@ def __init__( :param min_keep_steps: The fewest complete recent Agent steps to keep, even when they exceed the target. :param max_summary_tokens: The output-token budget reserved for each summary. Known built-in generators receive the corresponding runtime generation setting unless one is already configured on the generator. - :param summary_instruction: An instruction replacing the built-in summary prompt. + :param summary_instruction: What the model is told to preserve when it writes a summary. The default asks for + the user's goal, decisions and their reasoning, completed work, exact identifiers, the names of attachments + that cannot survive summarization, and the next step. The token budget is appended to whatever is given + here, so a replacement does not need to mention it. :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is logged and any successful partial compaction is returned. :raises ValueError: If `min_keep_steps` is negative or `max_summary_tokens` is not positive. @@ -183,7 +211,7 @@ def __init__( self.chat_generator = chat_generator self.min_keep_steps = min_keep_steps self.max_summary_tokens = max_summary_tokens - self.summary_instruction = summary_instruction or _DEFAULT_SUMMARY_INSTRUCTION + self.summary_instruction = summary_instruction self.raise_on_failure = raise_on_failure def compact( @@ -359,7 +387,9 @@ def _next_summary( def _prompt(self, messages: list[ChatMessage], indices: list[int], summary_tokens: int) -> list[ChatMessage]: """Build the bounded summarization instruction and the rendered transcript of the selected messages.""" - transcript = _rendered_conversation(_messages_at(messages=messages, indices=indices)) + transcript = _rendered_conversation( + _messages_at(messages=messages, indices=indices), placeholder=_attachment_placeholder + ) instruction = ( f"{self.summary_instruction}\n\nWrite a complete summary in no more than approximately " f"{summary_tokens} tokens. Prioritize completeness within that limit so the response is not cut off." diff --git a/haystack/token_counters/utils.py b/haystack/token_counters/utils.py index 9dbf9b77940..7f4d8f3ef71 100644 --- a/haystack/token_counters/utils.py +++ b/haystack/token_counters/utils.py @@ -3,11 +3,15 @@ # SPDX-License-Identifier: Apache-2.0 import json +from collections.abc import Callable from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent from haystack.dataclasses.chat_message import ChatMessageContentT, ToolCallResultContentT from haystack.tools import ToolsType, flatten_tools_or_toolsets +# Builds the stand-in for message content that has no text form, such as an image. +_PlaceholderFn = Callable[[ChatMessageContentT], str] + def _non_text_placeholder(content: ChatMessageContentT) -> str: """A short stand-in, such as ``, for message content that has no text form.""" @@ -18,26 +22,33 @@ def _non_text_placeholder(content: ChatMessageContentT) -> str: return f"<{type(content).__name__}>" -def _tool_result_text(result: ToolCallResultContentT) -> str: +def _tool_result_text(result: ToolCallResultContentT, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """A tool result as a single string, with placeholders standing in for any non-text parts.""" if isinstance(result, str): return result - return "".join(block.text if isinstance(block, TextContent) else _non_text_placeholder(block) for block in result) + return "".join(block.text if isinstance(block, TextContent) else placeholder(block) for block in result) -def _render_message(message: ChatMessage) -> str: +def _render_message(message: ChatMessage, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """ One message as one or more lines of plain text. Reasoning content is deliberately left out: providers discard it between turns, so it is not part of the context being measured. + + :param message: The message to render. + :param placeholder: Builds the stand-in for content that has no text form. The default is short and stable, which + is what a counter needs because the stand-in's own length is what gets measured. A caller rendering for a model + to read can pass one that describes the content instead. + :returns: The rendered message. """ role = message.role.value # A tool-result msg only carries tool_call_results, so it is rendered on its own and labelled with the tool that # produced it. if results := message.tool_call_results: return "\n".join( - f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] {_tool_result_text(result.result)}" + f"[tool:{result.origin.tool_name}{' (error)' if result.error else ''}] " + f"{_tool_result_text(result.result, placeholder=placeholder)}" for result in results ) @@ -50,13 +61,13 @@ def _render_message(message: ChatMessage) -> str: # Images and files cost tokens too, so they need a stand-in rather than being skipped. non_text: list[ChatMessageContentT] = [*message.images, *message.files] for content in non_text: - lines.append(f"[{role}] {_non_text_placeholder(content)}") + lines.append(f"[{role}] {placeholder(content)}") return "\n".join(lines) if lines else f"[{role}] " -def _rendered_conversation(messages: list[ChatMessage]) -> str: +def _rendered_conversation(messages: list[ChatMessage], *, placeholder: _PlaceholderFn = _non_text_placeholder) -> str: """The whole conversation as one plain-text block, which is what a counter measures.""" - return "\n".join(_render_message(message) for message in messages) + return "\n".join(_render_message(message, placeholder=placeholder) for message in messages) def _rendered_tools(tools: ToolsType | None) -> str: diff --git a/test/hooks/compaction/test_summarization.py b/test/hooks/compaction/test_summarization.py index c9d916fa56d..b2b5197125e 100644 --- a/test/hooks/compaction/test_summarization.py +++ b/test/hooks/compaction/test_summarization.py @@ -5,8 +5,9 @@ import pytest from haystack.components.generators.chat import MockChatGenerator -from haystack.dataclasses import ChatMessage, ChatRole +from haystack.dataclasses import ChatMessage, ChatRole, FileContent, ImageContent, TextContent, ToolCall from haystack.hooks.compaction import SummarizationCompactor +from haystack.hooks.compaction.summarization import _attachment_placeholder from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result @@ -77,6 +78,38 @@ def a_task_with_two_steps() -> list[ChatMessage]: ] +class TestAttachmentPlaceholder: + @pytest.mark.parametrize( + ("content", "expected"), + [ + pytest.param(ImageContent(base64_image="Zm9v", mime_type="image/png"), "", id="image"), + pytest.param( + ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"file_path": "/tmp/shot.png"}), + "", + id="image-named-by-meta", + ), + pytest.param( + FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="q3.pdf"), + "", + id="file", + ), + pytest.param( + FileContent(base64_data="Zm9v", mime_type="application/pdf", extra={"page": 4}), + "", + id="file-unnamed-with-extra", + ), + # A nested value could be arbitrarily large, so it is left out rather than bloating the prompt. + pytest.param( + ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"boxes": [[1, 2], [3, 4]]}), + "", + id="nested-metadata-left-out", + ), + ], + ) + def test_names_the_attachment(self, content, expected): + assert _attachment_placeholder(content) == expected + + class TestSummarizationCompactor: def test_summarizes_oldest_historical_turn(self): messages = two_turns_and_a_task() @@ -171,6 +204,28 @@ def test_min_keep_steps_wins_over_an_unaffordable_target(self, min_keep_steps, e result = compacted or messages assert sum(message.is_from(role=ChatRole.ASSISTANT) for message in result) == expected + def test_attachments_are_named_in_the_transcript(self): + image = ImageContent(base64_image="Zm9v", mime_type="image/png", meta={"file_path": "/tmp/shot.png"}) + pdf = FileContent(base64_data="Zm9v", mime_type="application/pdf", filename="q3.pdf") + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user(content_parts=["review this " * 20, pdf]), + tool_call("c1"), + # An attachment a tool returned is nested inside the tool result rather than on the message. + ChatMessage.from_tool( + tool_result=[TextContent(text="captured " * 20), image], + origin=ToolCall(tool_name="browse", arguments={}, id="c1"), + ), + ChatMessage.from_user("current task"), + ] + generator, prompts = summarizer("summary") + SummarizationCompactor(generator, max_summary_tokens=1).compact( + messages=messages, target_tokens=SMALLEST, token_counter=COUNTER + ) + # The summary cannot reproduce either attachment, so the transcript has to name them well enough to ask again. + assert "" in prompts[0] + assert "" in prompts[0] + def test_custom_summary_instruction_replaces_the_default(self): generator, prompts = summarizer("summary") SummarizationCompactor(generator, summary_instruction="Only list file paths.", max_summary_tokens=1).compact( diff --git a/test/token_counters/test_utils.py b/test/token_counters/test_utils.py index b3101da360b..8fb06a896a2 100644 --- a/test/token_counters/test_utils.py +++ b/test/token_counters/test_utils.py @@ -75,6 +75,17 @@ def test_joins_messages_with_newlines(self): def test_empty_conversation(self): assert _rendered_conversation([]) == "" + def test_a_custom_placeholder_reaches_nested_tool_results(self): + messages = [ + ChatMessage.from_user(content_parts=["look:", IMAGE]), + ChatMessage.from_tool( + tool_result=[TextContent(text="screenshot: "), IMAGE], + origin=ToolCall(tool_name="browse", arguments={}, id="c1"), + ), + ] + rendered = _rendered_conversation(messages, placeholder=lambda content: "") + assert rendered == "[user] look:\n[user] \n[tool:browse] screenshot: " + @tool def search(query: Annotated[str, "the search query"]) -> str: From 6ad75d1c1d2df3707e1fca40e819273d58ce25ca Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 10:39:16 +0200 Subject: [PATCH 37/38] improve default summarization prompt --- haystack/hooks/compaction/summarization.py | 49 +++++++++++++++------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/haystack/hooks/compaction/summarization.py b/haystack/hooks/compaction/summarization.py index 387146deb89..9bef7c06abf 100644 --- a/haystack/hooks/compaction/summarization.py +++ b/haystack/hooks/compaction/summarization.py @@ -39,20 +39,37 @@ _CURRENT_TASK_SUMMARIES = "current_task_summaries" _CURRENT_TASK_STEPS = "current_task_steps" -_DEFAULT_SUMMARY_INSTRUCTION = """You are compacting part of a conversation between a user and an AI agent so the \ -agent can keep working with fewer tokens. Write a self-contained summary that preserves: -- The user's goal, requirements, constraints, and preferences. -- Decisions and the reasoning behind them. -- Work already completed and important tool results. -- Exact file paths, URLs, identifiers, and references to stored data. -- Unresolved work and the immediate next step. +_DEFAULT_SUMMARY_INSTRUCTION = """You are compacting one portion of a conversation between a user and an AI agent so \ +the agent can keep working with fewer tokens. You are shown only the portion being replaced. The rest of the \ +conversation, including the user's current request, stays in place and is not shown to you. Summarize only what you \ +are given, and never say or imply that something did not happen just because it is absent from this portion. -Images and files appear only as and placeholders; their contents are not available to you \ -and will be lost. Keep the names and details shown in the placeholder, along with whatever the conversation says \ -about them, so they can be supplied again if they are needed. +Use these sections, in this order. Keep every section, and write "(none)" when this portion says nothing about it. -Fold any existing blocks into one summary. Record only what the conversation shows. Do not \ -infer or add advice. Use plain prose or short bullets, and do not address the user.""" +## Objective +What the user was trying to accomplish, if this portion shows it. + +## Decisions and constraints +Choices made and the reasoning behind them, and any requirements, preferences, or instructions the user gave. Note \ +options that were rejected and why. + +## Work completed +What was done, and what the tool results established. + +## Identifiers +Exact file paths, URLs, IDs, names, commands, and error strings, copied character for character. Images and files \ +appear only as and placeholders; their contents are not available to you and are lost once \ +this portion is replaced, so copy the placeholder details here. + +## Unresolved +Work still outstanding, and the immediate next step. + +Rules: +- Record only what this portion shows. Do not infer, do not give advice, and do not add anything that is not here. +- Copy identifiers exactly rather than describing them. They cannot be recovered once this portion is gone. +- Fold any blocks you are given into your own: keep what is still true, drop what is now \ +stale, and merge in the new facts. +- Use terse bullets. Do not address the user, and do not mention that you are summarizing.""" def _identifying_details(metadata: dict[str, Any]) -> list[str]: @@ -197,9 +214,11 @@ def __init__( :param max_summary_tokens: The output-token budget reserved for each summary. Known built-in generators receive the corresponding runtime generation setting unless one is already configured on the generator. :param summary_instruction: What the model is told to preserve when it writes a summary. The default asks for - the user's goal, decisions and their reasoning, completed work, exact identifiers, the names of attachments - that cannot survive summarization, and the next step. The token budget is appended to whatever is given - here, so a replacement does not need to mention it. + fixed sections covering the objective, decisions and constraints, completed work, exact identifiers, and + unresolved work, each written as `(none)` when the summarized portion says nothing about it. It also states + that only part of the conversation is shown, so the model does not conclude that something never happened + just because it is absent. The token budget is appended to whatever is given here, so a replacement does + not need to mention it. :param raise_on_failure: Whether a failed or non-shrinking summarization raises. By default the failure is logged and any successful partial compaction is returned. :raises ValueError: If `min_keep_steps` is negative or `max_summary_tokens` is not positive. From 2cdfceb95c6e0764fd744e44655b1f52f2cf9087 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 10:54:19 +0200 Subject: [PATCH 38/38] PR comment --- haystack/hooks/compaction/sliding_window.py | 6 +++++- test/hooks/compaction/test_sliding_window.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index c916e5958ef..7f99bb8e0b1 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -321,7 +321,7 @@ def compact( :param target_tokens: The size the kept conversation should come in under. :param token_counter: The `TokenCounter` to measure messages with. :returns: The conversation that survived, with an omission note if configured standing where the removed - messages used to sit; or None when there is nothing to remove. + messages used to sit; or None when there is nothing to remove but an earlier note. """ if token_counter.count(messages=messages) <= target_tokens: return None @@ -335,6 +335,10 @@ def compact( return None if not self.omission_note: return kept + # Swapping an earlier note for a new one frees nothing, so decline instead of rewriting the conversation again + # on every following step. + if all(_is_compaction_note(message=message) for message in removable): + return None # We prefer user over system since not all providers support multiple system messages note = ChatMessage.from_user( diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 332332b71b6..7757b59e7f1 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -292,6 +292,24 @@ def test_omission_note_can_be_customized(self, note, expected): assert compacted is not None assert compacted[2].text == expected + def test_returns_none_when_only_an_earlier_note_would_be_removed(self): + messages = [ + ChatMessage.from_system(text="rules"), + ChatMessage.from_user( + text=_DEFAULT_OMISSION_NOTE.replace("{num_removed}", "12"), + meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}}, + ), + ChatMessage.from_user(text="current task"), + ChatMessage.from_assistant(text="current step"), + ] + # Enough for everything but the earlier note, leaving that note as the only thing compaction could remove. + target_tokens = 16 + # Swapping one note for another frees nothing, so compaction must decline rather than run again every step. + assert ( + SlidingWindowCompactor().compact(messages=messages, target_tokens=target_tokens, token_counter=COUNTER) + is None + ) + @pytest.mark.parametrize( ("messages", "target_tokens"), [