From 48d98650f02202de462491de3b161afeb3678f0d Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee <10526848+sjrl@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:00:11 +0200 Subject: [PATCH] feat: Add tool result offloading (#11849) --- .../hooks/tool_result_offloading/__init__.py | 27 ++ .../hooks/tool_result_offloading/hooks.py | 354 ++++++++++++++++++ .../hooks/tool_result_offloading/policies.py | 70 ++++ .../hooks/tool_result_offloading/stores.py | 78 ++++ .../tool_result_offloading/types/__init__.py | 7 + .../tool_result_offloading/types/protocol.py | 76 ++++ pydoc/hooks_api.yml | 6 +- ...ol-result-offloading-1f9d7b7751f20d9d.yaml | 65 ++++ test/hooks/tool_result_offloading/__init__.py | 3 + .../tool_result_offloading/test_hooks.py | 271 ++++++++++++++ .../tool_result_offloading/test_policies.py | 23 ++ .../tool_result_offloading/test_stores.py | 48 +++ 12 files changed, 1026 insertions(+), 2 deletions(-) create mode 100644 haystack/hooks/tool_result_offloading/__init__.py create mode 100644 haystack/hooks/tool_result_offloading/hooks.py create mode 100644 haystack/hooks/tool_result_offloading/policies.py create mode 100644 haystack/hooks/tool_result_offloading/stores.py create mode 100644 haystack/hooks/tool_result_offloading/types/__init__.py create mode 100644 haystack/hooks/tool_result_offloading/types/protocol.py create mode 100644 releasenotes/notes/add-tool-result-offloading-1f9d7b7751f20d9d.yaml create mode 100644 test/hooks/tool_result_offloading/__init__.py create mode 100644 test/hooks/tool_result_offloading/test_hooks.py create mode 100644 test/hooks/tool_result_offloading/test_policies.py create mode 100644 test/hooks/tool_result_offloading/test_stores.py diff --git a/haystack/hooks/tool_result_offloading/__init__.py b/haystack/hooks/tool_result_offloading/__init__.py new file mode 100644 index 0000000000..b16d26d91e --- /dev/null +++ b/haystack/hooks/tool_result_offloading/__init__.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import sys +from typing import TYPE_CHECKING + +from lazy_imports import LazyImporter + +_import_structure = { + "hooks": ["ToolResultOffloadHook", "RESULT_STORE_CONTEXT_KEY"], + "policies": ["AlwaysOffload", "NeverOffload", "OffloadOverChars"], + "stores": ["FileSystemToolResultStore"], + "types": ["OffloadPolicy", "ToolResultStore"], +} + +if TYPE_CHECKING: + from .hooks import RESULT_STORE_CONTEXT_KEY as RESULT_STORE_CONTEXT_KEY + from .hooks import ToolResultOffloadHook as ToolResultOffloadHook + from .policies import AlwaysOffload as AlwaysOffload + from .policies import NeverOffload as NeverOffload + from .policies import OffloadOverChars as OffloadOverChars + from .stores import FileSystemToolResultStore as FileSystemToolResultStore + from .types import OffloadPolicy as OffloadPolicy + from .types import ToolResultStore as ToolResultStore +else: + sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure) diff --git a/haystack/hooks/tool_result_offloading/hooks.py b/haystack/hooks/tool_result_offloading/hooks.py new file mode 100644 index 0000000000..0a29dd58f6 --- /dev/null +++ b/haystack/hooks/tool_result_offloading/hooks.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import json +from typing import Any + +from haystack import logging +from haystack.components.agents.state.state import State +from haystack.components.agents.state.state_utils import replace_values +from haystack.core.serialization import default_from_dict, default_to_dict +from haystack.dataclasses import ChatMessage, TextContent +from haystack.dataclasses.chat_message import ToolCallResultContentT +from haystack.hooks.tool_result_offloading.types import OffloadPolicy, ToolResultStore +from haystack.utils.deserialization import deserialize_component_inplace + +logger = logging.getLogger(__name__) + +# Meta key marking an already-offloaded tool-result message (its value is the store reference). The offloaded pointer +# is itself a tool result in the trailing block the hook scans, so this marker stops a second offload hook registered +# under `after_tool` from offloading the pointer text again and writing a junk file. +_OFFLOADED_META_KEY = "tool_result_offloaded" + +# Key under which a per-run store override may be supplied via the Agent's `hook_context` (e.g. a request-scoped +# sandbox filesystem). +RESULT_STORE_CONTEXT_KEY = "tool_result_store" + + +def _result_store_key(tool_name: str, tool_call_id: str | None, step: int, index: int) -> str: + """ + Build a per-result store key that is stable and unique within a run. + + Combining the step, tool name, and tool call id keeps results from different tools and different steps from + colliding. When the tool call carries no id (it is optional and not every generator sets it), the result's + position in the step's batch is used instead, so two id-less calls to the same tool in the same step do not + collide. + + :param tool_name: The name of the tool that produced the result. + :param tool_call_id: The id of the originating tool call, or None when the call carried no id. + :param step: The Agent's current step count. + :param index: The result's position within this step's batch of tool results, used when `tool_call_id` is None. + :returns: A file-name-like key for the store, e.g. `2_web_search_call-123.txt`. + """ + return f"{step}_{tool_name}_{tool_call_id or f'call{index}'}.txt" + + +def _fresh_tool_results_start(messages: list[ChatMessage]) -> int: + """ + Return the index at which the trailing run of tool-result messages begins. + + The Agent appends the current step's tool results to the end of the conversation, so the trailing contiguous + block of tool-result messages is exactly the freshly produced batch; everything before it is history the hook + must not touch (results from earlier steps or ones the caller passed in). + + :param messages: The conversation, oldest to newest. + :returns: The index of the first message in the trailing tool-result block, or `len(messages)` when the last + message is not a tool result (no fresh results to offload). + """ + index = len(messages) + while index > 0 and messages[index - 1].tool_call_result is not None: + index -= 1 + return index + + +def _offloadable_text(content: ToolCallResultContentT) -> str | None: + """ + Return the text of a tool result if it can be offloaded as text, otherwise None. + + A plain string is returned as-is; a non-empty sequence made up entirely of `TextContent` blocks is concatenated + into a single string. Anything else (e.g. a result containing image or file content) returns None and is left in + context. + + :param content: The tool result content to inspect. + :returns: The offloadable text, or None when the content is not purely text. + """ + if isinstance(content, str): + return content + texts = [block.text for block in content if isinstance(block, TextContent)] + if texts and len(texts) == len(content): + return "".join(texts) + return None + + +def _serialize_offload_strategies(strategies: dict[str | tuple[str, ...], OffloadPolicy]) -> dict[str, Any]: + """ + Serialize an offload-strategies mapping to a plain, mapping-key-safe dictionary. + + Mapping keys must be strings, so a tuple of tool names (one policy shared across several tools) is encoded as a + JSON-array string (e.g. `("a", "b")` -> `'["a", "b"]'`); a single tool name or the `"*"` wildcard is kept as-is. + Each policy is serialized via its own `to_dict`, which embeds its type so it can be reconstructed regardless of + its concrete class. + + :param strategies: Mapping of tool name (or a tuple of tool names, or `"*"`) to its `OffloadPolicy`. + :returns: The same mapping with string keys and each policy serialized to a dictionary. + """ + return { + (json.dumps(list(key)) if isinstance(key, tuple) else key): policy.to_dict() + for key, policy in strategies.items() + } + + +def _deserialize_offload_strategies(data: dict[str, Any]) -> dict[str | tuple[str, ...], OffloadPolicy]: + """ + Deserialize an offload-strategies mapping from its serialized form. + + Reverses `_serialize_offload_strategies`: each policy is rebuilt from its stored type via + `deserialize_component_inplace`, and keys that were encoded as JSON-array strings become tuples of tool names + (single tool-name and `"*"` keys are kept as-is). + + :param data: Raw dictionary of serialized offload strategies, keyed by tool name(s). + :returns: The offload strategies with their original key and policy types restored. + """ + for raw_key in list(data): + deserialize_component_inplace(data, key=raw_key) + return { + (tuple(json.loads(raw_key)) if isinstance(raw_key, str) and raw_key.startswith("[") else raw_key): policy + for raw_key, policy in data.items() + } + + +class ToolResultOffloadHook: + """ + Offload tool results to a `ToolResultStore`, replacing them in the conversation with a compact pointer. + + This `after_tool` Agent hook writes the full result to the store so the next LLM call sees a reference instead of + the full result. Register it on an `Agent` under the `after_tool` hook point. Which tools offload, and under what + condition, is controlled per tool by `offload_strategies`: + + ```python + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIChatGenerator + from haystack.hooks.tool_result_offloading import ( + AlwaysOffload, + FileSystemToolResultStore, + NeverOffload, + OffloadOverChars, + ToolResultOffloadHook, + ) + + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root="tool_results"), + offload_strategies={ + "web_search": AlwaysOffload(), # force offload + "get_time": NeverOffload(), # opt out + ("read_file", "list_dir"): OffloadOverChars(4000), # tuple key: shared policy + "*": OffloadOverChars(8000), # wildcard default for any unlisted tool + }, + ) + agent = Agent( + chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"), + tools=[web_search, get_time, read_file, list_dir], + hooks={"after_tool": [hook]}, + ) + ``` + + A key may be a single tool name, a tuple of tool names sharing one policy, or the wildcard `"*"` which applies to + any tool without a more specific entry. More specific keys win. A tool with no matching key (and no `"*"`) is not + offloaded. + + Only successful, text tool output is offloaded. Error results (including `before_tool` human-in-the-loop + rejections) are always left in context. Non-text results (image or file content) are also left in context, and a + warning is logged when such a result has a matching offload policy; supporting only text is a deliberate choice + for now. Each result is offloaded at most once, even though the hook runs on every tool step. + + The hook keeps no mutable state, so a single instance can be shared across concurrent runs. The constructor + `store`, however, is shared by every run that does not override it — fine for single-user or local use, but in a + multi-user server give each run its own isolated store (a per-session directory or sandbox) via `hook_context` + under the key `RESULT_STORE_CONTEXT_KEY` + (`agent.run(messages=[...], hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store})`); it overrides the + constructor store for that run. Isolating the store per run keeps concurrent users from colliding on store keys or + reading each other's offloaded results — important especially when a bash/read tool is scoped to the store. + """ + + allowed_hook_points = ("after_tool",) + + def __init__( + self, + store: ToolResultStore, + offload_strategies: dict[str | tuple[str, ...], OffloadPolicy], + *, + preview_chars: int = 200, + ) -> None: + """ + Initialize the hook with a store and per-tool offload strategies. + + :param store: Where offloaded results are written. Can be overridden per run via `hook_context`. + :param offload_strategies: Mapping of tool name (or a tuple of tool names, or the wildcard `"*"`) to the + `OffloadPolicy` that decides whether that tool's results are offloaded. + :param preview_chars: Number of leading characters of the original result to include in the pointer left in + the conversation, so the model knows roughly what was offloaded. + """ + self.store = store + self.offload_strategies = offload_strategies + self.preview_chars = preview_chars + + def run(self, state: State) -> None: + """ + Offload the freshly produced tool results in `state.data["messages"]` according to `offload_strategies`. + + Considers only the trailing block of tool-result messages (the current step's results); earlier history is + left untouched. Offloads each of those messages its policy opts in for, and writes the rewritten conversation + back to `messages` only if at least one message changed. + + Results are written to the store this run resolves to: a per-run store passed in `state`'s `hook_context` + under `RESULT_STORE_CONTEXT_KEY` if present, otherwise the store the hook was constructed with. Supply the + per-run store when calling the Agent, e.g. + `agent.run(messages=[...], hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store})`. In a multi-user + server, pass an isolated store per run this way so concurrent users write to separate locations and never + read each other's results. + + The hook keeps no mutable state, so a single instance is safe to share across concurrent runs; isolation + comes entirely from giving each run its own store via `hook_context`. + + :param state: The Agent's live `State`. Reads the per-run store from `hook_context` and rewrites the offloaded + tool-result messages back into `messages`. + :returns: None. The hook mutates `state` in place. + """ + messages = state.data.get("messages") or [] + start = _fresh_tool_results_start(messages) + if start == len(messages): + return + store = self._resolve_store(state) + rewritten: list[ChatMessage] = list(messages[:start]) + changed = False + for index, message in enumerate(messages[start:]): + new_message = self._maybe_offload(message, store, state, index) + rewritten.append(new_message) + changed = changed or new_message is not message + if changed: + state.set("messages", rewritten, handler_override=replace_values) + + def _resolve_store(self, state: State) -> ToolResultStore: + """ + Return the store to write to for this run. + + :param state: The Agent's live `State`, whose `hook_context` may carry a per-run store override under + `RESULT_STORE_CONTEXT_KEY`. + :returns: The per-run store from `hook_context` if provided, otherwise the store the hook was built with. + """ + context = state.data.get("hook_context") or {} + return context.get(RESULT_STORE_CONTEXT_KEY, self.store) + + def _policy_for(self, tool_name: str) -> OffloadPolicy | None: + """ + Resolve the offload policy that applies to a tool, most specific first. + + Lookup order: an exact tool-name key, then any tuple key that contains the tool name, then the `"*"` wildcard. + + :param tool_name: The name of the tool whose policy to resolve. + :returns: The matching `OffloadPolicy`, or None when no key (and no `"*"`) applies. + """ + strategies = self.offload_strategies + if tool_name in strategies: + return strategies[tool_name] + for key, policy in strategies.items(): + if isinstance(key, tuple) and tool_name in key: + return policy + return strategies.get("*") + + def _maybe_offload(self, message: ChatMessage, store: ToolResultStore, state: State, index: int) -> ChatMessage: + """ + Offload a single tool-result message if its policy opts in, otherwise return it unchanged. + + A message is left as-is when it is not a tool result, when the result is an error (including `before_tool` + human-in-the-loop rejections), when it was already offloaded (e.g. another offload hook under `after_tool` + handled it), when no policy applies, when the result is non-text (contains image or file content), or when the + policy declines to offload. + + Otherwise the result text is written to `store` and the message is rebuilt with a pointer in place of the full + result, preserving its origin and error flag and marking it offloaded. + + :param message: The message to consider offloading. + :param store: The store to write the result to. + :param state: The Agent's live `State`, passed to the policy and used to derive the store key. + :param index: The message's position within this step's batch of tool results, used to build the store key. + :returns: An offloaded copy of the message, or the original message when it is not offloaded. + """ + result = message.tool_call_result + # Only successful tool output is offloaded - never errors, before_tool human-in-the-loop rejections, or a + # result already offloaded (guards against a second offload hook re-offloading the first one's pointer). + if result is None or result.error or message.meta.get(_OFFLOADED_META_KEY): + return message + + tool_name = result.origin.tool_name + policy = self._policy_for(tool_name) + + # If no policy applies, leave the result in context + if policy is None: + return message + + # A policy matched, so an offload was wanted. Offloading only supports text results (a string or a sequence + # of TextContent) for now, by design; leave image/file content in context and warn since the intent was to + # offload it. + text = _offloadable_text(result.result) + if text is None: + logger.warning( + "Tool '{tool}' produced a non-text result; leaving it in context. Result offloading currently " + "supports text results only.", + tool=tool_name, + ) + return message + + # If the policy declines to offload, leave the result in context + if not policy.should_offload(tool_name, text, state): + return message + + key = _result_store_key(tool_name, result.origin.id, state.data.get("step_count", 0), index) + reference = store.write(key=key, content=text) + return ChatMessage.from_tool( + tool_result=self._pointer(reference, text), + origin=result.origin, + error=result.error, + meta={**message.meta, _OFFLOADED_META_KEY: reference}, + ) + + def _pointer(self, reference: str, result: str) -> str: + """ + Build the compact pointer that replaces a full result in the conversation. + + :param reference: The store reference the result was written to. + :param result: The original result string, used for its length and a leading preview. + :returns: A one-line pointer carrying the reference, the result length, and a `preview_chars`-long preview. + """ + ellip = "..." if len(result) > self.preview_chars else "" + preview = result[: self.preview_chars] + return f"Tool result offloaded to '{reference}' ({len(result)} characters). Preview: {preview}{ellip}" + + def to_dict(self) -> dict[str, Any]: + """ + Serialize the hook, including its store and per-tool offload strategies. + + :returns: A dictionary representation of the hook. + """ + return default_to_dict( + self, + store=self.store.to_dict(), + offload_strategies=_serialize_offload_strategies(self.offload_strategies), + preview_chars=self.preview_chars, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ToolResultOffloadHook": + """ + Deserialize the hook, reconstructing its store and offload strategies. + + :param data: A dictionary representation produced by `to_dict`. + :returns: The deserialized `ToolResultOffloadHook`. + """ + init_params = data.get("init_parameters", {}) + if init_params.get("store") is not None: + deserialize_component_inplace(init_params, key="store") + if init_params.get("offload_strategies") is not None: + init_params["offload_strategies"] = _deserialize_offload_strategies(init_params["offload_strategies"]) + return default_from_dict(cls, data) diff --git a/haystack/hooks/tool_result_offloading/policies.py b/haystack/hooks/tool_result_offloading/policies.py new file mode 100644 index 0000000000..c718c7983f --- /dev/null +++ b/haystack/hooks/tool_result_offloading/policies.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.components.agents.state.state import State +from haystack.core.serialization import default_to_dict +from haystack.hooks.tool_result_offloading.types import OffloadPolicy + + +class AlwaysOffload(OffloadPolicy): + """Offload every result of the tool it is assigned to.""" + + def should_offload(self, tool_name: str, result: str, state: State) -> bool: # noqa: ARG002 + """ + Decide whether to offload the given tool result. + + :param tool_name: The name of the tool that produced the result (unused; this policy always offloads). + :param result: The tool result string (unused; this policy always offloads). + :param state: The Agent's live `State` (unused; this policy always offloads). + :returns: Always True. + """ + return True + + +class NeverOffload(OffloadPolicy): + """Never offload; keep the tool's full result in context. Use to opt a tool out of a wildcard default.""" + + def should_offload(self, tool_name: str, result: str, state: State) -> bool: # noqa: ARG002 + """ + Decide whether to offload the given tool result. + + :param tool_name: The name of the tool that produced the result (unused; this policy never offloads). + :param result: The tool result string (unused; this policy never offloads). + :param state: The Agent's live `State` (unused; this policy never offloads). + :returns: Always False. + """ + return False + + +class OffloadOverChars(OffloadPolicy): + """Offload a result only when its string length exceeds `threshold` characters.""" + + def __init__(self, threshold: int) -> None: + """ + Initialize the policy with its character threshold. + + :param threshold: Offload the result when its length in characters is strictly greater than this value. + """ + self.threshold = threshold + + def should_offload(self, tool_name: str, result: str, state: State) -> bool: # noqa: ARG002 + """ + Decide whether to offload the given tool result based on its length. + + :param tool_name: The name of the tool that produced the result (unused; only length is considered). + :param result: The tool result string whose length is compared against the threshold. + :param state: The Agent's live `State` (unused; only length is considered). + :returns: True when `result` is longer than `threshold` characters, otherwise False. + """ + return len(result) > self.threshold + + def to_dict(self) -> dict[str, Any]: + """ + Serialize the policy, including its threshold. + + :returns: A dictionary representation of the policy. + """ + return default_to_dict(self, threshold=self.threshold) diff --git a/haystack/hooks/tool_result_offloading/stores.py b/haystack/hooks/tool_result_offloading/stores.py new file mode 100644 index 0000000000..eb23c5426d --- /dev/null +++ b/haystack/hooks/tool_result_offloading/stores.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from typing import Any + +from haystack.core.serialization import default_from_dict, default_to_dict +from haystack.hooks.tool_result_offloading.types import ToolResultStore + + +class FileSystemToolResultStore(ToolResultStore): + """ + A `ToolResultStore` that writes offloaded tool results to files under a root directory on the local file system. + + ```python + from haystack.hooks.tool_result_offloading import FileSystemToolResultStore + + store = FileSystemToolResultStore(root="tool_results") + reference = store.write(key="search_1.txt", content="...") + store.read(reference) + ``` + """ + + def __init__(self, root: str | Path) -> None: + """ + Initialize the store with the root directory results are written under. + + :param root: Directory under which result files are written. Created on first write if it does not exist. + """ + self.root = Path(root) + + def write(self, *, key: str, content: str) -> str: + """ + Write `content` to `/`, creating parent directories, and return the file path. + + The resolved target must stay within the root directory: a `key` that escapes it (e.g. containing `../` or an + absolute path) is rejected, so a tool-provided key cannot write outside the store. + + :param key: Relative file name for the result within the store root. + :param content: The tool result to persist. + :returns: The absolute path the content was written to, as a string, for use with `read`. + :raises ValueError: If `key` resolves to a location outside the store root. + """ + root = self.root.resolve() + path = (root / key).resolve() + if not path.is_relative_to(root): + raise ValueError(f"Result key '{key}' resolves outside the store root '{root}'.") + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + return str(path) + + def read(self, reference: str) -> str: + """ + Read back the content previously written to `reference`. + + :param reference: A path returned by `write`. + :returns: The stored content. + """ + return Path(reference).read_text(encoding="utf-8") + + def to_dict(self) -> dict[str, Any]: + """ + Serialize the store, storing its root directory as a string. + + :returns: A dictionary representation of the store. + """ + return default_to_dict(self, root=str(self.root)) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "FileSystemToolResultStore": + """ + Deserialize the store from a dictionary. + + :param data: A dictionary representation produced by `to_dict`. + :returns: The deserialized `FileSystemToolResultStore`. + """ + return default_from_dict(cls, data) diff --git a/haystack/hooks/tool_result_offloading/types/__init__.py b/haystack/hooks/tool_result_offloading/types/__init__.py new file mode 100644 index 0000000000..28c28dfc25 --- /dev/null +++ b/haystack/hooks/tool_result_offloading/types/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from .protocol import OffloadPolicy, ToolResultStore + +__all__ = ["OffloadPolicy", "ToolResultStore"] diff --git a/haystack/hooks/tool_result_offloading/types/protocol.py b/haystack/hooks/tool_result_offloading/types/protocol.py new file mode 100644 index 0000000000..c396f4c9a3 --- /dev/null +++ b/haystack/hooks/tool_result_offloading/types/protocol.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, Protocol + +from haystack.components.agents.state.state import State +from haystack.core.serialization import default_from_dict, default_to_dict + + +class ToolResultStore(Protocol): + """ + A place a `ToolResultOffloadHook` writes offloaded tool results to, and reads them back from. + + Implementations decide where and how the content lives (local disk, an isolated sandbox filesystem, object + storage, ...). `write` returns an opaque reference string that the Agent puts in the conversation in place of the + full result; `read` resolves that reference back to the original content. + + Implement both `to_dict` and `from_dict` to make a custom store serializable; the default implementations below + cover stores whose constructor takes no arguments. + """ + + def write(self, *, key: str, content: str) -> str: + """ + Persist `content` under `key` and return an opaque reference to it. + + :param key: A stable, per-result identifier the hook derives from the tool call (e.g. a file name). + :param content: The tool result to persist. + :returns: A reference string (e.g. a path or URI) that `read` can later resolve. + """ + ... + + def read(self, reference: str) -> str: + """Return the content previously stored under `reference`.""" + ... + + def to_dict(self) -> dict[str, Any]: + """Serialize the store to a dictionary.""" + return default_to_dict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ToolResultStore": + """Deserialize the store from a dictionary.""" + return default_from_dict(cls, data) + + +class OffloadPolicy(Protocol): + """ + Decides, per tool result, whether the `ToolResultOffloadHook` offloads it to the store or leaves it in context. + + A `ToolResultOffloadHook` maps tool names to policies, so different tools can offload under different conditions + (always, never, or a custom rule such as a size threshold). + + Implement both `to_dict` and `from_dict` to make a custom policy serializable; the default implementations below + cover policies whose constructor takes no arguments. + """ + + def should_offload(self, tool_name: str, result: str, state: State) -> bool: + """ + Return whether the given tool result should be offloaded. + + :param tool_name: The name of the tool that produced the result. + :param result: The tool result as a string (the content that would otherwise stay in the conversation). + :param state: The Agent's live `State`, for policies that decide based on run context. + :returns: True to offload the result to the store, False to leave it in context. + """ + ... + + def to_dict(self) -> dict[str, Any]: + """Serialize the policy to a dictionary.""" + return default_to_dict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "OffloadPolicy": + """Deserialize the policy from a dictionary.""" + return default_from_dict(cls, data) diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml index 14f1e979d9..6555c46941 100644 --- a/pydoc/hooks_api.yml +++ b/pydoc/hooks_api.yml @@ -1,6 +1,7 @@ loaders: - search_path: [../haystack/hooks] - modules: ["protocol", "from_function"] + modules: ["protocol", "from_function", "tool_result_offloading/hooks", "tool_result_offloading/policies", + "tool_result_offloading/stores", "tool_result_offloading/types/protocol"] processors: - type: filter documented_only: true @@ -8,5 +9,6 @@ processors: renderer: title: Hooks id: hooks-api - description: Hooks that run at points in the Agent's run loop and influence it by mutating State. + description: Hooks that run at points in the Agent's run loop and influence it by mutating State, including + built-in tool result offloading. filename: hooks_api.md diff --git a/releasenotes/notes/add-tool-result-offloading-1f9d7b7751f20d9d.yaml b/releasenotes/notes/add-tool-result-offloading-1f9d7b7751f20d9d.yaml new file mode 100644 index 0000000000..a3653d9fe8 --- /dev/null +++ b/releasenotes/notes/add-tool-result-offloading-1f9d7b7751f20d9d.yaml @@ -0,0 +1,65 @@ +--- +features: + - | + Added tool result offloading for the ``Agent`` via the new ``haystack.hooks.tool_result_offloading`` module. + ``ToolResultOffloadHook`` is an ``after_tool`` hook that writes selected tool results to a store and replaces + them in the conversation with a compact pointer (a reference plus a short preview). The next LLM call sees the + reference instead of the full result, which helps manage the context window and is a step towards letting an + Agent operate on offloaded results with follow-up tools (e.g. a bash tool that reads the referenced files). + + Configure what gets offloaded, per tool, via ``offload_strategies``: + + - Map a tool name, a tuple of tool names, or the wildcard ``"*"`` to an ``OffloadPolicy``. More specific keys + win over ``"*"``; a tool with no matching key is not offloaded. + - Built-in policies: ``AlwaysOffload``, ``NeverOffload``, and ``OffloadOverChars`` (offload when the result + exceeds a character threshold). + - For custom conditions, implement the ``OffloadPolicy`` protocol (it provides default ``to_dict``/``from_dict``, + so a stateless policy only needs a ``should_offload`` method). + + Choose where results are stored: + + - ``FileSystemToolResultStore`` writes results to the local file system. + - Implement the ``ToolResultStore`` protocol to target other backends. + - The hook is stateless and can be shared across concurrent runs. In a multi-user server, pass an isolated store + per run via the Agent's ``hook_context`` (``hook_context={RESULT_STORE_CONTEXT_KEY: per_request_store}``) so + concurrent users write to separate locations and never read each other's results. + + What is offloaded: + + - Only successful, text results are offloaded. Error results (including ``before_tool`` human-in-the-loop + rejections) are left in context. + - Non-text results (image or file content) are left in context; supporting only text is a deliberate choice for + now. A warning is logged when a non-text result has a matching offload policy. + - Each result is offloaded at most once. + + .. code-block:: python + + from typing import Annotated + + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIChatGenerator + from haystack.dataclasses import ChatMessage + from haystack.hooks.tool_result_offloading import ( + FileSystemToolResultStore, + OffloadOverChars, + ToolResultOffloadHook, + ) + from haystack.tools import tool + + + @tool + def search(query: Annotated[str, "The search query"]) -> str: + """Search the web and return the (potentially large) results.""" + return f"... large result for {query} ..." + + + offload_hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root="tool_results"), + offload_strategies={"*": OffloadOverChars(4000)}, + ) + agent = Agent( + chat_generator=OpenAIChatGenerator(model="gpt-5.4-nano"), + tools=[search], + hooks={"after_tool": [offload_hook]}, + ) + agent.run(messages=[ChatMessage.from_user("Summarize today's tech news")]) diff --git a/test/hooks/tool_result_offloading/__init__.py b/test/hooks/tool_result_offloading/__init__.py new file mode 100644 index 0000000000..c1764a6e03 --- /dev/null +++ b/test/hooks/tool_result_offloading/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/test/hooks/tool_result_offloading/test_hooks.py b/test/hooks/tool_result_offloading/test_hooks.py new file mode 100644 index 0000000000..4229183e8c --- /dev/null +++ b/test/hooks/tool_result_offloading/test_hooks.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from typing import Annotated, Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from haystack import component +from haystack.components.agents import Agent +from haystack.components.agents.state.state import State +from haystack.dataclasses import ChatMessage, ImageContent, TextContent, ToolCall +from haystack.hooks.tool_result_offloading import ( + RESULT_STORE_CONTEXT_KEY, + AlwaysOffload, + FileSystemToolResultStore, + NeverOffload, + OffloadOverChars, + ToolResultOffloadHook, +) +from haystack.tools import Tool, Toolset, tool + + +@component +class MockChatGenerator: + @component.output_types(replies=list[ChatMessage]) + def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]: + return {"replies": [ChatMessage.from_assistant("done")]} + + @component.output_types(replies=list[ChatMessage]) + async def run_async( + self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs + ) -> dict[str, Any]: + return {"replies": [ChatMessage.from_assistant("done")]} + + +@tool +def big_tool(query: Annotated[str, "the query"]) -> str: + """Return a large result.""" + return "R" * 500 + + +def _state_with_messages(messages: list[ChatMessage]) -> State: + return State(schema={"messages": {"type": list[ChatMessage]}}, data={"messages": messages, "step_count": 1}) + + +def _tool_message(tool_name: str, result: str, *, error: bool = False, call_id: str = "c1") -> ChatMessage: + return ChatMessage.from_tool( + tool_result=result, origin=ToolCall(tool_name=tool_name, arguments={}, id=call_id), error=error + ) + + +class TestToolResultOffloadHookRouting: + def test_exact_tuple_and_wildcard_keys(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), + offload_strategies={ + "a": AlwaysOffload(), + ("b", "c"): AlwaysOffload(), + "d": NeverOffload(), + "*": AlwaysOffload(), + }, + ) + state = _state_with_messages( + [ + _tool_message("a", "A" * 50, call_id="1"), # exact -> offload + _tool_message("b", "B" * 50, call_id="2"), # tuple -> offload + _tool_message("d", "D" * 50, call_id="3"), # exact NeverOffload -> keep + _tool_message("z", "Z" * 50, call_id="4"), # wildcard -> offload + ] + ) + hook.run(state) + results = [m.tool_call_result.result for m in state.data["messages"]] + assert results[0].startswith("Tool result offloaded") + assert results[1].startswith("Tool result offloaded") + assert results[2] == "D" * 50 + assert results[3].startswith("Tool result offloaded") + + def test_tool_without_matching_key_is_not_offloaded(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"a": AlwaysOffload()} + ) + state = _state_with_messages([_tool_message("b", "B" * 50)]) + hook.run(state) + assert state.data["messages"][0].tool_call_result.result == "B" * 50 + + def test_over_chars_threshold(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": OffloadOverChars(10)} + ) + state = _state_with_messages( + [_tool_message("a", "x" * 10, call_id="1"), _tool_message("a", "x" * 11, call_id="2")] + ) + hook.run(state) + results = [m.tool_call_result.result for m in state.data["messages"]] + assert results[0] == "x" * 10 + assert results[1].startswith("Tool result offloaded") + + +class TestToolResultOffloadHookBehavior: + def test_error_results_are_never_offloaded(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} + ) + state = _state_with_messages([_tool_message("a", "boom", error=True)]) + hook.run(state) + assert state.data["messages"][0].tool_call_result.result == "boom" + + def test_offloads_sequence_of_text_content(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) + message = ChatMessage.from_tool( + tool_result=[TextContent("A" * 30), TextContent("B" * 30)], + origin=ToolCall(tool_name="a", arguments={}, id="1"), + ) + state = _state_with_messages([message]) + hook.run(state) + offloaded = state.data["messages"][0] + assert offloaded.tool_call_result.result.startswith("Tool result offloaded") + assert "60 characters" in offloaded.tool_call_result.result + assert store.read(offloaded.meta["tool_result_offloaded"]) == "A" * 30 + "B" * 30 + + def test_result_with_image_content_is_not_offloaded(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} + ) + content = [TextContent("caption"), ImageContent(base64_image="aGVsbG8=", mime_type="image/png")] + message = ChatMessage.from_tool(tool_result=content, origin=ToolCall(tool_name="a", arguments={}, id="1")) + state = _state_with_messages([message]) + hook.run(state) + assert state.data["messages"][0].tool_call_result.result == content + assert not list(Path(tmp_path).iterdir()) + + def test_id_less_parallel_calls_do_not_collide(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) + first = ChatMessage.from_tool(tool_result="FIRST" * 20, origin=ToolCall(tool_name="a", arguments={}, id=None)) + second = ChatMessage.from_tool(tool_result="SECOND" * 20, origin=ToolCall(tool_name="a", arguments={}, id=None)) + state = _state_with_messages([first, second]) + hook.run(state) + refs = [m.meta["tool_result_offloaded"] for m in state.data["messages"]] + assert refs[0] != refs[1] + assert store.read(refs[0]) == "FIRST" * 20 + assert store.read(refs[1]) == "SECOND" * 20 + assert len(list(Path(tmp_path).iterdir())) == 2 + + def test_only_trailing_tool_results_are_offloaded(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) + # A tool result from a prior turn, then an assistant message, then this step's fresh tool result. + history = _tool_message("old", "H" * 50, call_id="old1") + assistant = ChatMessage.from_assistant(tool_calls=[ToolCall("a", {}, id="c1")]) + fresh = _tool_message("a", "F" * 50, call_id="c1") + state = _state_with_messages([history, assistant, fresh]) + hook.run(state) + out = state.data["messages"] + assert out[0].tool_call_result.result == "H" * 50 + assert out[2].tool_call_result.result.startswith("Tool result offloaded") + assert len(list(Path(tmp_path).iterdir())) == 1 + + def test_second_offload_hook_does_not_reoffload_pointer(self, tmp_path): + # Two offload hooks under `after_tool` run in sequence on the same state. The first offloads the result and + # marks it; the `_OFFLOADED_META_KEY` marker stops the second from offloading the pointer text again. + store = FileSystemToolResultStore(root=tmp_path) + first_hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) + second_hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) + state = _state_with_messages([_tool_message("a", "A" * 50)]) + first_hook.run(state) + pointer = state.data["messages"][0].tool_call_result.result + second_hook.run(state) + assert state.data["messages"][0].tool_call_result.result == pointer + assert len(list(Path(tmp_path).iterdir())) == 1 + + def test_pointer_contains_reference_size_and_preview(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()}, preview_chars=5 + ) + state = _state_with_messages([_tool_message("a", "ABCDEFGH")]) + hook.run(state) + message = state.data["messages"][0] + reference = message.meta["tool_result_offloaded"] + pointer = message.tool_call_result.result + assert reference in pointer + assert "8 characters" in pointer + assert "ABCDE..." in pointer + + def test_concurrent_runs_are_isolated_by_per_run_store(self, tmp_path): + # One shared hook instance, two runs each supplying its own store via hook_context. Even with identical store + # keys (same tool, id and step), each run writes to and reads from its own store — no cross-run collision. + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path / "shared"), offload_strategies={"*": AlwaysOffload()} + ) + store_a = FileSystemToolResultStore(root=tmp_path / "a") + store_b = FileSystemToolResultStore(root=tmp_path / "b") + state_a = _state_with_messages([_tool_message("t", "AAA" * 20, call_id="x")]) + state_a.data["hook_context"] = {RESULT_STORE_CONTEXT_KEY: store_a} + state_b = _state_with_messages([_tool_message("t", "BBB" * 20, call_id="x")]) + state_b.data["hook_context"] = {RESULT_STORE_CONTEXT_KEY: store_b} + hook.run(state_a) + hook.run(state_b) + ref_a = state_a.data["messages"][0].meta["tool_result_offloaded"] + ref_b = state_b.data["messages"][0].meta["tool_result_offloaded"] + assert store_a.read(ref_a) == "AAA" * 20 + assert store_b.read(ref_b) == "BBB" * 20 + assert not (tmp_path / "shared").exists() + + def test_hook_context_store_overrides_constructor_store(self, tmp_path): + default_store = FileSystemToolResultStore(root=tmp_path / "default") + request_store = FileSystemToolResultStore(root=tmp_path / "request") + hook = ToolResultOffloadHook(store=default_store, offload_strategies={"*": AlwaysOffload()}) + state = _state_with_messages([_tool_message("a", "A" * 50)]) + state.data["hook_context"] = {RESULT_STORE_CONTEXT_KEY: request_store} + hook.run(state) + assert (tmp_path / "request").exists() + assert not (tmp_path / "default").exists() + + +class TestToolResultOffloadHookSerde: + def test_to_dict_from_dict_roundtrip(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), + offload_strategies={"a": AlwaysOffload(), ("b", "c"): OffloadOverChars(100), "*": NeverOffload()}, + preview_chars=42, + ) + restored = ToolResultOffloadHook.from_dict(hook.to_dict()) + assert restored.preview_chars == 42 + assert isinstance(restored.store, FileSystemToolResultStore) + assert set(restored.offload_strategies) == {"a", ("b", "c"), "*"} + assert isinstance(restored.offload_strategies[("b", "c")], OffloadOverChars) + assert restored.offload_strategies[("b", "c")].threshold == 100 + + +class TestToolResultOffloadHookInAgent: + def test_offloads_tool_result_seen_by_next_llm_call(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} + ) + agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]}) + agent.warm_up() + agent.chat_generator.run = MagicMock( + side_effect=[ + {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("big_tool", {"query": "x"})])]}, + {"replies": [ChatMessage.from_assistant("done")]}, + ] + ) + agent.run(messages=[ChatMessage.from_user("hi")]) + second_call_messages = agent.chat_generator.run.call_args_list[1].kwargs["messages"] + offloaded = [m for m in second_call_messages if m.tool_call_result is not None] + assert offloaded[0].tool_call_result.result.startswith("Tool result offloaded") + assert len(list(Path(tmp_path).iterdir())) == 1 + + +class TestToolResultOffloadHookInAgentAsync: + @pytest.mark.asyncio + async def test_offloads_tool_result_async(self, tmp_path): + hook = ToolResultOffloadHook( + store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} + ) + agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]}) + agent.warm_up() + agent.chat_generator.run_async = AsyncMock( + side_effect=[ + {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("big_tool", {"query": "x"})])]}, + {"replies": [ChatMessage.from_assistant("done")]}, + ] + ) + result = await agent.run_async(messages=[ChatMessage.from_user("hi")]) + offloaded = [m for m in result["messages"] if m.tool_call_result is not None] + assert offloaded[0].tool_call_result.result.startswith("Tool result offloaded") diff --git a/test/hooks/tool_result_offloading/test_policies.py b/test/hooks/tool_result_offloading/test_policies.py new file mode 100644 index 0000000000..aa2697b889 --- /dev/null +++ b/test/hooks/tool_result_offloading/test_policies.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from haystack.components.agents.state.state import State +from haystack.hooks.tool_result_offloading import AlwaysOffload, NeverOffload, OffloadOverChars + + +class TestOffloadPolicies: + def test_always_offload(self): + assert AlwaysOffload().should_offload("t", "anything", State(schema={})) is True + + def test_never_offload(self): + assert NeverOffload().should_offload("t", "x" * 10_000, State(schema={})) is False + + def test_offload_over_chars_is_strictly_greater(self): + policy = OffloadOverChars(threshold=10) + assert policy.should_offload("t", "x" * 10, State(schema={})) is False + assert policy.should_offload("t", "x" * 11, State(schema={})) is True + + def test_offload_over_chars_roundtrip(self): + restored = OffloadOverChars.from_dict(OffloadOverChars(threshold=42).to_dict()) + assert restored.threshold == 42 diff --git a/test/hooks/tool_result_offloading/test_stores.py b/test/hooks/tool_result_offloading/test_stores.py new file mode 100644 index 0000000000..6522639d08 --- /dev/null +++ b/test/hooks/tool_result_offloading/test_stores.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest + +from haystack.hooks.tool_result_offloading import FileSystemToolResultStore + + +class TestFileSystemToolResultStore: + def test_write_returns_path_and_persists_content(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + reference = store.write(key="a.txt", content="hello") + assert reference == str(tmp_path / "a.txt") + assert Path(reference).read_text(encoding="utf-8") == "hello" + + def test_write_creates_missing_directories(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path / "nested" / "dir") + reference = store.write(key="a.txt", content="hi") + assert Path(reference).read_text(encoding="utf-8") == "hi" + + def test_write_allows_nested_keys_within_root(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + reference = store.write(key="sub/dir/a.txt", content="ok") + assert Path(reference).read_text(encoding="utf-8") == "ok" + + def test_write_rejects_parent_traversal_key(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path / "root") + with pytest.raises(ValueError, match="outside the store root"): + store.write(key="../escape.txt", content="x") + assert not (tmp_path / "escape.txt").exists() + + def test_write_rejects_absolute_key(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path / "root") + with pytest.raises(ValueError, match="outside the store root"): + store.write(key=str(tmp_path / "outside.txt"), content="x") + + def test_read_round_trips_written_content(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + reference = store.write(key="a.txt", content="round trip") + assert store.read(reference) == "round trip" + + def test_to_dict_from_dict_roundtrip(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + restored = FileSystemToolResultStore.from_dict(store.to_dict()) + assert restored.root == Path(tmp_path)