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