mirror of
https://github.com/deepset-ai/haystack.git
synced 2026-10-12 02:12:50 +00:00
feat: Add tool result offloading (#11849)
This commit is contained in:
@@ -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
@@ -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)
|
||||
Reference in New Issue
Block a user