refactor: centralizing shared compaction utils

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_agent_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,20 +56,14 @@ 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,
-    )
-    return historical_groups, step_groups
+    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)
+    ]
+    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(
@@ -170,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,
@@ -180,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(
@@ -238,22 +164,24 @@ 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,
     )
 
     # 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

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,57 @@ def _agent_step_spans(messages: list[ChatMessage], start: int) -> list[tuple[int
     return spans
 
 
+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)]
+
+
+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.
+
+    :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))
+        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:

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,
@@ -61,55 +61,14 @@ 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",
             ),
         ],
     )
     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 = [