From c046447d3bffc93a8c34d88dd336c079a11999e4 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Tue, 11 Aug 2026 15:35:57 +0200 Subject: [PATCH 1/2] Moving shared utils over --- haystack/hooks/compaction/sliding_window.py | 106 +++-------------- haystack/hooks/compaction/utils.py | 67 +++++++++++ test/hooks/compaction/test_sliding_window.py | 43 +------ test/hooks/compaction/test_utils.py | 115 ++++++++++++++++++- 4 files changed, 197 insertions(+), 134 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 7f99bb8e0b1..492fb27a895 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -7,7 +7,16 @@ 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_at, + _messages_except, +) from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental @@ -23,86 +32,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 - ] - - -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) - return [message for index, message in enumerate(messages) if index not in left_out] + return _is_compaction_message(message=message, strategy=_STRATEGY, role=ChatRole.USER) def _flatten(groups: list[list[int]]) -> list[int]: @@ -124,19 +56,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/utils.py b/haystack/hooks/compaction/utils.py index 6a6d1402c52..590ddb84b2d 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -10,6 +10,41 @@ _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_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) + 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 +82,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/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index 7757b59e7f1..f5da9be89c6 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, 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.sliding_window import _DEFAULT_OMISSION_NOTE, _is_compaction_note from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from test.hooks.compaction.helpers import ( FakeCounter, @@ -69,47 +69,6 @@ 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 = [ - 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): - 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"), - ChatMessage.from_user("next task"), - ChatMessage.from_assistant("second step"), - ] - # 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 = [ diff --git a/test/hooks/compaction/test_utils.py b/test/hooks/compaction/test_utils.py index 05ce301aa2b..c17d09e404c 100644 --- a/test/hooks/compaction/test_utils.py +++ b/test/hooks/compaction/test_utils.py @@ -4,8 +4,17 @@ import pytest -from haystack.dataclasses import ChatMessage -from haystack.hooks.compaction.utils import _agent_step_spans, _estimated_context_tokens, _last_assistant_index +from haystack.dataclasses import ChatMessage, ChatRole +from haystack.hooks.compaction.utils import ( + _COMPACTION_META_KEY, + _agent_step_spans, + _current_step_groups, + _estimated_context_tokens, + _historical_turn_groups, + _historical_turn_spans, + _is_compaction_message, + _last_assistant_index, +) from haystack.tools import tool from test.hooks.compaction.helpers import FakeCounter, tool_call, tool_result @@ -71,6 +80,108 @@ def test_starts_at_the_requested_message(self): assert _agent_step_spans(messages=messages, start=2) == [(2, 3)] +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): + 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"), + ChatMessage.from_user("next task"), + ChatMessage.from_assistant("second step"), + ] + # 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 TestIsCompactionMessage: + @pytest.mark.parametrize( + ("strategy", "role", "expected"), + [ + pytest.param("sliding_window", None, True, id="matching-strategy-any-role"), + pytest.param("summarization", None, False, id="another-strategy"), + pytest.param("sliding_window", ChatRole.USER, True, id="matching-strategy-and-role"), + pytest.param("sliding_window", ChatRole.SYSTEM, False, id="matching-strategy-wrong-role"), + ], + ) + def test_matches_a_strategy_and_an_optional_role(self, strategy, role, expected): + note = ChatMessage.from_user(text="removed", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}}) + assert _is_compaction_message(message=note, strategy=strategy, role=role) is expected + + @pytest.mark.parametrize( + "message", + [ + pytest.param(ChatMessage.from_user(text="hi"), id="no-marker"), + # A marker that is not a mapping cannot carry a strategy, so it matches nothing. + pytest.param( + ChatMessage.from_user(text="odd", meta={_COMPACTION_META_KEY: "sliding_window"}), + id="marker-that-is-not-a-mapping", + ), + ], + ) + def test_does_not_match_without_a_usable_marker(self, message): + assert _is_compaction_message(message=message, strategy="sliding_window") is False + + +class TestHistoricalTurnGroups: + def test_covers_the_turns_between_the_system_block_and_the_task(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("old question"), + ChatMessage.from_assistant("old answer"), + ChatMessage.from_user("current task"), + ] + assert _historical_turn_groups(messages=messages, system_end=1, task_index=3) == [[1, 2]] + + def test_there_is_no_history_without_a_task_anchor(self): + # With no user message to anchor on, everything after the system block belongs to the current task instead. + messages = [ChatMessage.from_system("rules"), ChatMessage.from_assistant("step")] + assert _historical_turn_groups(messages=messages, system_end=1, task_index=None) == [] + + +class TestCurrentStepGroups: + def test_covers_the_steps_after_the_task_anchor(self): + messages = [ + ChatMessage.from_system("rules"), + ChatMessage.from_user("current task"), + tool_call("c1"), + tool_result("result", call_id="c1"), + ChatMessage.from_assistant("answer"), + ] + assert _current_step_groups(messages=messages, system_end=1, task_index=1) == [[2, 3], [4]] + + def test_starts_after_the_system_block_without_a_task_anchor(self): + messages = [ChatMessage.from_system("rules"), ChatMessage.from_assistant("step")] + assert _current_step_groups(messages=messages, system_end=1, task_index=None) == [[1]] + + class TestEstimatedContextTokens: def test_counts_only_what_the_generator_has_not_seen(self): counter = FakeCounter() From 082a93bedb420c6760220cd685622f60915fdff5 Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Wed, 12 Aug 2026 08:57:19 +0200 Subject: [PATCH 2/2] refactoring --- haystack/hooks/compaction/sliding_window.py | 28 +++++++++++--------- haystack/hooks/compaction/utils.py | 25 ++++++++++++++--- test/hooks/compaction/test_sliding_window.py | 2 +- test/hooks/compaction/test_utils.py | 24 ++++++++--------- 4 files changed, 50 insertions(+), 29 deletions(-) diff --git a/haystack/hooks/compaction/sliding_window.py b/haystack/hooks/compaction/sliding_window.py index 492fb27a895..01f0a9b358d 100644 --- a/haystack/hooks/compaction/sliding_window.py +++ b/haystack/hooks/compaction/sliding_window.py @@ -9,7 +9,7 @@ from haystack.hooks.compaction.types import Compactor from haystack.hooks.compaction.utils import ( _COMPACTION_META_KEY, - _current_step_groups, + _current_agent_step_groups, _historical_turn_groups, _is_compaction_message, _latest_user_index, @@ -62,8 +62,8 @@ def _removable_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 + agent_step_groups = _current_agent_step_groups(messages=messages, system_end=system_end, task_index=task_index) + return historical_groups, agent_step_groups def _first_group_to_keep( @@ -96,7 +96,7 @@ def _first_group_to_keep( def _first_turn_and_step_to_keep( messages: list[ChatMessage], historical_groups: list[list[int]], - step_groups: list[list[int]], + agent_step_groups: list[list[int]], available_tokens: int, token_counter: TokenCounter, min_keep_steps: int, @@ -106,24 +106,24 @@ def _first_turn_and_step_to_keep( :param messages: The full conversation containing the messages referenced by both group lists. :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 agent_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 `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. + :returns: The position in `historical_groups` and the position in `agent_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)) + messages=_messages_at(messages=messages, indices=_flatten(groups=agent_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 + messages=messages, groups=agent_step_groups, available_tokens=available_tokens, token_counter=token_counter ) # 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)) + return len(historical_groups), min(first_kept_step, max(len(agent_step_groups) - min_keep_steps, 0)) # 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( @@ -164,14 +164,16 @@ def _task_and_step_split( task_indices = [task_index] if task_index is not None else [] # 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) + historical_groups, agent_step_groups = _removable_groups( + messages=messages, system_end=system_end, task_index=task_index + ) # 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, + agent_step_groups=agent_step_groups, available_tokens=target_tokens - token_counter.count(messages=protected), token_counter=token_counter, min_keep_steps=min_keep_steps, @@ -179,7 +181,7 @@ def _task_and_step_split( # 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_step_indices = _flatten(groups=agent_step_groups[first_kept_step:]) 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 diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 590ddb84b2d..70177416237 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -82,8 +82,18 @@ 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.""" +def _current_agent_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. + + :param messages: The conversation to analyze, ordered oldest to newest. + :param system_end: The end of the leading system-message block. Steps are looked for from here when the + conversation has no user message to anchor on. + :param task_index: The index of the user message anchoring the current task, or None when there is none. Steps are + looked for from the message after it. + :returns: One group of message indices per step, ordered oldest step first. A step is an assistant message and all + immediately following tool results, so a group holds a tool call together with its results. + """ 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)] @@ -106,7 +116,16 @@ def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> 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.""" + """ + Return message-index groups for complete historical turns preceding the current task. + + :param messages: The conversation to analyze, ordered oldest to newest. + :param system_end: The end of the leading system-message block, where the history begins. + :param task_index: The index of the user message anchoring the current task, where the history ends. None when + there is no such message, in which case the history is empty and everything belongs to the current task. + :returns: One group of message indices per turn, ordered oldest turn first. A turn is a user message that + compaction did not produce, together with everything that follows it up to the next one. + """ historical_end = task_index if task_index is not None else system_end return [ list(range(start, end)) diff --git a/test/hooks/compaction/test_sliding_window.py b/test/hooks/compaction/test_sliding_window.py index f5da9be89c6..09f140dc2b6 100644 --- a/test/hooks/compaction/test_sliding_window.py +++ b/test/hooks/compaction/test_sliding_window.py @@ -61,7 +61,7 @@ class TestIsCompactionNote: pytest.param( ChatMessage.from_user(text="odd", meta={_COMPACTION_META_KEY: "sliding_window"}), False, - id="marker-that-is-not-a-mapping", + id="marker-that-is-not-a-dict", ), ], ) diff --git a/test/hooks/compaction/test_utils.py b/test/hooks/compaction/test_utils.py index c17d09e404c..8c7ce5780a5 100644 --- a/test/hooks/compaction/test_utils.py +++ b/test/hooks/compaction/test_utils.py @@ -8,7 +8,7 @@ from haystack.hooks.compaction.utils import ( _COMPACTION_META_KEY, _agent_step_spans, - _current_step_groups, + _current_agent_step_groups, _estimated_context_tokens, _historical_turn_groups, _historical_turn_spans, @@ -131,7 +131,7 @@ class TestIsCompactionMessage: pytest.param("sliding_window", ChatRole.SYSTEM, False, id="matching-strategy-wrong-role"), ], ) - def test_matches_a_strategy_and_an_optional_role(self, strategy, role, expected): + def test_strategy_and_role(self, strategy, role, expected): note = ChatMessage.from_user(text="removed", meta={_COMPACTION_META_KEY: {"strategy": "sliding_window"}}) assert _is_compaction_message(message=note, strategy=strategy, role=role) is expected @@ -139,19 +139,19 @@ def test_matches_a_strategy_and_an_optional_role(self, strategy, role, expected) "message", [ pytest.param(ChatMessage.from_user(text="hi"), id="no-marker"), - # A marker that is not a mapping cannot carry a strategy, so it matches nothing. + # A marker that is not a dict cannot carry a strategy, so it matches nothing. pytest.param( ChatMessage.from_user(text="odd", meta={_COMPACTION_META_KEY: "sliding_window"}), - id="marker-that-is-not-a-mapping", + id="marker-that-is-not-a-dict", ), ], ) - def test_does_not_match_without_a_usable_marker(self, message): + def test_unusable_marker(self, message): assert _is_compaction_message(message=message, strategy="sliding_window") is False class TestHistoricalTurnGroups: - def test_covers_the_turns_between_the_system_block_and_the_task(self): + def test_basic(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("old question"), @@ -160,14 +160,14 @@ def test_covers_the_turns_between_the_system_block_and_the_task(self): ] assert _historical_turn_groups(messages=messages, system_end=1, task_index=3) == [[1, 2]] - def test_there_is_no_history_without_a_task_anchor(self): + def test_missing_task_anchor(self): # With no user message to anchor on, everything after the system block belongs to the current task instead. messages = [ChatMessage.from_system("rules"), ChatMessage.from_assistant("step")] assert _historical_turn_groups(messages=messages, system_end=1, task_index=None) == [] -class TestCurrentStepGroups: - def test_covers_the_steps_after_the_task_anchor(self): +class TestCurrentAgentStepGroups: + def test_basic(self): messages = [ ChatMessage.from_system("rules"), ChatMessage.from_user("current task"), @@ -175,11 +175,11 @@ def test_covers_the_steps_after_the_task_anchor(self): tool_result("result", call_id="c1"), ChatMessage.from_assistant("answer"), ] - assert _current_step_groups(messages=messages, system_end=1, task_index=1) == [[2, 3], [4]] + assert _current_agent_step_groups(messages=messages, system_end=1, task_index=1) == [[2, 3], [4]] - def test_starts_after_the_system_block_without_a_task_anchor(self): + def test_missing_task_anchor(self): messages = [ChatMessage.from_system("rules"), ChatMessage.from_assistant("step")] - assert _current_step_groups(messages=messages, system_end=1, task_index=None) == [[1]] + assert _current_agent_step_groups(messages=messages, system_end=1, task_index=None) == [[1]] class TestEstimatedContextTokens: