Coverage for haystack/hooks/compaction/utils.py: 100%
71 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from collections.abc import Callable
7from haystack.dataclasses import ChatMessage, ChatRole
8from haystack.token_counters import TokenCounter
10# Meta key marking a message that a compactor produced. Its value records which strategy ran.
11_COMPACTION_META_KEY = "context_compaction"
14def _leading_system_end(messages: list[ChatMessage]) -> int:
15 """Return the end of the leading system-message block, excluding system messages created by compaction."""
16 for index, message in enumerate(messages):
17 if not message.is_from(role=ChatRole.SYSTEM) or _COMPACTION_META_KEY in message.meta:
18 return index
19 return len(messages)
22def _latest_user_index(messages: list[ChatMessage]) -> int | None:
23 """Return the latest user message not produced by compaction."""
24 for index in reversed(range(len(messages))):
25 message = messages[index]
26 if message.is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in message.meta:
27 return index
28 return None
31def _messages_at(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]:
32 """Return the messages at the given indices, in the order the indices are given."""
33 return [messages[index] for index in indices]
36def _messages_except(messages: list[ChatMessage], indices: list[int]) -> list[ChatMessage]:
37 """Return the messages the given indices leave out, in conversation order."""
38 left_out = set(indices)
39 return [message for index, message in enumerate(messages) if index not in left_out]
42def _is_compaction_message(message: ChatMessage, strategy: str, role: ChatRole | None = None) -> bool:
43 """Return whether a message was produced by a compaction strategy and optionally has the requested role."""
44 marker = message.meta.get(_COMPACTION_META_KEY)
45 has_role = role is None or message.is_from(role=role)
46 return has_role and isinstance(marker, dict) and marker.get("strategy") == strategy
49def _last_assistant_index(messages: list[ChatMessage]) -> int:
50 """Return the index of the last assistant message, or -1 if none exists."""
51 for index in reversed(range(len(messages))):
52 if messages[index].is_from(role=ChatRole.ASSISTANT):
53 return index
54 return -1
57def _agent_step_spans(messages: list[ChatMessage], start: int) -> list[tuple[int, int]]:
58 """
59 Return spans containing an assistant message and all immediately following tool results.
61 :param messages: The conversation to analyze, oldest to newest.
62 :param start: The index to start searching for steps from. We recommend starting after the latest user message, or
63 after the leading system messages when there is no user message. Otherwise, the returned spans may include
64 steps that are not part of the current task.
65 :returns: A list of spans, where each span is a tuple of (start_index, end_index) and end_index is exclusive.
66 """
67 # e.g. a span of (2, 5) means messages[2:5] are part of the same step
68 spans: list[tuple[int, int]] = []
69 index = start
70 while index < len(messages):
71 # Only assistant messages can start a step; skip any other messages.
72 if not messages[index].is_from(role=ChatRole.ASSISTANT):
73 index += 1
74 continue
75 # Find the end of the step by looking for the first message that is not a tool result.
76 end = index + 1
77 while end < len(messages) and messages[end].tool_call_results:
78 end += 1
79 # Record the span and continue searching for the next step.
80 spans.append((index, end))
81 # Reset the index to the end of the current step to avoid overlapping spans.
82 index = end
83 return spans
86def _current_agent_step_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]:
87 """
88 Return message-index groups for complete Agent steps belonging to the current task.
90 :param messages: The conversation to analyze, ordered oldest to newest.
91 :param system_end: The end of the leading system-message block. Steps are looked for from here when the
92 conversation has no user message to anchor on.
93 :param task_index: The index of the user message anchoring the current task, or None when there is none. Steps are
94 looked for from the message after it.
95 :returns: One group of message indices per step, ordered oldest step first. A step is an assistant message and all
96 immediately following tool results, so a group holds a tool call together with its results.
97 """
98 step_start = task_index + 1 if task_index is not None else system_end
99 return [list(range(start, end)) for start, end in _agent_step_spans(messages=messages, start=step_start)]
102def _historical_turn_spans(messages: list[ChatMessage], start: int, end: int) -> list[tuple[int, int]]:
103 """
104 Return spans for complete real-user turns in a bounded section of conversation history.
106 Each turn begins with a user message not created by compaction and continues up to the next such message.
107 """
108 user_indices = [
109 index
110 for index in range(start, end)
111 if messages[index].is_from(role=ChatRole.USER) and _COMPACTION_META_KEY not in messages[index].meta
112 ]
113 return [
114 (index, user_indices[position + 1] if position + 1 < len(user_indices) else end)
115 for position, index in enumerate(user_indices)
116 ]
119def _historical_turn_groups(messages: list[ChatMessage], system_end: int, task_index: int | None) -> list[list[int]]:
120 """
121 Return message-index groups for complete historical turns preceding the current task.
123 :param messages: The conversation to analyze, ordered oldest to newest.
124 :param system_end: The end of the leading system-message block, where the history begins.
125 :param task_index: The index of the user message anchoring the current task, where the history ends. None when
126 there is no such message, in which case the history is empty and everything belongs to the current task.
127 :returns: One group of message indices per turn, ordered oldest turn first. A turn is a user message that
128 compaction did not produce, together with everything that follows it up to the next one.
129 """
130 historical_end = task_index if task_index is not None else system_end
131 return [
132 list(range(start, end))
133 for start, end in _historical_turn_spans(messages=messages, start=system_end, end=historical_end)
134 ]
137def _replace_tool_results_until_target(
138 messages: list[ChatMessage],
139 *,
140 target_tokens: int,
141 token_counter: TokenCounter,
142 min_keep_steps: int,
143 replace: Callable[[ChatMessage, int], tuple[ChatMessage, int] | None],
144) -> list[ChatMessage] | None:
145 """
146 Replace tool results oldest first until the conversation fits within `target_tokens`.
148 Results from the most recent `min_keep_steps` tool-calling Agent steps are never passed to `replace`. After
149 measuring the initial conversation, the running total is updated with each replacement's saved tokens to avoid
150 recounting the full context.
152 :param messages: The conversation to compact, oldest to newest.
153 :param target_tokens: The token count the compacted conversation should fit within.
154 :param token_counter: The `TokenCounter` used to measure the initial conversation.
155 :param min_keep_steps: Number of most recent tool-calling Agent steps whose results are left unchanged.
156 :param replace: Called with a tool-result message and its position. Returns the replacement and the tokens it
157 saves, or None to leave the message unchanged.
158 :returns: The conversation with replaced tool results, or None when nothing was replaced.
159 """
160 current_tokens = token_counter.count(messages=messages)
161 if current_tokens <= target_tokens:
162 return None
164 # Filter the shared Agent-step spans to tool-calling steps; parallel results remain grouped in one span.
165 result_steps = [
166 list(range(start + 1, end))
167 for start, end in _agent_step_spans(messages=messages, start=0)
168 # An assistant-only span has no result to protect and must not consume one of `min_keep_steps`.
169 if end > start + 1
170 ]
171 protected_positions = {position for step in result_steps[-min_keep_steps:] for position in step}
173 # Replace entries in a new list so the caller-owned input list remains unchanged.
174 compacted = list(messages)
175 changed = False
176 # Go oldest first and stop at the target, so as much recent output as possible stays in context.
177 for index, message in enumerate(messages):
178 if message.tool_call_result is None or index in protected_positions:
179 continue
180 replacement = replace(message, index)
181 if replacement is None:
182 continue
183 compacted[index], saved_tokens = replacement
184 current_tokens -= saved_tokens
185 changed = True
186 if current_tokens <= target_tokens:
187 break
189 return compacted if changed else None