Coverage for haystack/hooks/compaction/tool_result_offloading.py: 98%
50 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
5import asyncio
6from typing import Any
7from uuid import uuid4
9from haystack.core.serialization import default_from_dict, default_to_dict
10from haystack.dataclasses import ChatMessage
11from haystack.hooks.compaction.types import Compactor
12from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _replace_tool_results_until_target
13from haystack.hooks.tool_result_offloading.types import ToolResultStore
14from haystack.hooks.tool_result_offloading.utils import (
15 _OFFLOADED_META_KEY,
16 _offloadable_content_blocks,
17 _offloaded_message,
18)
19from haystack.token_counters import TokenCounter
20from haystack.utils.deserialization import deserialize_component_inplace
21from haystack.utils.experimental import _experimental
24@_experimental
25class ToolResultOffloadCompactor(Compactor):
26 """
27 Writes older tool results to a `ToolResultStore` and leaves a reference in their place.
29 Every tool call keeps its matching result, and the model can read the full output back with a read tool you scope
30 to the same store. `CompactionHook` only calls the compactor once the conversation reaches `compact_at`, so fresh
31 output stays in context until then. Use `ToolResultOffloadHook` to offload a tool's output as soon as it arrives.
33 <!-- test-ignore -->
34 ```python
35 from typing import Annotated
37 from haystack.components.agents import Agent
38 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
39 from haystack.hooks.compaction import CompactionHook, ToolResultOffloadCompactor
40 from haystack.hooks.tool_result_offloading import FileSystemToolResultStore
41 from haystack.tools import tool
43 store = FileSystemToolResultStore(root="tool_results")
46 @tool
47 def read_offloaded_result(path: Annotated[str, "Path of an offloaded tool result"]) -> str:
48 '''Read back the full content of an offloaded tool result.'''
49 content = store.read(path)
50 if isinstance(content, bytes):
51 return f"'{path}' holds {len(content)} bytes of binary content and cannot be read as text."
52 return content
55 # `web_search` stands in for your own tool that returns long results.
56 hook = CompactionHook(compactor=ToolResultOffloadCompactor(store=store), context_window=400_000)
57 agent = Agent(
58 chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.4-nano"),
59 tools=[web_search, read_offloaded_result],
60 hooks={"before_llm": [hook]},
61 )
62 ```
64 The compactor always writes to the store it was created with; unlike `ToolResultOffloadHook`, it does not read a
65 per-run store from `hook_context`. In a multi-user server, create an Agent per run with its own compactor and
66 store, so users never read each other's results.
68 For results with images or files, we recommend passing a provider token counter such as `OpenAITokenCounter` to
69 `CompactionHook`, which measures their real size. Its default `ApproximateTokenCounter` charges a flat
70 `tokens_per_image` and `tokens_per_file` instead, which can undercount such results and keep them below
71 `min_tokens`.
72 """
74 def __init__(
75 self, store: ToolResultStore, *, min_keep_steps: int = 1, min_tokens: int = 200, preview_chars: int = 200
76 ) -> None:
77 """
78 Initialize the compactor.
80 :param store: Where offloaded results are written. Image and file results are only written to a store that sets
81 `supports_binary_content`; otherwise they stay in the conversation.
82 :param min_keep_steps: Number of most recent tool-calling Agent steps whose results are never offloaded, even
83 when the target is missed. Must be at least 1, so the model always sees the latest results.
84 :param min_tokens: Only offload tool-result messages larger than this many tokens.
85 :param preview_chars: Number of leading characters of each offloaded text kept in its reference. Image and file
86 results are described by MIME type and size instead.
87 :raises ValueError: If `min_keep_steps` is less than 1, or `min_tokens` or `preview_chars` is negative.
88 """
89 if min_keep_steps < 1:
90 raise ValueError(
91 f"`min_keep_steps` must be at least 1, got {min_keep_steps}. The most recent tool-calling step "
92 f"contains results the model may still need."
93 )
94 if min_tokens < 0:
95 raise ValueError(f"`min_tokens` must be at least 0, got {min_tokens}.")
96 if preview_chars < 0:
97 raise ValueError(f"`preview_chars` must be at least 0, got {preview_chars}.")
98 self.store = store
99 self.min_keep_steps = min_keep_steps
100 self.min_tokens = min_tokens
101 self.preview_chars = preview_chars
103 def compact(
104 self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter
105 ) -> list[ChatMessage] | None:
106 """
107 Offload tool results to the store, oldest first, until the conversation fits within `target_tokens`.
109 Results from the most recent `min_keep_steps` tool-calling Agent steps are never offloaded.
111 :param messages: The conversation to compact, oldest to newest.
112 :param target_tokens: The token count the compacted conversation should fit within.
113 :param token_counter: The `TokenCounter` used to measure the conversation before and after each replacement.
114 :returns: The conversation with older tool results replaced by references, or None when nothing was offloaded.
115 """
116 return _replace_tool_results_until_target(
117 messages=messages,
118 target_tokens=target_tokens,
119 token_counter=token_counter,
120 min_keep_steps=self.min_keep_steps,
121 replace=lambda message, _index: self._offload(message=message, token_counter=token_counter),
122 )
124 async def compact_async(
125 self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter
126 ) -> list[ChatMessage] | None:
127 """
128 Run `compact` in a thread so store writes and token counting do not block the event loop.
130 :param messages: The conversation to compact, oldest to newest.
131 :param target_tokens: The token count the compacted conversation should fit within.
132 :param token_counter: The `TokenCounter` used to measure the conversation before and after each replacement.
133 :returns: The conversation with older tool results replaced by references, or None when nothing was offloaded.
134 """
135 return await asyncio.to_thread(
136 self.compact, messages=messages, target_tokens=target_tokens, token_counter=token_counter
137 )
139 def _offload(self, message: ChatMessage, token_counter: TokenCounter) -> tuple[ChatMessage, int] | None:
140 """
141 Write one tool result to the store and report how many tokens its reference saves.
143 :param message: The tool-result message to consider.
144 :param token_counter: The `TokenCounter` that measures the result and its reference.
145 :returns: The offloaded message and the number of tokens it saves, or None when the result stays in context.
146 """
147 result = message.tool_call_result
148 # Skip errors, results another compactor already rewrote, and results that are already a reference.
149 if (
150 result is None
151 or result.error
152 or message.meta.get(_OFFLOADED_META_KEY)
153 or _COMPACTION_META_KEY in message.meta
154 ):
155 return None
157 original_tokens = token_counter.count(messages=[message])
158 if original_tokens <= self.min_tokens:
159 return None
161 # An empty result, or image or file content a text-only store cannot take, stays in context.
162 content_blocks = _offloadable_content_blocks(result=result, store=self.store)
163 if content_blocks is None:
164 return None
166 # A random suffix keeps store keys unique across runs that share the store, with or without a tool call id.
167 offloaded = _offloaded_message(
168 message=message,
169 content_blocks=content_blocks,
170 store=self.store,
171 key_prefix=f"compacted_{result.origin.tool_name}_{uuid4().hex}",
172 preview_chars=self.preview_chars,
173 additional_meta={
174 _COMPACTION_META_KEY: {"strategy": "tool_result_offloading", "original_tokens": original_tokens}
175 },
176 )
177 saved_tokens = original_tokens - token_counter.count(messages=[offloaded])
178 # A reference can cost more than a small result. That replacement is dropped and its store entry left unused.
179 return (offloaded, saved_tokens) if saved_tokens > 0 else None
181 def to_dict(self) -> dict[str, Any]:
182 """
183 Serialize the compactor, including its store.
185 :returns: A dictionary representation of the compactor.
186 """
187 return default_to_dict(
188 self,
189 store=self.store.to_dict(),
190 min_keep_steps=self.min_keep_steps,
191 min_tokens=self.min_tokens,
192 preview_chars=self.preview_chars,
193 )
195 @classmethod
196 def from_dict(cls, data: dict[str, Any]) -> "ToolResultOffloadCompactor":
197 """
198 Deserialize the compactor, reconstructing its store.
200 :param data: A dictionary representation produced by `to_dict`.
201 :returns: The deserialized `ToolResultOffloadCompactor`.
202 """
203 init_params = data.get("init_parameters", {})
204 if init_params.get("store") is not None:
205 deserialize_component_inplace(init_params, key="store")
206 return default_from_dict(cls, data)