Coverage for haystack/hooks/tool_result_offloading/hooks.py: 99%
95 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
5import json
6from typing import Any
8from haystack import logging
9from haystack.components.agents.state.state import State
10from haystack.components.agents.state.state_utils import replace_values
11from haystack.core.serialization import default_from_dict, default_to_dict
12from haystack.dataclasses import ChatMessage, TextContent
13from haystack.dataclasses.chat_message import ToolCallResultContentT
14from haystack.hooks.tool_result_offloading.types import OffloadPolicy, ToolResultStore
15from haystack.utils.deserialization import deserialize_component_inplace
17logger = logging.getLogger(__name__)
19# Meta key marking an already-offloaded tool-result message (its value is the store reference). The offloaded pointer
20# is itself a tool result in the trailing block the hook scans, so this marker stops a second offload hook registered
21# under `after_tool` from offloading the pointer text again and writing a junk file.
22_OFFLOADED_META_KEY = "tool_result_offloaded"
24# Key under which a per-run store override may be supplied via the Agent's `hook_context` (e.g. a request-scoped
25# sandbox filesystem).
26RESULT_STORE_CONTEXT_KEY = "tool_result_store"
29def _result_store_key(tool_name: str, tool_call_id: str | None, step: int, index: int) -> str:
30 """
31 Build a per-result store key that is stable and unique within a run.
33 Combining the step, tool name, and tool call id keeps results from different tools and different steps from
34 colliding. When the tool call carries no id (it is optional and not every generator sets it), the result's
35 position in the step's batch is used instead, so two id-less calls to the same tool in the same step do not
36 collide.
38 :param tool_name: The name of the tool that produced the result.
39 :param tool_call_id: The id of the originating tool call, or None when the call carried no id.
40 :param step: The Agent's current step count.
41 :param index: The result's position within this step's batch of tool results, used when `tool_call_id` is None.
42 :returns: A file-name-like key for the store, e.g. `2_web_search_call-123.txt`.
43 """
44 return f"{step}_{tool_name}_{tool_call_id or f'call{index}'}.txt"
47def _fresh_tool_results_start(messages: list[ChatMessage]) -> int:
48 """
49 Return the index at which the trailing run of tool-result messages begins.
51 The Agent appends the current step's tool results to the end of the conversation, so the trailing contiguous
52 block of tool-result messages is exactly the freshly produced batch; everything before it is history the hook
53 must not touch (results from earlier steps or ones the caller passed in).
55 :param messages: The conversation, oldest to newest.
56 :returns: The index of the first message in the trailing tool-result block, or `len(messages)` when the last
57 message is not a tool result (no fresh results to offload).
58 """
59 index = len(messages)
60 while index > 0 and messages[index - 1].tool_call_result is not None:
61 index -= 1
62 return index
65def _offloadable_text(content: ToolCallResultContentT) -> str | None:
66 """
67 Return the text of a tool result if it can be offloaded as text, otherwise None.
69 A plain string is returned as-is; a non-empty sequence made up entirely of `TextContent` blocks is concatenated
70 into a single string. Anything else (e.g. a result containing image or file content) returns None and is left in
71 context.
73 :param content: The tool result content to inspect.
74 :returns: The offloadable text, or None when the content is not purely text.
75 """
76 if isinstance(content, str):
77 return content
78 texts = [block.text for block in content if isinstance(block, TextContent)]
79 if texts and len(texts) == len(content):
80 return "".join(texts)
81 return None
84def _serialize_offload_strategies(strategies: dict[str | tuple[str, ...], OffloadPolicy]) -> dict[str, Any]:
85 """
86 Serialize an offload-strategies mapping to a plain, mapping-key-safe dictionary.
88 Mapping keys must be strings, so a tuple of tool names (one policy shared across several tools) is encoded as a
89 JSON-array string (e.g. `("a", "b")` -> `'["a", "b"]'`); a single tool name or the `"*"` wildcard is kept as-is.
90 Each policy is serialized via its own `to_dict`, which embeds its type so it can be reconstructed regardless of
91 its concrete class.
93 :param strategies: Mapping of tool name (or a tuple of tool names, or `"*"`) to its `OffloadPolicy`.
94 :returns: The same mapping with string keys and each policy serialized to a dictionary.
95 """
96 return {
97 (json.dumps(list(key)) if isinstance(key, tuple) else key): policy.to_dict()
98 for key, policy in strategies.items()
99 }
102def _deserialize_offload_strategies(data: dict[str, Any]) -> dict[str | tuple[str, ...], OffloadPolicy]:
103 """
104 Deserialize an offload-strategies mapping from its serialized form.
106 Reverses `_serialize_offload_strategies`: each policy is rebuilt from its stored type via
107 `deserialize_component_inplace`, and keys that were encoded as JSON-array strings become tuples of tool names
108 (single tool-name and `"*"` keys are kept as-is).
110 :param data: Raw dictionary of serialized offload strategies, keyed by tool name(s).
111 :returns: The offload strategies with their original key and policy types restored.
112 """
113 for raw_key in list(data):
114 deserialize_component_inplace(data, key=raw_key)
115 return {
116 (tuple(json.loads(raw_key)) if isinstance(raw_key, str) and raw_key.startswith("[") else raw_key): policy
117 for raw_key, policy in data.items()
118 }
121class ToolResultOffloadHook:
122 """
123 Offload tool results to a `ToolResultStore`, replacing them in the conversation with a compact pointer.
125 This `after_tool` Agent hook writes the full result to the store so the next LLM call sees a reference instead of
126 the full result. Register it on an `Agent` under the `after_tool` hook point. Which tools offload, and under what
127 condition, is controlled per tool by `offload_strategies`:
129 <!-- test-concept -->
130 ```python
131 from haystack.components.agents import Agent
132 from haystack.components.generators.chat import OpenAIChatGenerator
133 from haystack.hooks.tool_result_offloading import (
134 AlwaysOffload,
135 FileSystemToolResultStore,
136 NeverOffload,
137 OffloadOverChars,
138 ToolResultOffloadHook,
139 )
141 hook = ToolResultOffloadHook(
142 store=FileSystemToolResultStore(root="tool_results"),
143 offload_strategies={
144 "web_search": AlwaysOffload(), # force offload
145 "get_time": NeverOffload(), # opt out
146 ("read_file", "list_dir"): OffloadOverChars(4000), # tuple key: shared policy
147 "*": OffloadOverChars(8000), # wildcard default for any unlisted tool
148 },
149 )
150 agent = Agent(
151 chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"),
152 tools=[web_search, get_time, read_file, list_dir],
153 hooks={"after_tool": [hook]},
154 )
155 ```
157 A key may be a single tool name, a tuple of tool names sharing one policy, or the wildcard `"*"` which applies to
158 any tool without a more specific entry. More specific keys win. A tool with no matching key (and no `"*"`) is not
159 offloaded.
161 Only successful, text tool output is offloaded. Error results (including `before_tool` human-in-the-loop
162 rejections) are always left in context. Non-text results (image or file content) are also left in context, and a
163 warning is logged when such a result has a matching offload policy; supporting only text is a deliberate choice
164 for now. Each result is offloaded at most once, even though the hook runs on every tool step.
166 The hook keeps no mutable state, so a single instance can be shared across concurrent runs. The constructor
167 `store`, however, is shared by every run that does not override it — fine for single-user or local use, but in a
168 multi-user server give each run its own isolated store (a per-session directory or sandbox) via `hook_context`
169 under the key `RESULT_STORE_CONTEXT_KEY`
170 (`agent.run(messages=[...], hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store})`); it overrides the
171 constructor store for that run. Isolating the store per run keeps concurrent users from colliding on store keys or
172 reading each other's offloaded results — important especially when a bash/read tool is scoped to the store.
173 """
175 allowed_hook_points = ("after_tool",)
177 def __init__(
178 self,
179 store: ToolResultStore,
180 offload_strategies: dict[str | tuple[str, ...], OffloadPolicy],
181 *,
182 preview_chars: int = 200,
183 ) -> None:
184 """
185 Initialize the hook with a store and per-tool offload strategies.
187 :param store: Where offloaded results are written. Can be overridden per run via `hook_context`.
188 :param offload_strategies: Mapping of tool name (or a tuple of tool names, or the wildcard `"*"`) to the
189 `OffloadPolicy` that decides whether that tool's results are offloaded.
190 :param preview_chars: Number of leading characters of the original result to include in the pointer left in
191 the conversation, so the model knows roughly what was offloaded.
192 """
193 self.store = store
194 self.offload_strategies = offload_strategies
195 self.preview_chars = preview_chars
197 def run(self, state: State) -> None:
198 """
199 Offload the freshly produced tool results in `state.data["messages"]` according to `offload_strategies`.
201 Considers only the trailing block of tool-result messages (the current step's results); earlier history is
202 left untouched. Offloads each of those messages its policy opts in for, and writes the rewritten conversation
203 back to `messages` only if at least one message changed.
205 Results are written to the store this run resolves to: a per-run store passed in `state`'s `hook_context`
206 under `RESULT_STORE_CONTEXT_KEY` if present, otherwise the store the hook was constructed with. Supply the
207 per-run store when calling the Agent, e.g.
208 `agent.run(messages=[...], hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store})`. In a multi-user
209 server, pass an isolated store per run this way so concurrent users write to separate locations and never
210 read each other's results.
212 The hook keeps no mutable state, so a single instance is safe to share across concurrent runs; isolation
213 comes entirely from giving each run its own store via `hook_context`.
215 :param state: The Agent's live `State`. Reads the per-run store from `hook_context` and rewrites the offloaded
216 tool-result messages back into `messages`.
217 :returns: None. The hook mutates `state` in place.
218 """
219 messages = state.data.get("messages") or []
220 start = _fresh_tool_results_start(messages)
221 if start == len(messages):
222 return
223 store = self._resolve_store(state)
224 rewritten: list[ChatMessage] = list(messages[:start])
225 changed = False
226 for index, message in enumerate(messages[start:]):
227 new_message = self._maybe_offload(message, store, state, index)
228 rewritten.append(new_message)
229 changed = changed or new_message is not message
230 if changed:
231 state.set("messages", rewritten, handler_override=replace_values)
233 def _resolve_store(self, state: State) -> ToolResultStore:
234 """
235 Return the store to write to for this run.
237 :param state: The Agent's live `State`, whose `hook_context` may carry a per-run store override under
238 `RESULT_STORE_CONTEXT_KEY`.
239 :returns: The per-run store from `hook_context` if provided, otherwise the store the hook was built with.
240 """
241 context = state.data.get("hook_context") or {}
242 return context.get(RESULT_STORE_CONTEXT_KEY, self.store)
244 def _policy_for(self, tool_name: str) -> OffloadPolicy | None:
245 """
246 Resolve the offload policy that applies to a tool, most specific first.
248 Lookup order: an exact tool-name key, then any tuple key that contains the tool name, then the `"*"` wildcard.
250 :param tool_name: The name of the tool whose policy to resolve.
251 :returns: The matching `OffloadPolicy`, or None when no key (and no `"*"`) applies.
252 """
253 strategies = self.offload_strategies
254 if tool_name in strategies:
255 return strategies[tool_name]
256 for key, policy in strategies.items():
257 if isinstance(key, tuple) and tool_name in key:
258 return policy
259 return strategies.get("*")
261 def _maybe_offload(self, message: ChatMessage, store: ToolResultStore, state: State, index: int) -> ChatMessage:
262 """
263 Offload a single tool-result message if its policy opts in, otherwise return it unchanged.
265 A message is left as-is when it is not a tool result, when the result is an error (including `before_tool`
266 human-in-the-loop rejections), when it was already offloaded (e.g. another offload hook under `after_tool`
267 handled it), when no policy applies, when the result is non-text (contains image or file content), or when the
268 policy declines to offload.
270 Otherwise the result text is written to `store` and the message is rebuilt with a pointer in place of the full
271 result, preserving its origin and error flag and marking it offloaded.
273 :param message: The message to consider offloading.
274 :param store: The store to write the result to.
275 :param state: The Agent's live `State`, passed to the policy and used to derive the store key.
276 :param index: The message's position within this step's batch of tool results, used to build the store key.
277 :returns: An offloaded copy of the message, or the original message when it is not offloaded.
278 """
279 result = message.tool_call_result
280 # Only successful tool output is offloaded - never errors, before_tool human-in-the-loop rejections, or a
281 # result already offloaded (guards against a second offload hook re-offloading the first one's pointer).
282 if result is None or result.error or message.meta.get(_OFFLOADED_META_KEY):
283 return message
285 tool_name = result.origin.tool_name
286 policy = self._policy_for(tool_name)
288 # If no policy applies, leave the result in context
289 if policy is None:
290 return message
292 # A policy matched, so an offload was wanted. Offloading only supports text results (a string or a sequence
293 # of TextContent) for now, by design; leave image/file content in context and warn since the intent was to
294 # offload it.
295 text = _offloadable_text(result.result)
296 if text is None:
297 logger.warning(
298 "Tool '{tool}' produced a non-text result; leaving it in context. Result offloading currently "
299 "supports text results only.",
300 tool=tool_name,
301 )
302 return message
304 # If the policy declines to offload, leave the result in context
305 if not policy.should_offload(tool_name, text, state):
306 return message
308 key = _result_store_key(tool_name, result.origin.id, state.data.get("step_count", 0), index)
309 reference = store.write(key=key, content=text)
310 return ChatMessage.from_tool(
311 tool_result=self._pointer(reference, text),
312 origin=result.origin,
313 error=result.error,
314 meta={**message.meta, _OFFLOADED_META_KEY: reference},
315 )
317 def _pointer(self, reference: str, result: str) -> str:
318 """
319 Build the compact pointer that replaces a full result in the conversation.
321 :param reference: The store reference the result was written to.
322 :param result: The original result string, used for its length and a leading preview.
323 :returns: A one-line pointer carrying the reference, the result length, and a `preview_chars`-long preview.
324 """
325 ellip = "..." if len(result) > self.preview_chars else ""
326 preview = result[: self.preview_chars]
327 return f"Tool result offloaded to '{reference}' ({len(result)} characters). Preview: {preview}{ellip}"
329 def to_dict(self) -> dict[str, Any]:
330 """
331 Serialize the hook, including its store and per-tool offload strategies.
333 :returns: A dictionary representation of the hook.
334 """
335 return default_to_dict(
336 self,
337 store=self.store.to_dict(),
338 offload_strategies=_serialize_offload_strategies(self.offload_strategies),
339 preview_chars=self.preview_chars,
340 )
342 @classmethod
343 def from_dict(cls, data: dict[str, Any]) -> "ToolResultOffloadHook":
344 """
345 Deserialize the hook, reconstructing its store and offload strategies.
347 :param data: A dictionary representation produced by `to_dict`.
348 :returns: The deserialized `ToolResultOffloadHook`.
349 """
350 init_params = data.get("init_parameters", {})
351 if init_params.get("store") is not None:
352 deserialize_component_inplace(init_params, key="store")
353 if init_params.get("offload_strategies") is not None:
354 init_params["offload_strategies"] = _deserialize_offload_strategies(init_params["offload_strategies"])
355 return default_from_dict(cls, data)