feat: Add tool result offloading (#11849)

This commit is contained in:
Sebastian Husch Lee
2026-07-07 15:00:11 +02:00
committed by GitHub
parent a1c5ec63dd
commit 48d98650f0
12 changed files with 1026 additions and 2 deletions
@@ -0,0 +1,27 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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)
@@ -0,0 +1,354 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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)
@@ -0,0 +1,70 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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)
@@ -0,0 +1,78 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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 `<root>/<key>`, 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)
@@ -0,0 +1,7 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
from .protocol import OffloadPolicy, ToolResultStore
__all__ = ["OffloadPolicy", "ToolResultStore"]
@@ -0,0 +1,76 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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)
+4 -2
View File
@@ -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
@@ -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")])
@@ -0,0 +1,3 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,271 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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")
@@ -0,0 +1,23 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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
@@ -0,0 +1,48 @@
# SPDX-FileCopyrightText: 2022-present deepset GmbH <[email protected]>
#
# 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)