Coverage for haystack/hooks/compaction/tool_result_pruning.py: 100%
36 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +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, _replace_tool_results_until_target
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 return _replace_tool_results_until_target(
100 messages=messages,
101 target_tokens=target_tokens,
102 token_counter=token_counter,
103 min_keep_steps=self.min_keep_steps,
104 replace=lambda message, _index: self._prune(message=message, token_counter=token_counter),
105 )
107 def _prune(self, message: ChatMessage, token_counter: TokenCounter) -> tuple[ChatMessage, int] | None:
108 """
109 Rewrite one tool-result message with a placeholder, or return None to leave it as it is.
111 The result's `origin` is carried over so the message keeps pointing at the tool call it answers, and its error
112 flag is preserved.
114 :param message: The tool-result message to consider.
115 :param token_counter: The `TokenCounter` used to determine whether the result exceeds `min_tokens`.
116 :returns: The rewritten message and number of tokens it saves, or None when this result is not prunable.
117 """
118 result = message.tool_call_result
119 # Guard against messages that are not tool results or have already been compacted.
120 if result is None or _COMPACTION_META_KEY in message.meta:
121 return None
123 # Guard against messages that have meta keys indicating they should be skipped.
124 if any(key in message.meta for key in self.skip_meta_keys):
125 return None
127 original_tokens = token_counter.count(messages=[message])
128 # If the result is small enough, leave it alone.
129 if original_tokens <= self.min_tokens:
130 return None
132 # `replace` permits custom placeholders to contain unrelated braces, such as JSON examples.
133 placeholder = self.placeholder.replace("{tool_name}", result.origin.tool_name)
134 pruned = ChatMessage.from_tool(
135 tool_result=placeholder,
136 origin=result.origin,
137 error=result.error,
138 meta={
139 **message.meta,
140 _COMPACTION_META_KEY: {"strategy": "tool_result_pruning", "original_tokens": original_tokens},
141 },
142 )
143 saved_tokens = original_tokens - token_counter.count(messages=[pruned])
144 # A custom placeholder can be larger than the original result; only return replacements that reduce context.
145 return (pruned, saved_tokens) if saved_tokens > 0 else None
147 def to_dict(self) -> dict[str, Any]:
148 """
149 Serialize the compactor.
151 :returns: A dictionary representation of the compactor.
152 """
153 return default_to_dict(
154 self,
155 min_keep_steps=self.min_keep_steps,
156 min_tokens=self.min_tokens,
157 placeholder=self.placeholder,
158 skip_meta_keys=list(self.skip_meta_keys),
159 )