Coverage for haystack/hooks/tool_result_offloading/hooks.py: 97%
76 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import json
6from typing import Any
8from haystack.components.agents.state.state import State
9from haystack.components.agents.state.state_utils import replace_values
10from haystack.core.serialization import default_from_dict, default_to_dict
11from haystack.dataclasses import ChatMessage
12from haystack.hooks.tool_result_offloading.types import OffloadPolicy, ToolResultStore
13from haystack.hooks.tool_result_offloading.utils import (
14 _OFFLOADED_META_KEY,
15 _content_block_payload,
16 _offloadable_content_blocks,
17 _offloaded_message,
18)
19from haystack.utils.deserialization import deserialize_component_inplace
21# Key under which a per-run store override may be supplied via the Agent's `hook_context` (e.g. a request-scoped
22# sandbox filesystem).
23RESULT_STORE_CONTEXT_KEY = "tool_result_store"
26def _fresh_tool_results_start(messages: list[ChatMessage]) -> int:
27 """
28 Return the index at which the trailing run of tool-result messages begins.
30 The Agent appends the current step's tool results to the end of the conversation, so the trailing contiguous
31 block of tool-result messages is exactly the freshly produced batch; everything before it is history the hook
32 must not touch (results from earlier steps or ones the caller passed in).
34 :param messages: The conversation, oldest to newest.
35 :returns: The index of the first message in the trailing tool-result block, or `len(messages)` when the last
36 message is not a tool result (no fresh results to offload).
37 """
38 index = len(messages)
39 while index > 0 and messages[index - 1].tool_call_result is not None:
40 index -= 1
41 return index
44def _serialize_offload_strategies(strategies: dict[str | tuple[str, ...], OffloadPolicy]) -> dict[str, Any]:
45 """
46 Serialize an offload-strategies mapping to a plain, mapping-key-safe dictionary.
48 Mapping keys must be strings, so a tuple of tool names (one policy shared across several tools) is encoded as a
49 JSON-array string (e.g. `("a", "b")` -> `'["a", "b"]'`); a single tool name or the `"*"` wildcard is kept as-is.
50 Each policy is serialized via its own `to_dict`, which embeds its type so it can be reconstructed regardless of
51 its concrete class.
53 :param strategies: Mapping of tool name (or a tuple of tool names, or `"*"`) to its `OffloadPolicy`.
54 :returns: The same mapping with string keys and each policy serialized to a dictionary.
55 """
56 return {
57 (json.dumps(list(key)) if isinstance(key, tuple) else key): policy.to_dict()
58 for key, policy in strategies.items()
59 }
62def _deserialize_offload_strategies(data: dict[str, Any]) -> dict[str | tuple[str, ...], OffloadPolicy]:
63 """
64 Deserialize an offload-strategies mapping from its serialized form.
66 Reverses `_serialize_offload_strategies`: each policy is rebuilt from its stored type via
67 `deserialize_component_inplace`, and keys that were encoded as JSON-array strings become tuples of tool names
68 (single tool-name and `"*"` keys are kept as-is).
70 :param data: Raw dictionary of serialized offload strategies, keyed by tool name(s).
71 :returns: The offload strategies with their original key and policy types restored.
72 """
73 for raw_key in list(data):
74 deserialize_component_inplace(data, key=raw_key)
75 return {
76 (tuple(json.loads(raw_key)) if isinstance(raw_key, str) and raw_key.startswith("[") else raw_key): policy
77 for raw_key, policy in data.items()
78 }
81class ToolResultOffloadHook:
82 """
83 Offload tool results to a `ToolResultStore`, replacing them in the conversation with a compact pointer.
85 This `after_tool` Agent hook writes the full result to the store so the next LLM call sees a reference instead of
86 the full result. Register it on an `Agent` under the `after_tool` hook point. Which tools offload, and under what
87 condition, is controlled per tool by `offload_strategies`:
89 <!-- test-concept -->
90 ```python
91 from haystack.components.agents import Agent
92 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
93 from haystack.hooks.tool_result_offloading import (
94 AlwaysOffload,
95 FileSystemToolResultStore,
96 NeverOffload,
97 OffloadOverChars,
98 ToolResultOffloadHook,
99 )
101 hook = ToolResultOffloadHook(
102 store=FileSystemToolResultStore(root="tool_results"),
103 offload_strategies={
104 "web_search": AlwaysOffload(), # force offload
105 "get_time": NeverOffload(), # opt out
106 ("read_file", "list_dir"): OffloadOverChars(4000), # tuple key: shared policy
107 "*": OffloadOverChars(8000), # wildcard default for any unlisted tool
108 },
109 )
110 agent = Agent(
111 chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.6-luna"),
112 tools=[web_search, get_time, read_file, list_dir],
113 hooks={"after_tool": [hook]},
114 )
115 ```
117 A key may be a single tool name, a tuple of tool names sharing one policy, or the wildcard `"*"` which applies to
118 any tool without a more specific entry. More specific keys win. A tool with no matching key (and no `"*"`) is not
119 offloaded.
121 Only successful tool output is offloaded; error results are always left in context. Each part of a result is
122 written to its own store entry and the pointer says where each one went. Image and file content is only offloaded
123 to a store that sets `supports_binary_content`; with a text-only store the result stays in context and a warning
124 is logged. Each result is offloaded at most once, even though the hook runs on every tool step.
126 The hook keeps no mutable state, so a single instance can be shared across concurrent runs. The constructor
127 `store`, however, is shared by every run that does not override it — fine for single-user or local use, but in a
128 multi-user server give each run its own isolated store (a per-session directory or sandbox) via `hook_context`
129 under the key `RESULT_STORE_CONTEXT_KEY`
130 (`agent.run(messages=[...], hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store})`); it overrides the
131 constructor store for that run. Isolating the store per run keeps concurrent users from colliding on store keys or
132 reading each other's offloaded results — important especially when a bash/read tool is scoped to the store.
133 """
135 allowed_hook_points = ("after_tool",)
137 def __init__(
138 self,
139 store: ToolResultStore,
140 offload_strategies: dict[str | tuple[str, ...], OffloadPolicy],
141 *,
142 preview_chars: int = 200,
143 ) -> None:
144 """
145 Initialize the hook with a store and per-tool offload strategies.
147 :param store: Where offloaded results are written. Can be overridden per run via `hook_context`.
148 :param offload_strategies: Mapping of tool name (or a tuple of tool names, or the wildcard `"*"`) to the
149 `OffloadPolicy` that decides whether that tool's results are offloaded.
150 :param preview_chars: Number of leading characters of each offloaded text to include in the pointer left in
151 the conversation, so the model knows roughly what was offloaded. Image and file blocks are described by
152 their MIME type and size instead.
153 """
154 self.store = store
155 self.offload_strategies = offload_strategies
156 self.preview_chars = preview_chars
158 def run(self, state: State) -> None:
159 """
160 Offload the freshly produced tool results in `state.data["messages"]` according to `offload_strategies`.
162 Considers only the trailing block of tool-result messages (the current step's results); earlier history is
163 left untouched. Offloads each of those messages its policy opts in for, and writes the rewritten conversation
164 back to `messages` only if at least one message changed.
166 Results are written to the store this run resolves to: a per-run store passed in `state`'s `hook_context`
167 under `RESULT_STORE_CONTEXT_KEY` if present, otherwise the store the hook was constructed with. Supply the
168 per-run store when calling the Agent, e.g.
169 `agent.run(messages=[...], hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store})`. In a multi-user
170 server, pass an isolated store per run this way so concurrent users write to separate locations and never
171 read each other's results.
173 The hook keeps no mutable state, so a single instance is safe to share across concurrent runs; isolation
174 comes entirely from giving each run its own store via `hook_context`.
176 :param state: The Agent's live `State`. Reads the per-run store from `hook_context` and rewrites the offloaded
177 tool-result messages back into `messages`.
178 :returns: None. The hook mutates `state` in place.
179 """
180 messages = state.data.get("messages") or []
181 start = _fresh_tool_results_start(messages=messages)
182 if start == len(messages):
183 return
185 # The hook instance is shared across concurrent runs, so a run isolates itself by carrying its own store in
186 # `hook_context`.
187 hook_context = state.data.get("hook_context") or {}
188 store = hook_context.get(RESULT_STORE_CONTEXT_KEY, self.store)
190 rewritten: list[ChatMessage] = list(messages[:start])
191 changed = False
192 for index, message in enumerate(messages[start:]):
193 new_message = self._maybe_offload(message=message, store=store, state=state, index=index)
194 rewritten.append(new_message)
195 changed = changed or new_message is not message
197 # Only write back to state when at least one message changed
198 if changed:
199 state.set(key="messages", value=rewritten, handler_override=replace_values)
201 def _policy_for(self, tool_name: str) -> OffloadPolicy | None:
202 """
203 Resolve the offload policy that applies to a tool, most specific first.
205 Lookup order: an exact tool-name key, then any tuple key that contains the tool name, then the `"*"` wildcard.
207 :param tool_name: The name of the tool whose policy to resolve.
208 :returns: The matching `OffloadPolicy`, or None when no key (and no `"*"`) applies.
209 """
210 strategies = self.offload_strategies
211 if tool_name in strategies:
212 return strategies[tool_name]
213 for key, policy in strategies.items():
214 if isinstance(key, tuple) and tool_name in key:
215 return policy
216 return strategies.get("*")
218 def _maybe_offload(self, message: ChatMessage, store: ToolResultStore, state: State, index: int) -> ChatMessage:
219 """
220 Offload a single tool-result message if its policy opts in, otherwise return it unchanged.
222 A message is left as-is when it is not a tool result, when the result is an error (including `before_tool`
223 human-in-the-loop rejections), when it was already offloaded (e.g. another offload hook under `after_tool`
224 handled it), when no policy applies, when the result is empty (no content, or nothing but empty text), when
225 the result carries image or file content that `store` cannot store, or when the policy declines to offload.
227 Otherwise the result is written to `store` and the message is rebuilt with a pointer in place of the full
228 result, preserving its origin and error flag and marking it offloaded. Each part of the result goes to its own
229 store entry.
231 :param message: The message to consider offloading.
232 :param store: The store to write the result to.
233 :param state: The Agent's live `State`, passed to the policy and used to derive the store key.
234 :param index: The message's position within this step's batch of tool results, used to build the store key.
235 :returns: An offloaded copy of the message, or the original message when it is not offloaded.
236 """
237 result = message.tool_call_result
238 # Only successful tool output is offloaded - never errors, before_tool human-in-the-loop rejections, or a
239 # result already offloaded (guards against a second offload hook re-offloading the first one's pointer).
240 if result is None or result.error or message.meta.get(_OFFLOADED_META_KEY):
241 return message
243 tool_name = result.origin.tool_name
244 policy = self._policy_for(tool_name=tool_name)
246 # If no policy applies, leave the result in context
247 if policy is None:
248 return message
250 # An empty result, or image or file content a text-only store cannot take, stays in context.
251 content_blocks = _offloadable_content_blocks(result=result, store=store)
252 if content_blocks is None:
253 return message
255 # The policy sizes up the result by the string it occupies in the conversation, which for an image or a file
256 # block is its base64 payload.
257 payload = "".join(_content_block_payload(content_block=content_block) for content_block in content_blocks)
258 if not policy.should_offload(tool_name=tool_name, result=payload, state=state):
259 return message
261 # Step, tool name, and call id keep results from different tools and steps from colliding. A tool call id is
262 # optional, so an id-less call falls back to its position in this step's batch.
263 step = state.data.get("step_count", 0)
264 return _offloaded_message(
265 message=message,
266 content_blocks=content_blocks,
267 store=store,
268 key_prefix=f"{step}_{tool_name}_{result.origin.id or f'call{index}'}",
269 preview_chars=self.preview_chars,
270 )
272 def to_dict(self) -> dict[str, Any]:
273 """
274 Serialize the hook, including its store and per-tool offload strategies.
276 :returns: A dictionary representation of the hook.
277 """
278 return default_to_dict(
279 self,
280 store=self.store.to_dict(),
281 offload_strategies=_serialize_offload_strategies(strategies=self.offload_strategies),
282 preview_chars=self.preview_chars,
283 )
285 @classmethod
286 def from_dict(cls, data: dict[str, Any]) -> "ToolResultOffloadHook":
287 """
288 Deserialize the hook, reconstructing its store and offload strategies.
290 :param data: A dictionary representation produced by `to_dict`.
291 :returns: The deserialized `ToolResultOffloadHook`.
292 """
293 init_params = data.get("init_parameters", {})
294 if init_params.get("store") is not None:
295 deserialize_component_inplace(init_params, key="store")
296 if init_params.get("offload_strategies") is not None:
297 init_params["offload_strategies"] = _deserialize_offload_strategies(data=init_params["offload_strategies"])
298 return default_from_dict(cls, data)