Coverage for haystack/hooks/compaction/tool_result_pruning.py: 100%
55 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from typing import Any
7from haystack.core.serialization import default_to_dict
8from haystack.dataclasses import ChatMessage
9from haystack.hooks.compaction.types import Compactor
10from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _agent_step_spans
11from haystack.token_counters import TokenCounter
12from haystack.utils.experimental import _experimental
14_DEFAULT_PLACEHOLDER = "[Tool result removed to free up context. Call `{tool_name}` again if you need it.]"
17@_experimental
18class ToolResultPruningCompactor(Compactor):
19 """
20 Replaces the content of older tool results with a short placeholder, keeping the conversation's shape intact.
22 Tool output usually dominates a long Agent run, and most of it stops being useful once the model has acted on it.
23 This compactor rewrites those results in place rather than removing messages, so every tool call keeps its matching
24 result and the model can see what it ran and re-run it if needed.
26 <!-- test-ignore -->
27 ```python
28 from haystack.components.agents import Agent
29 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
30 from haystack.hooks.compaction import CompactionHook, ToolResultPruningCompactor
32 hook = CompactionHook(
33 compactor=ToolResultPruningCompactor(min_keep_steps=1),
34 context_window=400_000,
35 compact_at=0.7,
36 compact_to=0.4,
37 )
38 agent = Agent(
39 chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.4-nano"),
40 tools=[web_search],
41 hooks={"before_llm": [hook]},
42 )
43 ```
44 """
46 def __init__(
47 self,
48 *,
49 min_keep_steps: int = 1,
50 min_tokens: int = 200,
51 placeholder: str = _DEFAULT_PLACEHOLDER,
52 skip_meta_keys: tuple[str, ...] = ("tool_result_offloaded",),
53 ) -> None:
54 """
55 Initialize the compactor with the rules deciding which results it prunes.
57 :param min_keep_steps: The minimum number of recent tool-calling Agent steps whose results remain untouched,
58 even when they exceed the target. Must be at least 1, which ensures the current result batch remains intact
59 until the model has acted on it.
60 :param min_tokens: Only prune tool-result messages that use more than this many tokens. Small results cost
61 little and are often the ones worth keeping.
62 :param placeholder: The text left in place of a pruned result, replacing the built-in one. May contain
63 `{tool_name}`, which is filled in with the name of the tool that produced the result.
64 :param skip_meta_keys: Results whose `meta` contains any of these keys are left alone. The default covers
65 results that a `ToolResultOffloadHook` already replaced with a reference to stored content: pruning one of
66 those would destroy the reference the model needs to read it back.
67 :raises ValueError: If `min_keep_steps` is less than 1 or `min_tokens` is negative.
68 """
69 if min_keep_steps < 1:
70 raise ValueError(
71 f"`min_keep_steps` must be at least 1, got {min_keep_steps}. The most recent tool-calling step "
72 f"contains results the model may still need."
73 )
74 if min_tokens < 0:
75 raise ValueError(f"`min_tokens` must be at least 0, got {min_tokens}.")
76 self.min_keep_steps = min_keep_steps
77 self.min_tokens = min_tokens
78 self.placeholder = placeholder
79 # Normalized to a tuple so a round trip through `to_dict`, which has to emit a list, restores the same type.
80 self.skip_meta_keys = tuple(skip_meta_keys)
82 def compact(
83 self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter
84 ) -> list[ChatMessage] | None:
85 """
86 Replace the content of prunable tool results with a placeholder.
88 Results are considered oldest first and pruning stops as soon as the conversation reaches `target_tokens`.
89 This keeps as much original output as possible. Results from the most recent `min_keep_steps` tool-calling
90 Agent steps are never considered, even when the target cannot otherwise be reached. After measuring the initial
91 conversation, the running total is updated with per-result token deltas to avoid repeatedly counting the full
92 context.
94 :param messages: The conversation to compact, oldest to newest.
95 :param target_tokens: The size the compacted conversation should come in under.
96 :param token_counter: The `TokenCounter` used to measure the conversation before and after each replacement.
97 :returns: The conversation with older tool results replaced, or None when no result was prunable.
98 """
99 current_tokens = token_counter.count(messages=messages)
100 if current_tokens <= target_tokens:
101 return None
103 # Filter the shared Agent-step spans to tool-calling steps; parallel results remain grouped in one span.
104 result_steps = [
105 list(range(start + 1, end))
106 for start, end in _agent_step_spans(messages=messages, start=0)
107 # An assistant-only span has no result to protect and must not consume one of `min_keep_steps`.
108 if end > start + 1
109 ]
111 protected_positions = {position for step in result_steps[-self.min_keep_steps :] for position in step}
113 # Replace entries in a new list so the caller-owned input list remains unchanged.
114 compacted = list(messages)
115 changed = False
116 # Iterate oldest-first so we can stop at the target while leaving as much recent output intact as possible.
117 for index, message in enumerate(messages):
118 if message.tool_call_result is None or index in protected_positions:
119 continue
120 replacement = self._prune(message=message, token_counter=token_counter)
121 if replacement is None:
122 continue
123 pruned, saved_tokens = replacement
125 # Update the compacted list and the running token count
126 compacted[index] = pruned
127 current_tokens -= saved_tokens
128 changed = True
130 # Once we reach the target, stop pruning to preserve as much recent output as possible.
131 if current_tokens <= target_tokens:
132 break
134 # If no candidates were pruned, return None to indicate no change.
135 return compacted if changed else None
137 def _prune(self, message: ChatMessage, token_counter: TokenCounter) -> tuple[ChatMessage, int] | None:
138 """
139 Rewrite one tool-result message with a placeholder, or return None to leave it as it is.
141 The result's `origin` is carried over so the message keeps pointing at the tool call it answers, and its error
142 flag is preserved.
144 :param message: The tool-result message to consider.
145 :param token_counter: The `TokenCounter` used to determine whether the result exceeds `min_tokens`.
146 :returns: The rewritten message and number of tokens it saves, or None when this result is not prunable.
147 """
148 result = message.tool_call_result
149 # Guard against messages that are not tool results or have already been compacted.
150 if result is None or _COMPACTION_META_KEY in message.meta:
151 return None
153 # Guard against messages that have meta keys indicating they should be skipped.
154 if any(key in message.meta for key in self.skip_meta_keys):
155 return None
157 original_tokens = token_counter.count(messages=[message])
158 # If the result is small enough, leave it alone.
159 if original_tokens <= self.min_tokens:
160 return None
162 # `replace` permits custom placeholders to contain unrelated braces, such as JSON examples.
163 placeholder = self.placeholder.replace("{tool_name}", result.origin.tool_name)
164 pruned = ChatMessage.from_tool(
165 tool_result=placeholder,
166 origin=result.origin,
167 error=result.error,
168 meta={
169 **message.meta,
170 _COMPACTION_META_KEY: {"strategy": "tool_result_pruning", "original_tokens": original_tokens},
171 },
172 )
173 saved_tokens = original_tokens - token_counter.count(messages=[pruned])
174 # A custom placeholder can be larger than the original result; only return replacements that reduce context.
175 return (pruned, saved_tokens) if saved_tokens > 0 else None
177 def to_dict(self) -> dict[str, Any]:
178 """
179 Serialize the compactor.
181 :returns: A dictionary representation of the compactor.
182 """
183 return default_to_dict(
184 self,
185 min_keep_steps=self.min_keep_steps,
186 min_tokens=self.min_tokens,
187 placeholder=self.placeholder,
188 skip_meta_keys=list(self.skip_meta_keys),
189 )