Coverage for haystack/hooks/compaction/hooks.py: 99%
111 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 import logging, tracing
8from haystack.components.agents.state.state import State
9from haystack.components.agents.state.state_utils import replace_values
10from haystack.core.serialization import (
11 component_to_dict,
12 default_from_dict,
13 default_to_dict,
14 generate_qualified_class_name,
15)
16from haystack.dataclasses import ChatMessage
17from haystack.hooks.compaction.types import Compactor
18from haystack.hooks.compaction.utils import _last_assistant_index
19from haystack.token_counters import ApproximateTokenCounter, TokenCounter
20from haystack.tools import ToolsType
21from haystack.utils.deserialization import deserialize_component_inplace
22from haystack.utils.experimental import _experimental
24logger = logging.getLogger(__name__)
27def _estimated_context_tokens(
28 messages: list[ChatMessage], context_tokens: int, token_counter: TokenCounter, tools: ToolsType | None = None
29) -> int:
30 """
31 Estimate the size of the whole conversation.
33 :param messages: The conversation, oldest to newest.
34 :param context_tokens: The `context_tokens` state key, anchored on the provider's own token counting. It accounts
35 for the conversation through the last assistant message, or for the whole conversation when there is none.
36 :param token_counter: The counter to measure the unaccounted messages with.
37 :param tools: Tools whose schemas are sent alongside the messages. These are counted when provider usage is absent.
38 :returns: The estimated total token count.
39 """
40 # Nothing sent yet, or a generator that reports no usage, so count everything.
41 if context_tokens == 0:
42 return token_counter.count(messages=messages, tools=tools)
43 # `context_tokens` accounts for the conversation through the last assistant message, so only the tool result
44 # messages after it still need estimating. If there is no assistant message, `context_tokens` accounts for the
45 # whole conversation, so nothing more needs counting.
46 last_assistant_index = _last_assistant_index(messages=messages)
47 if last_assistant_index < 0:
48 return context_tokens
49 tool_result_messages = messages[last_assistant_index + 1 :]
50 return context_tokens + token_counter.count(messages=tool_result_messages)
53@_experimental
54class CompactionHook:
55 """
56 Compacts an Agent's conversation once it fills too much of the model's context window.
58 This `before_llm` Agent hook estimates the size of the conversation before each chat-generator call and, once it
59 reaches `compact_at` of the window, hands it to a `Compactor` to bring back down to `compact_to`. Register it on an
60 `Agent` under the `before_llm` hook point:
62 <!-- test-ignore -->
63 ```python
64 from haystack.components.agents import Agent
65 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
66 from haystack.hooks.compaction import CompactionHook, SlidingWindowCompactor
68 hook = CompactionHook(
69 compactor=SlidingWindowCompactor(),
70 context_window=400_000,
71 compact_at=0.7,
72 compact_to=0.4,
73 )
74 agent = Agent(
75 chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.4-nano"),
76 tools=[web_search],
77 hooks={"before_llm": [hook]},
78 max_agent_steps=50,
79 )
80 ```
82 Size is measured by anchoring on the `context_tokens` state key - the chat generator's own count of the request it
83 was sent plus its reply, which already covers the system prompt, the tool schemas, and the provider's chat-template
84 overhead - and counting only the messages appended since that call. The estimate is therefore exact for the bulk of
85 the conversation and approximate only for its most recent messages.
87 Compaction is lossy by nature, so the Agent works from a shorter record of the run afterwards. What survives is up
88 to the compactor.
89 """
91 allowed_hook_points = ("before_llm",)
93 def __init__(
94 self,
95 compactor: Compactor,
96 *,
97 context_window: int,
98 compact_at: float = 0.7,
99 compact_to: float = 0.4,
100 token_counter: TokenCounter | None = None,
101 ) -> None:
102 """
103 Initialize the hook with a compactor and the window it has to fit in.
105 :param compactor: The `Compactor` that rewrites the conversation.
106 :param context_window: The model's context window in tokens. Everything else is a fraction of this, so moving to
107 a different model means changing only this number.
108 :param compact_at: The fraction of the window at which compaction starts. Leave room above it for the reply and
109 the tool results it triggers, which land on top of what was measured.
110 :param compact_to: The fraction of the window compaction aims to bring the conversation down to. Lower means
111 compacting less often but losing more each time.
112 :param token_counter: The `TokenCounter` used to size the messages the chat generator has not reported on yet.
113 Defaults to `ApproximateTokenCounter`, which needs no extra dependency.
114 :raises ValueError: If `context_window` is not positive, or the fractions are not
115 `0 < compact_to < compact_at <= 1`.
116 """
117 if context_window < 1:
118 raise ValueError(f"`context_window` must be a positive number of tokens, got {context_window}.")
119 if not 0 < compact_to < compact_at <= 1:
120 raise ValueError(
121 f"`compact_at` and `compact_to` must satisfy 0 < compact_to < compact_at <= 1, got "
122 f"compact_at={compact_at} and compact_to={compact_to}. A target at or above the trigger would leave "
123 f"the conversation over the trigger after compacting, so it would be attempted again every step."
124 )
125 self.compactor = compactor
126 self.context_window = context_window
127 self.compact_at = compact_at
128 self.compact_to = compact_to
129 self.token_counter = token_counter or ApproximateTokenCounter()
131 def run(self, state: State) -> None:
132 """
133 Compact `state.data["messages"]` if the conversation fills too much of the window.
135 :param state: The Agent's live `State`. Read to decide whether to compact, and rewritten in place when the
136 compactor returns a compacted conversation.
137 :returns: None. The hook mutates `state` in place.
138 """
139 messages = state.data.get("messages", [])
140 context_tokens = state.data.get("context_tokens", 0)
141 target = self._target_tokens(messages=messages, context_tokens=context_tokens, tools=state.data.get("tools"))
142 if target is None:
143 return
144 target_tokens, estimated_overhead = target
145 compacted = self.compactor.compact(
146 messages=messages, target_tokens=target_tokens, token_counter=self.token_counter
147 )
148 self._apply(
149 state=state,
150 compacted=compacted,
151 before=len(messages),
152 target_tokens=target_tokens,
153 original_context_tokens=context_tokens,
154 estimated_overhead=estimated_overhead,
155 )
157 async def run_async(self, state: State) -> None:
158 """
159 Asynchronously compact `state.data["messages"]` if the conversation fills too much of the window.
161 :param state: The Agent's live `State`. Read to decide whether to compact, and rewritten in place when the
162 compactor returns a compacted conversation.
163 :returns: None. The hook mutates `state` in place.
164 """
165 messages = state.data.get("messages", [])
166 context_tokens = state.data.get("context_tokens", 0)
167 target = self._target_tokens(messages=messages, context_tokens=context_tokens, tools=state.data.get("tools"))
168 if target is None:
169 return
170 target_tokens, estimated_overhead = target
171 compacted = await self.compactor.compact_async(
172 messages=messages, target_tokens=target_tokens, token_counter=self.token_counter
173 )
174 self._apply(
175 state=state,
176 compacted=compacted,
177 before=len(messages),
178 target_tokens=target_tokens,
179 original_context_tokens=context_tokens,
180 estimated_overhead=estimated_overhead,
181 )
183 def _target_tokens(
184 self, messages: list[ChatMessage], context_tokens: int, tools: ToolsType | None = None
185 ) -> tuple[int, int] | None:
186 """
187 The size a compactor should bring `messages` down to.
189 :param messages: The conversation, oldest to newest.
190 :param context_tokens: The `context_tokens` state key. It is the current token count of the conversation except
191 for the most recent tool result messages.
192 :param tools: Tools whose schemas are sent alongside the messages.
193 :returns: The target token amount the messages should be compacted to and the estimated non-message overhead,
194 or None when the conversation is not yet large enough to compact.
195 """
196 # Estimate the total context size, including the provider's own count of the conversation through the last
197 # assistant message and the tool schemas.
198 estimated = _estimated_context_tokens(
199 messages=messages, context_tokens=context_tokens, token_counter=self.token_counter, tools=tools
200 )
201 triggered = estimated >= self.context_window * self.compact_at
203 span = tracing.tracer.current_span()
204 if span is not None:
205 span.set_tags(
206 tags={
207 "haystack.agent.hook.compaction.strategy": generate_qualified_class_name(cls=type(self.compactor)),
208 "haystack.agent.hook.compaction.estimated_context_tokens": estimated,
209 "haystack.agent.hook.compaction.triggered": triggered,
210 }
211 )
213 # The conversation is not yet large enough to compact, so leave it alone.
214 if not triggered:
215 return None
217 # Calculate the non-message overhead, such as tool schemas and the provider's chat-template overhead.
218 message_tokens = self.token_counter.count(messages=messages)
219 overhead = estimated - message_tokens
220 if overhead < 0:
221 logger.warning(
222 "The TokenCounter estimated more tokens for the messages ({message_tokens}) than the estimated total "
223 "context size ({estimated}). It may be overestimating; consider configuring a more accurate "
224 "TokenCounter.",
225 message_tokens=message_tokens,
226 estimated=estimated,
227 )
228 target_tokens = max(int(self.context_window * self.compact_to) - overhead, 0)
230 if span is not None:
231 span.set_tag(key="haystack.agent.hook.compaction.target_tokens", value=target_tokens)
233 # Return the target token amount the messages should be compacted to and the overhead needed to re-estimate the
234 # context size after compaction.
235 return target_tokens, overhead
237 def _apply(
238 self,
239 state: State,
240 compacted: list[ChatMessage] | None,
241 before: int,
242 target_tokens: int,
243 original_context_tokens: int,
244 estimated_overhead: int,
245 ) -> None:
246 """
247 Write a compacted conversation back into `State`, or do nothing if the compactor declined.
249 :param state: The Agent's live `State`, rewritten in place.
250 :param compacted: What the compactor returned.
251 :param before: How many messages the conversation held before compaction, for the log line.
252 :param target_tokens: The target the compactor was given, for the log line.
253 :param original_context_tokens: The provider-reported context count before compaction, or 0 when unavailable.
254 :param estimated_overhead: The estimated non-message overhead to include in the updated context count.
255 """
256 span = tracing.tracer.current_span()
257 if span is not None:
258 span.set_tag(key="haystack.agent.hook.compaction.compacted", value=compacted is not None)
260 # If the compactor returned None, it declined to compact. Leave the conversation alone.
261 if compacted is None:
262 return
264 state.set("messages", compacted, handler_override=replace_values)
265 # If the original value was 0, leave it unchanged so later steps keep recounting the full request locally.
266 if original_context_tokens != 0:
267 # Re-estimate the provider-accounted context through the last assistant message, including overhead. If we
268 # added the trailing tool results, a second registered hook could double count them. If there is no
269 # assistant message, count everything.
270 last_assistant_index = _last_assistant_index(messages=compacted)
271 accounted = compacted if last_assistant_index < 0 else compacted[: last_assistant_index + 1]
272 state.set("context_tokens", self.token_counter.count(messages=accounted) + estimated_overhead)
273 logger.debug(
274 "Compacted the Agent's conversation at step {step} from {before} to {after} messages, targeting {target} "
275 "tokens.",
276 step=state.data.get("step_count", 0),
277 before=before,
278 after=len(compacted),
279 target=target_tokens,
280 )
282 def warm_up(self) -> None:
283 """Warm up the token counter and the compactor, which may hold resources such as a Chat Generator."""
284 if hasattr(self.token_counter, "warm_up"):
285 self.token_counter.warm_up()
286 if hasattr(self.compactor, "warm_up"):
287 self.compactor.warm_up()
289 async def warm_up_async(self) -> None:
290 """Warm up the token counter and the compactor on the serving event loop."""
291 if hasattr(self.token_counter, "warm_up"):
292 self.token_counter.warm_up()
293 if hasattr(self.compactor, "warm_up_async"):
294 await self.compactor.warm_up_async()
295 elif hasattr(self.compactor, "warm_up"):
296 self.compactor.warm_up()
298 def close(self) -> None:
299 """Release the token counter's and the compactor's resources."""
300 for resource in (self.token_counter, self.compactor):
301 if hasattr(resource, "close"):
302 resource.close()
304 async def close_async(self) -> None:
305 """Release the token counter's and the compactor's async resources."""
306 for resource in (self.token_counter, self.compactor):
307 if hasattr(resource, "close_async"):
308 await resource.close_async()
309 elif hasattr(resource, "close"):
310 resource.close()
312 def to_dict(self) -> dict[str, Any]:
313 """
314 Serialize the hook, including its compactor and token counter.
316 :returns: A dictionary representation of the hook.
317 """
318 return default_to_dict(
319 self,
320 compactor=component_to_dict(obj=self.compactor, name="compactor"),
321 context_window=self.context_window,
322 compact_at=self.compact_at,
323 compact_to=self.compact_to,
324 token_counter=component_to_dict(obj=self.token_counter, name="token_counter"),
325 )
327 @classmethod
328 def from_dict(cls, data: dict[str, Any]) -> "CompactionHook":
329 """
330 Deserialize the hook, reconstructing its compactor and token counter.
332 :param data: A dictionary representation produced by `to_dict`.
333 :returns: The deserialized `CompactionHook`.
334 """
335 init_params = data.get("init_parameters", {})
336 for key in ("compactor", "token_counter"):
337 if init_params.get(key) is not None:
338 deserialize_component_inplace(data=init_params, key=key)
339 return default_from_dict(cls, data=data)