diff --git a/haystack/hooks/compaction/__init__.py b/haystack/hooks/compaction/__init__.py index fb0ae0a163..b5357752ee 100644 --- a/haystack/hooks/compaction/__init__.py +++ b/haystack/hooks/compaction/__init__.py @@ -11,6 +11,7 @@ _import_structure = { "hooks": ["CompactionHook"], "sliding_window": ["SlidingWindowCompactor"], "summarization": ["SummarizationCompactor"], + "tool_result_offloading": ["ToolResultOffloadCompactor"], "tool_result_pruning": ["ToolResultPruningCompactor"], "types": ["Compactor"], } @@ -19,6 +20,7 @@ if TYPE_CHECKING: from .hooks import CompactionHook as CompactionHook from .sliding_window import SlidingWindowCompactor as SlidingWindowCompactor from .summarization import SummarizationCompactor as SummarizationCompactor + from .tool_result_offloading import ToolResultOffloadCompactor as ToolResultOffloadCompactor from .tool_result_pruning import ToolResultPruningCompactor as ToolResultPruningCompactor from .types import Compactor as Compactor else: diff --git a/haystack/hooks/compaction/tool_result_offloading.py b/haystack/hooks/compaction/tool_result_offloading.py new file mode 100644 index 0000000000..8a69cd9965 --- /dev/null +++ b/haystack/hooks/compaction/tool_result_offloading.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from typing import Any +from uuid import uuid4 + +from haystack.core.serialization import default_from_dict, default_to_dict +from haystack.dataclasses import ChatMessage +from haystack.hooks.compaction.types import Compactor +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _replace_tool_results_until_target +from haystack.hooks.tool_result_offloading.types import ToolResultStore +from haystack.hooks.tool_result_offloading.utils import ( + _OFFLOADED_META_KEY, + _offloadable_content_blocks, + _offloaded_message, +) +from haystack.token_counters import TokenCounter +from haystack.utils.deserialization import deserialize_component_inplace +from haystack.utils.experimental import _experimental + + +@_experimental +class ToolResultOffloadCompactor(Compactor): + """ + Writes older tool results to a `ToolResultStore` and leaves a reference in their place. + + Every tool call keeps its matching result, and the model can read the full output back with a read tool you scope + to the same store. `CompactionHook` only calls the compactor once the conversation reaches `compact_at`, so fresh + output stays in context until then. Use `ToolResultOffloadHook` to offload a tool's output as soon as it arrives. + + + ```python + from typing import Annotated + + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIResponsesChatGenerator + from haystack.hooks.compaction import CompactionHook, ToolResultOffloadCompactor + from haystack.hooks.tool_result_offloading import FileSystemToolResultStore + from haystack.tools import tool + + store = FileSystemToolResultStore(root="tool_results") + + + @tool + def read_offloaded_result(path: Annotated[str, "Path of an offloaded tool result"]) -> str: + '''Read back the full content of an offloaded tool result.''' + content = store.read(path) + if isinstance(content, bytes): + return f"'{path}' holds {len(content)} bytes of binary content and cannot be read as text." + return content + + + # `web_search` stands in for your own tool that returns long results. + hook = CompactionHook(compactor=ToolResultOffloadCompactor(store=store), context_window=400_000) + agent = Agent( + chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.4-nano"), + tools=[web_search, read_offloaded_result], + hooks={"before_llm": [hook]}, + ) + ``` + + The compactor always writes to the store it was created with; unlike `ToolResultOffloadHook`, it does not read a + per-run store from `hook_context`. In a multi-user server, create an Agent per run with its own compactor and + store, so users never read each other's results. + + For results with images or files, we recommend passing a provider token counter such as `OpenAITokenCounter` to + `CompactionHook`, which measures their real size. Its default `ApproximateTokenCounter` charges a flat + `tokens_per_image` and `tokens_per_file` instead, which can undercount such results and keep them below + `min_tokens`. + """ + + def __init__( + self, store: ToolResultStore, *, min_keep_steps: int = 1, min_tokens: int = 200, preview_chars: int = 200 + ) -> None: + """ + Initialize the compactor. + + :param store: Where offloaded results are written. Image and file results are only written to a store that sets + `supports_binary_content`; otherwise they stay in the conversation. + :param min_keep_steps: Number of most recent tool-calling Agent steps whose results are never offloaded, even + when the target is missed. Must be at least 1, so the model always sees the latest results. + :param min_tokens: Only offload tool-result messages larger than this many tokens. + :param preview_chars: Number of leading characters of each offloaded text kept in its reference. Image and file + results are described by MIME type and size instead. + :raises ValueError: If `min_keep_steps` is less than 1, or `min_tokens` or `preview_chars` is negative. + """ + if min_keep_steps < 1: + raise ValueError( + f"`min_keep_steps` must be at least 1, got {min_keep_steps}. The most recent tool-calling step " + f"contains results the model may still need." + ) + if min_tokens < 0: + raise ValueError(f"`min_tokens` must be at least 0, got {min_tokens}.") + if preview_chars < 0: + raise ValueError(f"`preview_chars` must be at least 0, got {preview_chars}.") + self.store = store + self.min_keep_steps = min_keep_steps + self.min_tokens = min_tokens + self.preview_chars = preview_chars + + def compact( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Offload tool results to the store, oldest first, until the conversation fits within `target_tokens`. + + Results from the most recent `min_keep_steps` tool-calling Agent steps are never offloaded. + + :param messages: The conversation to compact, oldest to newest. + :param target_tokens: The token count the compacted conversation should fit within. + :param token_counter: The `TokenCounter` used to measure the conversation before and after each replacement. + :returns: The conversation with older tool results replaced by references, or None when nothing was offloaded. + """ + return _replace_tool_results_until_target( + messages=messages, + target_tokens=target_tokens, + token_counter=token_counter, + min_keep_steps=self.min_keep_steps, + replace=lambda message, _index: self._offload(message=message, token_counter=token_counter), + ) + + async def compact_async( + self, messages: list[ChatMessage], target_tokens: int, token_counter: TokenCounter + ) -> list[ChatMessage] | None: + """ + Run `compact` in a thread so store writes and token counting do not block the event loop. + + :param messages: The conversation to compact, oldest to newest. + :param target_tokens: The token count the compacted conversation should fit within. + :param token_counter: The `TokenCounter` used to measure the conversation before and after each replacement. + :returns: The conversation with older tool results replaced by references, or None when nothing was offloaded. + """ + return await asyncio.to_thread( + self.compact, messages=messages, target_tokens=target_tokens, token_counter=token_counter + ) + + def _offload(self, message: ChatMessage, token_counter: TokenCounter) -> tuple[ChatMessage, int] | None: + """ + Write one tool result to the store and report how many tokens its reference saves. + + :param message: The tool-result message to consider. + :param token_counter: The `TokenCounter` that measures the result and its reference. + :returns: The offloaded message and the number of tokens it saves, or None when the result stays in context. + """ + result = message.tool_call_result + # Skip errors, results another compactor already rewrote, and results that are already a reference. + if ( + result is None + or result.error + or message.meta.get(_OFFLOADED_META_KEY) + or _COMPACTION_META_KEY in message.meta + ): + return None + + original_tokens = token_counter.count(messages=[message]) + if original_tokens <= self.min_tokens: + return None + + # An empty result, or image or file content a text-only store cannot take, stays in context. + content_blocks = _offloadable_content_blocks(result=result, store=self.store) + if content_blocks is None: + return None + + # A random suffix keeps store keys unique across runs that share the store, with or without a tool call id. + offloaded = _offloaded_message( + message=message, + content_blocks=content_blocks, + store=self.store, + key_prefix=f"compacted_{result.origin.tool_name}_{uuid4().hex}", + preview_chars=self.preview_chars, + additional_meta={ + _COMPACTION_META_KEY: {"strategy": "tool_result_offloading", "original_tokens": original_tokens} + }, + ) + saved_tokens = original_tokens - token_counter.count(messages=[offloaded]) + # A reference can cost more than a small result. That replacement is dropped and its store entry left unused. + return (offloaded, saved_tokens) if saved_tokens > 0 else None + + def to_dict(self) -> dict[str, Any]: + """ + Serialize the compactor, including its store. + + :returns: A dictionary representation of the compactor. + """ + return default_to_dict( + self, + store=self.store.to_dict(), + min_keep_steps=self.min_keep_steps, + min_tokens=self.min_tokens, + preview_chars=self.preview_chars, + ) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ToolResultOffloadCompactor": + """ + Deserialize the compactor, reconstructing its store. + + :param data: A dictionary representation produced by `to_dict`. + :returns: The deserialized `ToolResultOffloadCompactor`. + """ + init_params = data.get("init_parameters", {}) + if init_params.get("store") is not None: + deserialize_component_inplace(init_params, key="store") + return default_from_dict(cls, data) diff --git a/haystack/hooks/compaction/tool_result_pruning.py b/haystack/hooks/compaction/tool_result_pruning.py index df4abb2af9..ad45b82795 100644 --- a/haystack/hooks/compaction/tool_result_pruning.py +++ b/haystack/hooks/compaction/tool_result_pruning.py @@ -7,7 +7,7 @@ from typing import Any from haystack.core.serialization import default_to_dict from haystack.dataclasses import ChatMessage from haystack.hooks.compaction.types import Compactor -from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _agent_step_spans +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY, _replace_tool_results_until_target from haystack.token_counters import TokenCounter from haystack.utils.experimental import _experimental @@ -96,43 +96,13 @@ class ToolResultPruningCompactor(Compactor): :param token_counter: The `TokenCounter` used to measure the conversation before and after each replacement. :returns: The conversation with older tool results replaced, or None when no result was prunable. """ - current_tokens = token_counter.count(messages=messages) - if current_tokens <= target_tokens: - return None - - # Filter the shared Agent-step spans to tool-calling steps; parallel results remain grouped in one span. - result_steps = [ - list(range(start + 1, end)) - for start, end in _agent_step_spans(messages=messages, start=0) - # An assistant-only span has no result to protect and must not consume one of `min_keep_steps`. - if end > start + 1 - ] - - protected_positions = {position for step in result_steps[-self.min_keep_steps :] for position in step} - - # Replace entries in a new list so the caller-owned input list remains unchanged. - compacted = list(messages) - changed = False - # Iterate oldest-first so we can stop at the target while leaving as much recent output intact as possible. - for index, message in enumerate(messages): - if message.tool_call_result is None or index in protected_positions: - continue - replacement = self._prune(message=message, token_counter=token_counter) - if replacement is None: - continue - pruned, saved_tokens = replacement - - # Update the compacted list and the running token count - compacted[index] = pruned - current_tokens -= saved_tokens - changed = True - - # Once we reach the target, stop pruning to preserve as much recent output as possible. - if current_tokens <= target_tokens: - break - - # If no candidates were pruned, return None to indicate no change. - return compacted if changed else None + return _replace_tool_results_until_target( + messages=messages, + target_tokens=target_tokens, + token_counter=token_counter, + min_keep_steps=self.min_keep_steps, + replace=lambda message, _index: self._prune(message=message, token_counter=token_counter), + ) def _prune(self, message: ChatMessage, token_counter: TokenCounter) -> tuple[ChatMessage, int] | None: """ diff --git a/haystack/hooks/compaction/utils.py b/haystack/hooks/compaction/utils.py index 5ba4648f3d..dcca987d99 100644 --- a/haystack/hooks/compaction/utils.py +++ b/haystack/hooks/compaction/utils.py @@ -2,7 +2,10 @@ # # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable + from haystack.dataclasses import ChatMessage, ChatRole +from haystack.token_counters import TokenCounter # Meta key marking a message that a compactor produced. Its value records which strategy ran. _COMPACTION_META_KEY = "context_compaction" @@ -129,3 +132,58 @@ def _historical_turn_groups(messages: list[ChatMessage], system_end: int, task_i list(range(start, end)) for start, end in _historical_turn_spans(messages=messages, start=system_end, end=historical_end) ] + + +def _replace_tool_results_until_target( + messages: list[ChatMessage], + *, + target_tokens: int, + token_counter: TokenCounter, + min_keep_steps: int, + replace: Callable[[ChatMessage, int], tuple[ChatMessage, int] | None], +) -> list[ChatMessage] | None: + """ + Replace tool results oldest first until the conversation fits within `target_tokens`. + + Results from the most recent `min_keep_steps` tool-calling Agent steps are never passed to `replace`. After + measuring the initial conversation, the running total is updated with each replacement's saved tokens to avoid + recounting the full context. + + :param messages: The conversation to compact, oldest to newest. + :param target_tokens: The token count the compacted conversation should fit within. + :param token_counter: The `TokenCounter` used to measure the initial conversation. + :param min_keep_steps: Number of most recent tool-calling Agent steps whose results are left unchanged. + :param replace: Called with a tool-result message and its position. Returns the replacement and the tokens it + saves, or None to leave the message unchanged. + :returns: The conversation with replaced tool results, or None when nothing was replaced. + """ + current_tokens = token_counter.count(messages=messages) + if current_tokens <= target_tokens: + return None + + # Filter the shared Agent-step spans to tool-calling steps; parallel results remain grouped in one span. + result_steps = [ + list(range(start + 1, end)) + for start, end in _agent_step_spans(messages=messages, start=0) + # An assistant-only span has no result to protect and must not consume one of `min_keep_steps`. + if end > start + 1 + ] + protected_positions = {position for step in result_steps[-min_keep_steps:] for position in step} + + # Replace entries in a new list so the caller-owned input list remains unchanged. + compacted = list(messages) + changed = False + # Go oldest first and stop at the target, so as much recent output as possible stays in context. + for index, message in enumerate(messages): + if message.tool_call_result is None or index in protected_positions: + continue + replacement = replace(message, index) + if replacement is None: + continue + compacted[index], saved_tokens = replacement + current_tokens -= saved_tokens + changed = True + if current_tokens <= target_tokens: + break + + return compacted if changed else None diff --git a/haystack/hooks/tool_result_offloading/hooks.py b/haystack/hooks/tool_result_offloading/hooks.py index 36d978e2e3..b9d091505e 100644 --- a/haystack/hooks/tool_result_offloading/hooks.py +++ b/haystack/hooks/tool_result_offloading/hooks.py @@ -2,30 +2,22 @@ # # SPDX-License-Identifier: Apache-2.0 -import base64 import json -import mimetypes -from pathlib import Path 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, FileContent, ImageContent, TextContent +from haystack.dataclasses import ChatMessage from haystack.hooks.tool_result_offloading.types import OffloadPolicy, ToolResultStore +from haystack.hooks.tool_result_offloading.utils import ( + _OFFLOADED_META_KEY, + _content_block_payload, + _offloadable_content_blocks, + _offloaded_message, +) from haystack.utils.deserialization import deserialize_component_inplace -# Extension used for a binary block whose MIME type is unknown or maps to no known extension. -_FALLBACK_EXTENSION = ".bin" - -logger = logging.getLogger(__name__) - -# Meta key marking an already-offloaded tool-result message; its value is the list of store references written. -# Stops a second `after_tool` offload hook from offloading the pointer text again, since the pointer is itself a -# tool result. -_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" @@ -49,20 +41,6 @@ def _fresh_tool_results_start(messages: list[ChatMessage]) -> int: return index -def _content_block_payload(content_block: TextContent | ImageContent | FileContent) -> str: - """ - Return the string a content block contributes to the conversation. - - For an image or a file this is the base64 payload, which is what actually occupies the context window. - - :param content_block: The content block to inspect. - :returns: The content block's text or base64 payload. - """ - if isinstance(content_block, TextContent): - return content_block.text - return content_block.base64_image if isinstance(content_block, ImageContent) else content_block.base64_data - - 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. @@ -269,26 +247,9 @@ class ToolResultOffloadHook: if policy is None: return message - # A plain string result is handled as a single text block, so everything below has one shape to work with. - content_blocks: list[TextContent | ImageContent | FileContent] = ( - [TextContent(text=result.result)] if isinstance(result.result, str) else list(result.result) - ) - # A result made up of nothing but empty text has nothing worth storing, so it stays in context. `all` also - # covers a result with no content blocks at all. - if all(isinstance(content_block, TextContent) and not content_block.text for content_block in content_blocks): - return message - - # Check whether the store can store binary content before offloading an image or file result. A text-only store - # leaves the result in context and logs a warning. - if not getattr(store, "supports_binary_content", False) and not all( - isinstance(content_block, TextContent) for content_block in content_blocks - ): - logger.warning( - "Tool '{tool}' produced a result with image or file content, but {store} does not support binary " - "content; leaving the result in context.", - tool=tool_name, - store=type(store).__name__, - ) + # An empty result, or image or file content a text-only store cannot take, stays in context. + content_blocks = _offloadable_content_blocks(result=result, store=store) + if content_blocks is None: return message # The policy sizes up the result by the string it occupies in the conversation, which for an image or a file @@ -300,81 +261,14 @@ class ToolResultOffloadHook: # Step, tool name, and call id keep results from different tools and steps from colliding. A tool call id is # optional, so an id-less call falls back to its position in this step's batch. step = state.data.get("step_count", 0) - prefix = f"{step}_{tool_name}_{result.origin.id or f'call{index}'}" - references, pointer = self._offload_content_blocks(content_blocks=content_blocks, store=store, prefix=prefix) - - return ChatMessage.from_tool( - tool_result=pointer, - origin=result.origin, - error=result.error, - meta={**message.meta, _OFFLOADED_META_KEY: references}, + return _offloaded_message( + message=message, + content_blocks=content_blocks, + store=store, + key_prefix=f"{step}_{tool_name}_{result.origin.id or f'call{index}'}", + preview_chars=self.preview_chars, ) - def _offload_content_blocks( - self, content_blocks: list[TextContent | ImageContent | FileContent], store: ToolResultStore, prefix: str - ) -> tuple[list[str], str]: - """ - Write a result's content blocks to the store and build the pointer that replaces them in the conversation. - - Every content block goes to its own store entry. A single block keeps `prefix` as its key and gets a one-line - pointer; several blocks get position-suffixed keys and one numbered pointer line each. - - :param content_blocks: The result's content blocks, in order. - :param store: The store to write to. - :param prefix: The result's store key prefix, as described above. - :returns: The store references written, and the pointer text for the conversation. - """ - references: list[str] = [] - descriptions: list[str] = [] - single = len(content_blocks) == 1 - for position, content_block in enumerate(content_blocks): - reference, description = self._offload_content_block( - content_block=content_block, store=store, key_prefix=prefix if single else f"{prefix}_{position}" - ) - references.append(reference) - descriptions.append(description) - - if len(descriptions) == 1: - return references, f"Tool result offloaded to {descriptions[0]}" - - numbered = [f"{position}. {description}" for position, description in enumerate(descriptions, start=1)] - return references, "\n".join([f"Tool result offloaded to {len(descriptions)} files:", *numbered]) - - def _offload_content_block( - self, content_block: TextContent | ImageContent | FileContent, store: ToolResultStore, key_prefix: str - ) -> tuple[str, str]: - """ - Write a single content block to the store and describe where it went. - - :param content_block: The content block to offload. - :param store: The store to write to. - :param key_prefix: The content block's store key without its extension. - :returns: The store reference the content block was written to, and a one-line description for the pointer. - """ - if isinstance(content_block, TextContent): - text = content_block.text - reference = store.write(key=f"{key_prefix}.txt", content=text) - # An ellipsis marks a preview that was cut short, so the model can tell it is not the whole text. - preview = f"{text[: self.preview_chars]}{'...' if len(text) > self.preview_chars else ''}" - return reference, f"text ({len(text)} characters) at '{reference}'. Preview: {preview}" - - if isinstance(content_block, ImageContent): - data = base64.b64decode(content_block.base64_image) - label = content_block.mime_type or "image" - filename = None - else: - data = base64.b64decode(content_block.base64_data) - label = content_block.mime_type or "file" - filename = content_block.filename - - # What the tool called the file wins over its MIME type. Only the suffix is taken. - mime_extension = mimetypes.guess_extension(content_block.mime_type) if content_block.mime_type else None - extension = (Path(filename).suffix if filename else "") or mime_extension or _FALLBACK_EXTENSION - reference = store.write(key=f"{key_prefix}{extension}", content=data) - - named = f" named '{filename}'" if filename else "" - return reference, f"{label}{named} ({len(data)} bytes) at '{reference}'" - def to_dict(self) -> dict[str, Any]: """ Serialize the hook, including its store and per-tool offload strategies. diff --git a/haystack/hooks/tool_result_offloading/types/protocol.py b/haystack/hooks/tool_result_offloading/types/protocol.py index 59c25e4ca4..acb4154bc7 100644 --- a/haystack/hooks/tool_result_offloading/types/protocol.py +++ b/haystack/hooks/tool_result_offloading/types/protocol.py @@ -10,7 +10,7 @@ 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. + Where a `ToolResultOffloadHook` or `ToolResultOffloadCompactor` writes tool results and a read tool fetches them. Implementations decide where and how the content lives (local disk, an isolated sandbox filesystem, object storage, ...). `write` returns a reference string that the Agent puts in the conversation in place of the full @@ -32,8 +32,8 @@ class ToolResultStore(Protocol): """ Persist `content` under `key` and return a reference to it. - :param key: A stable, per-result identifier the hook derives from the tool call (e.g. a file name). It carries - an extension matching the content, so a store that maps keys to files can use it as-is. + :param key: A stable, per-result identifier derived from the tool call (e.g. a file name). It carries an + extension matching the content, so a store that maps keys to files can use it as-is. :param content: The tool result to persist. Text arrives as a string. Image and file content arrives as the decoded bytes of its base64 payload, and only when the store sets `supports_binary_content` to True - a text-only store may narrow this parameter to `str`. diff --git a/haystack/hooks/tool_result_offloading/utils.py b/haystack/hooks/tool_result_offloading/utils.py new file mode 100644 index 0000000000..9486e73561 --- /dev/null +++ b/haystack/hooks/tool_result_offloading/utils.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import base64 +import mimetypes +from pathlib import Path +from typing import Any + +from haystack import logging +from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCallResult +from haystack.hooks.tool_result_offloading.types import ToolResultStore + +logger = logging.getLogger(__name__) + +# Meta key marking an already-offloaded tool-result message; its value is the list of store references written. +# Stops a second offload from offloading the pointer text again, since the pointer is itself a tool result. +_OFFLOADED_META_KEY = "tool_result_offloaded" + +# Extension used for a binary block whose MIME type is unknown or maps to no known extension. +_FALLBACK_EXTENSION = ".bin" + + +def _content_block_payload(content_block: TextContent | ImageContent | FileContent) -> str: + """ + Return the string a content block contributes to the conversation. + + For an image or a file this is the base64 payload, which is what actually occupies the context window. + + :param content_block: The content block to inspect. + :returns: The content block's text or base64 payload. + """ + if isinstance(content_block, TextContent): + return content_block.text + return content_block.base64_image if isinstance(content_block, ImageContent) else content_block.base64_data + + +def _offloadable_content_blocks( + result: ToolCallResult, store: ToolResultStore +) -> list[TextContent | ImageContent | FileContent] | None: + """ + Return a result's content blocks, or None when the result is empty or `store` cannot hold its image or file content. + + A warning is logged when a text-only store is given image or file content. + + :param result: The tool result to inspect. + :param store: The store the result would be written to. + :returns: The result's content blocks, or None when the result cannot or should not be offloaded. + """ + # A plain string result is handled as a single text block, so callers have one shape to work with. + content_blocks: list[TextContent | ImageContent | FileContent] = ( + [TextContent(text=result.result)] if isinstance(result.result, str) else list(result.result) + ) + + # `all` also covers a result with no content blocks at all. + if all(isinstance(content_block, TextContent) and not content_block.text for content_block in content_blocks): + return None + + if not getattr(store, "supports_binary_content", False) and not all( + isinstance(content_block, TextContent) for content_block in content_blocks + ): + logger.warning( + "Tool '{tool}' produced a result with image or file content, but {store} does not support binary " + "content; leaving the result in context.", + tool=result.origin.tool_name, + store=type(store).__name__, + ) + return None + + return content_blocks + + +def _offloaded_message( + message: ChatMessage, + *, + content_blocks: list[TextContent | ImageContent | FileContent], + store: ToolResultStore, + key_prefix: str, + preview_chars: int, + additional_meta: dict[str, Any] | None = None, +) -> ChatMessage: + """ + Write a tool result to the store and return a message with the same origin and error flag that points to it. + + :param message: The tool-result message being offloaded. Must carry a tool result. + :param content_blocks: The result's content blocks, as returned by `_offloadable_content_blocks`. + :param store: The store to write the content blocks to. + :param key_prefix: The result's store key prefix, unique within the run. + :param preview_chars: Number of leading characters of each offloaded text to include in the pointer. + :param additional_meta: Extra metadata to record on the offloaded message, such as a compaction marker. + :returns: A new tool-result message whose content points at the stored result. + """ + result = message.tool_call_result + if result is None: + raise ValueError("Only tool-result messages can be offloaded.") + + references, pointer = _offload_content_blocks( + content_blocks=content_blocks, store=store, prefix=key_prefix, preview_chars=preview_chars + ) + return ChatMessage.from_tool( + tool_result=pointer, + origin=result.origin, + error=result.error, + meta={**message.meta, **(additional_meta or {}), _OFFLOADED_META_KEY: references}, + ) + + +def _offload_content_blocks( + content_blocks: list[TextContent | ImageContent | FileContent], + store: ToolResultStore, + prefix: str, + preview_chars: int, +) -> tuple[list[str], str]: + """ + Write a result's content blocks to the store and build the pointer that replaces them in the conversation. + + Every content block goes to its own store entry. A single block keeps `prefix` as its key and gets a one-line + pointer; several blocks get position-suffixed keys and one numbered pointer line each. + + :param content_blocks: The result's content blocks, in order. + :param store: The store to write to. + :param prefix: The result's store key prefix, as described above. + :param preview_chars: Number of leading characters of each offloaded text to include in the pointer. + :returns: The store references written, and the pointer text for the conversation. + """ + references: list[str] = [] + descriptions: list[str] = [] + single = len(content_blocks) == 1 + for position, content_block in enumerate(content_blocks): + reference, description = _offload_content_block( + content_block=content_block, + store=store, + key_prefix=prefix if single else f"{prefix}_{position}", + preview_chars=preview_chars, + ) + references.append(reference) + descriptions.append(description) + + if len(descriptions) == 1: + return references, f"Tool result offloaded to {descriptions[0]}" + + numbered = [f"{position}. {description}" for position, description in enumerate(descriptions, start=1)] + return references, "\n".join([f"Tool result offloaded to {len(descriptions)} files:", *numbered]) + + +def _offload_content_block( + content_block: TextContent | ImageContent | FileContent, store: ToolResultStore, key_prefix: str, preview_chars: int +) -> tuple[str, str]: + """ + Write a single content block to the store and describe where it went. + + :param content_block: The content block to offload. + :param store: The store to write to. + :param key_prefix: The content block's store key without its extension. + :param preview_chars: Number of leading characters of an offloaded text to include in its description. + :returns: The store reference the content block was written to, and a one-line description for the pointer. + """ + if isinstance(content_block, TextContent): + text = content_block.text + reference = store.write(key=f"{key_prefix}.txt", content=text) + # An ellipsis marks a preview that was cut short, so the model can tell it is not the whole text. + preview = f"{text[:preview_chars]}{'...' if len(text) > preview_chars else ''}" + return reference, f"text ({len(text)} characters) at '{reference}'. Preview: {preview}" + + if isinstance(content_block, ImageContent): + data = base64.b64decode(content_block.base64_image) + label = content_block.mime_type or "image" + filename = None + else: + data = base64.b64decode(content_block.base64_data) + label = content_block.mime_type or "file" + filename = content_block.filename + + # What the tool called the file wins over its MIME type. Only the suffix is taken. + mime_extension = mimetypes.guess_extension(content_block.mime_type) if content_block.mime_type else None + extension = (Path(filename).suffix if filename else "") or mime_extension or _FALLBACK_EXTENSION + reference = store.write(key=f"{key_prefix}{extension}", content=data) + + named = f" named '{filename}'" if filename else "" + return reference, f"{label}{named} ({len(data)} bytes) at '{reference}'" diff --git a/pydoc/hooks_api.yml b/pydoc/hooks_api.yml index 9f35e24193..61226f62c3 100644 --- a/pydoc/hooks_api.yml +++ b/pydoc/hooks_api.yml @@ -1,7 +1,7 @@ loaders: - search_path: [../haystack/hooks] modules: ["protocol", "from_function", "budget/hooks", "compaction/hooks", "compaction/sliding_window", - "compaction/summarization", "compaction/tool_result_pruning", + "compaction/summarization", "compaction/tool_result_offloading", "compaction/tool_result_pruning", "compaction/types/protocol", "human_in_the_loop/dataclasses", "human_in_the_loop/hooks", "human_in_the_loop/policies", "human_in_the_loop/strategies", "human_in_the_loop/user_interfaces", "tool_result_offloading/hooks", "tool_result_offloading/policies", "tool_result_offloading/stores", diff --git a/releasenotes/notes/add-tool-result-offload-compactor-d8f8754dd98785fd.yaml b/releasenotes/notes/add-tool-result-offload-compactor-d8f8754dd98785fd.yaml new file mode 100644 index 0000000000..797859f03a --- /dev/null +++ b/releasenotes/notes/add-tool-result-offload-compactor-d8f8754dd98785fd.yaml @@ -0,0 +1,36 @@ +--- +features: + - | + Added the experimental ``ToolResultOffloadCompactor`` for ``CompactionHook``. Once compaction is triggered, it + writes older tool results to a ``ToolResultStore`` and leaves a reference in their place, so the Agent can read + the full output back with a read tool you scope to the same store. Fresh results stay in context until then. + + .. code-block:: python + + from typing import Annotated + + from haystack.components.agents import Agent + from haystack.components.generators.chat import OpenAIResponsesChatGenerator + from haystack.hooks.compaction import CompactionHook, ToolResultOffloadCompactor + from haystack.hooks.tool_result_offloading import FileSystemToolResultStore + from haystack.tools import tool + + store = FileSystemToolResultStore(root="tool_results") + + + @tool + def read_offloaded_result(path: Annotated[str, "Path of an offloaded tool result"]) -> str: + """Read back the full content of an offloaded tool result.""" + content = store.read(path) + if isinstance(content, bytes): + return f"'{path}' holds {len(content)} bytes of binary content and cannot be read as text." + return content + + + # web_search stands in for your own tool that returns long results. + hook = CompactionHook(compactor=ToolResultOffloadCompactor(store=store), context_window=400_000) + agent = Agent( + chat_generator=OpenAIResponsesChatGenerator(model="gpt-5.4-nano"), + tools=[web_search, read_offloaded_result], + hooks={"before_llm": [hook]}, + ) diff --git a/test/hooks/compaction/helpers.py b/test/hooks/compaction/helpers.py index 6966aade7f..51caadfad4 100644 --- a/test/hooks/compaction/helpers.py +++ b/test/hooks/compaction/helpers.py @@ -54,6 +54,15 @@ def tool_result(result: str, *, call_id: str = "c1", name: str = "search", error ) +def conversation(*results: str) -> list[ChatMessage]: + """A user task followed by one Agent step per given result.""" + messages = [ChatMessage.from_user("task")] + for index, result in enumerate(results): + call_id = f"c{index}" + messages.extend([tool_call(call_id), tool_result(result, call_id=call_id)]) + return messages + + def make_state(messages: list[ChatMessage], **data: Any) -> State: """A State shaped like the Agent's, holding `messages` and whatever run metadata a test overrides.""" base = { diff --git a/test/hooks/compaction/test_tool_result_offloading.py b/test/hooks/compaction/test_tool_result_offloading.py new file mode 100644 index 0000000000..80fe27a09b --- /dev/null +++ b/test/hooks/compaction/test_tool_result_offloading.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import base64 +from pathlib import Path + +import pytest + +from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCall +from haystack.dataclasses.chat_message import ChatMessageContentT +from haystack.hooks.compaction import CompactionHook, ToolResultOffloadCompactor +from haystack.hooks.compaction.utils import _COMPACTION_META_KEY +from haystack.hooks.tool_result_offloading import FileSystemToolResultStore +from haystack.hooks.tool_result_offloading.utils import _content_block_payload +from haystack.token_counters.utils import _rendered_conversation +from haystack.tools import ToolsType +from test.hooks.compaction.helpers import FakeCounter, conversation, tool_call, tool_result + +pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") + +COUNTER = FakeCounter(chars_per_token=1) + +# A 1x1 PNG, padded so its payload clearly outweighs the reference replacing it. +PNG_BYTES = ( + base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==") + + b"\x00" * 1000 +) + + +def _payload_placeholder(content: ChatMessageContentT) -> str: + """Render an image or a file as its base64 payload.""" + assert isinstance(content, (TextContent, ImageContent, FileContent)) + return _content_block_payload(content) + + +class PayloadCounter(FakeCounter): + """A counter that measures an image or a file by its base64 payload, the way a provider-side counter would.""" + + def count(self, messages: list[ChatMessage], tools: ToolsType | None = None) -> int: + return len(_rendered_conversation(messages, placeholder=_payload_placeholder)) // self.chars_per_token + + +def _step(content: list[TextContent | ImageContent | FileContent], call_id: str) -> list[ChatMessage]: + """An Agent step whose single tool result carries the given content blocks.""" + call = tool_call(call_id) + return [call, ChatMessage.from_tool(tool_result=content, origin=call.tool_calls[0])] + + +INELIGIBLE_CONVERSATION = [ + ChatMessage.from_user("task"), + tool_call("small"), + tool_result("small", call_id="small"), + tool_call("error"), + tool_result("error" * 100, call_id="error", error=True), + tool_call("offloaded"), + ChatMessage.from_tool( + tool_result="offloaded" * 100, + origin=ToolCall(tool_name="search", arguments={}, id="offloaded"), + meta={"tool_result_offloaded": ["stored-result"]}, + ), + tool_call("compacted"), + ChatMessage.from_tool( + tool_result="compacted" * 100, + origin=ToolCall(tool_name="search", arguments={}, id="compacted"), + meta={_COMPACTION_META_KEY: {"strategy": "other"}}, + ), + tool_call("newest"), + tool_result("newest", call_id="newest"), +] + + +class TestToolResultOffloadCompactor: + def test_offloads_older_results(self, tmp_path): + messages = conversation("a" * 400, "newest") + store = FileSystemToolResultStore(root=tmp_path) + compacted = ToolResultOffloadCompactor(store=store, min_tokens=0, preview_chars=5).compact( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + result = compacted[2].tool_call_result + assert result is not None + reference = compacted[2].meta["tool_result_offloaded"][0] + assert result.result == f"Tool result offloaded to text (400 characters) at '{reference}'. Preview: aaaaa..." + assert store.read(reference) == "a" * 400 + assert Path(reference).name.startswith("compacted_search_") + assert compacted[2].meta[_COMPACTION_META_KEY] == { + "strategy": "tool_result_offloading", + "original_tokens": COUNTER.count(messages=[messages[2]]), + } + + def test_repeated_runs_do_not_overwrite_results(self, tmp_path): + def id_less_conversation(result: str) -> list[ChatMessage]: + call = ToolCall(tool_name="search", arguments={}) + return [ + ChatMessage.from_user("task"), + ChatMessage.from_assistant(tool_calls=[call]), + ChatMessage.from_tool(tool_result=result, origin=call), + *conversation("newest")[1:], + ] + + store = FileSystemToolResultStore(root=tmp_path) + compactor = ToolResultOffloadCompactor(store=store, min_tokens=0, preview_chars=0) + first = compactor.compact(messages=id_less_conversation("a" * 400), target_tokens=1, token_counter=COUNTER) + second = compactor.compact(messages=id_less_conversation("b" * 400), target_tokens=1, token_counter=COUNTER) + + assert first is not None + assert second is not None + first_reference = first[2].meta["tool_result_offloaded"][0] + second_reference = second[2].meta["tool_result_offloaded"][0] + assert first_reference != second_reference + assert store.read(first_reference) == "a" * 400 + assert store.read(second_reference) == "b" * 400 + + def test_offloads_binary_results(self, tmp_path): + messages = [ + ChatMessage.from_user("task"), + *_step([ImageContent(base64_image=base64.b64encode(PNG_BYTES).decode(), mime_type="image/png")], "old"), + *conversation("newest")[1:], + ] + store = FileSystemToolResultStore(root=tmp_path) + compacted = ToolResultOffloadCompactor(store=store, min_tokens=0).compact( + messages=messages, target_tokens=1, token_counter=PayloadCounter(chars_per_token=1) + ) + + assert compacted is not None + assert store.read(compacted[2].meta["tool_result_offloaded"][0]) == PNG_BYTES + assert [Path(path).suffix for path in tmp_path.iterdir()] == [".png"] + + def test_skips_ineligible_results(self, tmp_path): + compacted = ToolResultOffloadCompactor(store=FileSystemToolResultStore(root=tmp_path), min_tokens=100).compact( + messages=INELIGIBLE_CONVERSATION, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is None + assert not list(tmp_path.iterdir()) + + @pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"min_keep_steps": 0}, "`min_keep_steps` must be at least 1"), + ({"min_tokens": -1}, "`min_tokens` must be at least 0"), + ({"preview_chars": -1}, "`preview_chars` must be at least 0"), + ], + ) + def test_rejects_invalid_settings(self, tmp_path, kwargs, message): + with pytest.raises(ValueError, match=message): + ToolResultOffloadCompactor(store=FileSystemToolResultStore(root=tmp_path), **kwargs) + + def test_serialization_round_trip(self, tmp_path): + hook = CompactionHook( + compactor=ToolResultOffloadCompactor( + store=FileSystemToolResultStore(root=tmp_path), min_keep_steps=2, min_tokens=12, preview_chars=42 + ), + context_window=10_000, + ) + restored = CompactionHook.from_dict(data=hook.to_dict()).compactor + + assert isinstance(restored, ToolResultOffloadCompactor) + assert isinstance(restored.store, FileSystemToolResultStore) + assert restored.store.root == tmp_path + assert restored.min_keep_steps == 2 + assert restored.min_tokens == 12 + assert restored.preview_chars == 42 + + +class TestToolResultOffloadCompactorAsync: + @pytest.mark.asyncio + async def test_compact_async_matches_compact(self, tmp_path): + messages = conversation("old" * 200, "newest") + store = FileSystemToolResultStore(root=tmp_path) + compacted = await ToolResultOffloadCompactor(store=store, min_tokens=0, preview_chars=0).compact_async( + messages=messages, target_tokens=1, token_counter=COUNTER + ) + + assert compacted is not None + assert store.read(compacted[2].meta["tool_result_offloaded"][0]) == "old" * 200 + assert compacted[3:] == messages[3:] diff --git a/test/hooks/compaction/test_tool_result_pruning.py b/test/hooks/compaction/test_tool_result_pruning.py index 22d6224d7d..d851060267 100644 --- a/test/hooks/compaction/test_tool_result_pruning.py +++ b/test/hooks/compaction/test_tool_result_pruning.py @@ -10,96 +10,16 @@ from haystack.hooks.compaction.tool_result_pruning import _DEFAULT_PLACEHOLDER from haystack.hooks.compaction.utils import _COMPACTION_META_KEY from haystack.hooks.tool_result_offloading import AlwaysOffload, FileSystemToolResultStore, ToolResultOffloadHook from haystack.token_counters import ApproximateTokenCounter -from test.hooks.compaction.helpers import FakeCounter, make_state, tool_call, tool_result +from test.hooks.compaction.helpers import FakeCounter, conversation, make_state, tool_call, tool_result pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") COUNTER = FakeCounter(chars_per_token=1) -def _conversation(*results: str) -> list[ChatMessage]: - messages = [ChatMessage.from_user("task")] - for index, result in enumerate(results): - call_id = f"c{index}" - messages.extend([tool_call(call_id), tool_result(result, call_id=call_id)]) - return messages - - class TestToolResultPruningCompactor: - def test_stops_pruning_after_reaching_target(self): - # Conversation with seven messages: one user message followed by three tool-call/result pairs. - messages = _conversation("a" * 400, "b" * 400, "c" * 400) - compactor = ToolResultPruningCompactor(min_keep_steps=1, min_tokens=0) - # We pre-calculated the target token count to just be enough to only remove the oldest tool result based on the - # count from the FakeCounter. - target_tokens = 1041 - compacted = compactor.compact(messages=messages, target_tokens=target_tokens, token_counter=COUNTER) - assert compacted is not None - assert compacted[2].tool_call_result is not None - assert compacted[2].tool_call_result.result == _DEFAULT_PLACEHOLDER.replace("{tool_name}", "search") - assert compacted[4:] == messages[4:] - # Ensure compaction returns new messages instead of changing the input. - original_results = [messages[index].tool_call_result for index in (2, 4, 6)] - assert all(tool_call_result is not None for tool_call_result in original_results) - assert [tool_call_result.result for tool_call_result in original_results if tool_call_result is not None] == [ - "a" * 400, - "b" * 400, - "c" * 400, - ] - - def test_prunes_multiple_results_in_one_call(self): - messages = _conversation("a" * 400, "b" * 400, "newest") - compacted = ToolResultPruningCompactor(min_keep_steps=1, min_tokens=0).compact( - messages=messages, target_tokens=1, token_counter=COUNTER - ) - assert compacted is not None - for index in (2, 4): - tool_call_result = compacted[index].tool_call_result - assert tool_call_result is not None - assert tool_call_result.result == _DEFAULT_PLACEHOLDER.replace("{tool_name}", "search") - assert compacted[5:] == messages[5:] - - def test_keeps_min_keep_steps(self): - # Three tool-calling steps, with three parallel results in the middle step. - messages = [ - ChatMessage.from_user("task"), - tool_call("old"), - tool_result("old" * 200, call_id="old"), - tool_call("parallel-1", "parallel-2", "parallel-3"), - tool_result("first" * 200, call_id="parallel-1"), - tool_result("second" * 200, call_id="parallel-2"), - tool_result("third" * 200, call_id="parallel-3"), - tool_call("newest"), - tool_result("newest" * 200, call_id="newest"), - ] - # Keeping two steps protects the complete parallel step and the newest step, leaving only `old` eligible. - compacted = ToolResultPruningCompactor(min_keep_steps=2, min_tokens=0).compact( - messages=messages, target_tokens=1, token_counter=COUNTER - ) - assert compacted is not None - assert compacted[2] != messages[2] - assert compacted[3:] == messages[3:] - - def test_returns_none_when_the_conversation_already_fits(self): - messages = _conversation("a" * 400, "b" * 400) - assert ( - ToolResultPruningCompactor(min_keep_steps=1, min_tokens=0).compact( - messages=messages, target_tokens=COUNTER.count(messages=messages), token_counter=COUNTER - ) - is None - ) - - def test_returns_none_when_all_steps_are_protected(self): - messages = _conversation("only result") - assert ( - ToolResultPruningCompactor(min_keep_steps=1, min_tokens=0).compact( - messages=messages, target_tokens=1, token_counter=COUNTER - ) - is None - ) - def test_returns_none_when_the_placeholder_would_not_save_tokens(self): - messages = _conversation("short", "newest") + messages = conversation("short", "newest") compactor = ToolResultPruningCompactor( min_keep_steps=1, min_tokens=0, placeholder="a placeholder much longer than the result" ) @@ -124,7 +44,7 @@ class TestToolResultPruningCompactor: assert compacted[2].tool_call_result.result == _DEFAULT_PLACEHOLDER.replace("{tool_name}", "image_generator") def test_skips_results_below_min_tokens(self): - messages = _conversation("small", "newest") + messages = conversation("small", "newest") compacted = ToolResultPruningCompactor(min_keep_steps=1, min_tokens=200).compact( messages=messages, target_tokens=1, token_counter=COUNTER ) diff --git a/test/hooks/compaction/test_utils.py b/test/hooks/compaction/test_utils.py index 03b685b094..65d7311e32 100644 --- a/test/hooks/compaction/test_utils.py +++ b/test/hooks/compaction/test_utils.py @@ -13,8 +13,9 @@ from haystack.hooks.compaction.utils import ( _historical_turn_spans, _is_compaction_message, _last_assistant_index, + _replace_tool_results_until_target, ) -from test.hooks.compaction.helpers import tool_call, tool_result +from test.hooks.compaction.helpers import FakeCounter, conversation, tool_call, tool_result pytestmark = pytest.mark.filterwarnings("ignore::haystack.utils.experimental.ExperimentalWarning") @@ -172,3 +173,76 @@ class TestCurrentAgentStepGroups: def test_missing_task_anchor(self): messages = [ChatMessage.from_system("rules"), ChatMessage.from_assistant("step")] assert _current_agent_step_groups(messages=messages, system_end=1, task_index=None) == [[1]] + + +COUNTER = FakeCounter(chars_per_token=1) + + +def _shorten(message: ChatMessage, index: int) -> tuple[ChatMessage, int]: + """Replace a tool result with a short marker naming its position.""" + assert message.tool_call_result is not None + shortened = ChatMessage.from_tool(tool_result=f"replaced {index}", origin=message.tool_call_result.origin) + return shortened, COUNTER.count(messages=[message]) - COUNTER.count(messages=[shortened]) + + +def _replaced_positions(messages: list[ChatMessage]) -> list[int]: + return [ + index + for index, message in enumerate(messages) + if message.tool_call_result is not None and message.tool_call_result.result == f"replaced {index}" + ] + + +class TestReplaceToolResultsUntilTarget: + def test_stops_after_reaching_target(self): + messages = conversation("a" * 400, "b" * 400, "newest") + compacted = _replace_tool_results_until_target( + messages=messages, + target_tokens=COUNTER.count(messages=messages) - 1, + token_counter=COUNTER, + min_keep_steps=1, + replace=_shorten, + ) + + assert compacted is not None + assert _replaced_positions(compacted) == [2] + # The caller-owned input list is unchanged. + assert _replaced_positions(messages) == [] + + @pytest.mark.parametrize( + ("min_keep_steps", "expected"), + [pytest.param(1, [2, 4, 5], id="keeps_newest_step"), pytest.param(2, [2], id="keeps_parallel_step_together")], + ) + def test_keeps_min_keep_steps(self, min_keep_steps, expected): + messages = [ + *conversation("old" * 200), + tool_call("parallel-1", "parallel-2"), + tool_result("first" * 200, call_id="parallel-1"), + tool_result("second" * 200, call_id="parallel-2"), + tool_call("newest"), + tool_result("newest" * 200, call_id="newest"), + ] + compacted = _replace_tool_results_until_target( + messages=messages, target_tokens=1, token_counter=COUNTER, min_keep_steps=min_keep_steps, replace=_shorten + ) + + assert compacted is not None + assert _replaced_positions(compacted) == expected + + @pytest.mark.parametrize( + ("messages", "target_tokens", "replace"), + [ + pytest.param(conversation("a" * 400, "newest"), 10_000, _shorten, id="conversation_already_fits"), + pytest.param(conversation("only result"), 1, _shorten, id="every_step_is_protected"), + pytest.param( + conversation("a" * 400, "newest"), 1, lambda message, index: None, id="nothing_is_replaceable" + ), + ], + ) + def test_returns_none(self, messages, target_tokens, replace): + assert ( + _replace_tool_results_until_target( + messages=messages, target_tokens=target_tokens, token_counter=COUNTER, min_keep_steps=1, replace=replace + ) + is None + ) diff --git a/test/hooks/tool_result_offloading/test_hooks.py b/test/hooks/tool_result_offloading/test_hooks.py index 8cc4a5adf7..a7b1dcc5d7 100644 --- a/test/hooks/tool_result_offloading/test_hooks.py +++ b/test/hooks/tool_result_offloading/test_hooks.py @@ -12,7 +12,7 @@ import pytest from haystack.components.agents import Agent from haystack.components.agents.state.state import State from haystack.components.generators.chat import MockChatGenerator -from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCall +from haystack.dataclasses import ChatMessage, ImageContent, ToolCall from haystack.hooks.tool_result_offloading import ( RESULT_STORE_CONTEXT_KEY, AlwaysOffload, @@ -41,11 +41,10 @@ def _tool_message(tool_name: str, result: str, *, error: bool = False, call_id: ) -# A 1x1 PNG and a minimal PDF. +# A 1x1 PNG. PNG_BYTES = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" ) -PDF_BYTES = b"%PDF-1.4\n1 0 obj\n<< /Type /Catalog >>\nendobj\n\xde\xad\xbe\xef\n%%EOF\n" class TextOnlyToolResultStore(ToolResultStore): @@ -68,12 +67,6 @@ def _image_block() -> ImageContent: return ImageContent(base64_image=base64.b64encode(PNG_BYTES).decode("utf-8"), mime_type="image/png") -def _file_block(filename: str | None = "report.pdf") -> FileContent: - return FileContent( - base64_data=base64.b64encode(PDF_BYTES).decode("utf-8"), mime_type="application/pdf", filename=filename - ) - - class TestToolResultOffloadHookRouting: def test_exact_tuple_and_wildcard_keys(self, tmp_path): hook = ToolResultOffloadHook( @@ -130,50 +123,6 @@ class TestToolResultOffloadHookBehavior: hook.run(state) assert state.data["messages"][0].tool_call_result.result == "boom" - def test_mixed_result_offloads_every_block_to_its_own_entry(self, tmp_path): - store = FileSystemToolResultStore(root=tmp_path) - hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}, preview_chars=4) - content: list[TextContent | ImageContent | FileContent] = [ - TextContent("caption"), - _image_block(), - _file_block(), - ] - message = ChatMessage.from_tool(tool_result=content, origin=ToolCall(tool_name="a", arguments={}, id="1")) - state = _state_with_messages([message]) - hook.run(state) - offloaded = state.data["messages"][0] - references = offloaded.meta["tool_result_offloaded"] - assert len(references) == 3 - assert [Path(reference).name for reference in references] == ["1_a_1_0.txt", "1_a_1_1.png", "1_a_1_2.pdf"] - assert store.read(references[0]) == "caption" - assert store.read(references[1]) == PNG_BYTES - assert store.read(references[2]) == PDF_BYTES - pointer = offloaded.tool_call_result.result - assert pointer.startswith("Tool result offloaded to 3 files:") - assert f"1. text (7 characters) at '{references[0]}'. Preview: capt..." in pointer - assert f"2. image/png ({len(PNG_BYTES)} bytes) at '{references[1]}'" in pointer - assert f"3. application/pdf named 'report.pdf' ({len(PDF_BYTES)} bytes) at '{references[2]}'" in pointer - - @pytest.mark.parametrize( - "block, expected_extension", - [ - pytest.param(FileContent(base64_data="aGk=", mime_type="text/csv", filename=None), ".csv", id="mime_type"), - pytest.param( - FileContent(base64_data="aGk=", mime_type="application/pdf", filename="notes.md"), ".md", id="filename" - ), - pytest.param(FileContent(base64_data="aGk=", mime_type=None, filename=None), ".bin", id="fallback"), - ], - ) - def test_binary_block_extension(self, tmp_path, block, expected_extension): - hook = ToolResultOffloadHook( - store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} - ) - message = ChatMessage.from_tool(tool_result=[block], origin=ToolCall(tool_name="a", arguments={}, id="1")) - state = _state_with_messages([message]) - hook.run(state) - reference = state.data["messages"][0].meta["tool_result_offloaded"][0] - assert Path(reference).suffix == expected_extension - def test_policy_sizes_result_by_its_base64_payload(self, tmp_path): # The base64 payload is what occupies the context window, so a threshold below it must trigger an offload # while one above it must not. @@ -195,33 +144,6 @@ class TestToolResultOffloadHookBehavior: assert over_state.data["messages"][0].tool_call_result.result.startswith("Tool result offloaded") assert under_state.data["messages"][0].tool_call_result.result == [image] - @pytest.mark.parametrize( - "empty_result", - ["", [], [TextContent(text="")], [TextContent(text=""), TextContent(text="")]], - ids=["empty_string", "empty_list", "single_empty_text_block", "several_empty_text_blocks"], - ) - def test_empty_result_is_not_offloaded(self, tmp_path, empty_result): - hook = ToolResultOffloadHook( - store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} - ) - message = ChatMessage.from_tool(tool_result=empty_result, 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 == empty_result - assert not list(Path(tmp_path).iterdir()) - - def test_empty_text_alongside_image_is_offloaded(self, tmp_path): - hook = ToolResultOffloadHook( - store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} - ) - image = ImageContent(base64_image=base64.b64encode(b"PNGDATA").decode(), mime_type="image/png") - message = ChatMessage.from_tool( - tool_result=[TextContent(text=""), image], 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.startswith("Tool result offloaded") - def test_id_less_parallel_calls_do_not_collide(self, tmp_path): store = FileSystemToolResultStore(root=tmp_path) hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) @@ -262,17 +184,6 @@ class TestToolResultOffloadHookBehavior: 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"][0] - pointer = message.tool_call_result.result - assert pointer == f"Tool result offloaded to text (8 characters) at '{reference}'. Preview: ABCDE..." - 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. @@ -315,15 +226,6 @@ class TestToolResultOffloadHookWithTextOnlyStore: assert store.read(offloaded.meta["tool_result_offloaded"][0]) == "A" * 50 assert not caplog.records - def test_image_and_file_results_stay_in_context(self, caplog): - hook = ToolResultOffloadHook(store=TextOnlyToolResultStore(), offload_strategies={"*": AlwaysOffload()}) - content: list[TextContent | ImageContent | FileContent] = [TextContent("caption"), _file_block()] - 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 "does not support binary content" in caplog.text - class TestToolResultOffloadHookSerde: def test_to_dict_from_dict_roundtrip(self, tmp_path): diff --git a/test/hooks/tool_result_offloading/test_utils.py b/test/hooks/tool_result_offloading/test_utils.py new file mode 100644 index 0000000000..ddfbfcd68e --- /dev/null +++ b/test/hooks/tool_result_offloading/test_utils.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import base64 +from pathlib import Path + +import pytest + +from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCall +from haystack.dataclasses.chat_message import ToolCallResultContentT +from haystack.hooks.tool_result_offloading import FileSystemToolResultStore, ToolResultStore +from haystack.hooks.tool_result_offloading.utils import _offloadable_content_blocks, _offloaded_message + +# A 1x1 PNG and a minimal PDF. +PNG_BYTES = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) +PDF_BYTES = b"%PDF-1.4\n1 0 obj\n<< /Type /Catalog >>\nendobj\n\xde\xad\xbe\xef\n%%EOF\n" +IMAGE_BLOCK = ImageContent(base64_image=base64.b64encode(PNG_BYTES).decode("utf-8"), mime_type="image/png") +FILE_BLOCK = FileContent( + base64_data=base64.b64encode(PDF_BYTES).decode("utf-8"), mime_type="application/pdf", filename="report.pdf" +) + + +class TextOnlyToolResultStore(ToolResultStore): + """A store that only holds text, leaving `supports_binary_content` at its False default.""" + + def __init__(self) -> None: + self.data: dict[str, str] = {} + + def write(self, *, key: str, content: str | bytes) -> str: + if isinstance(content, bytes): + raise TypeError("Binary content not supported") + self.data[key] = content + return key + + def read(self, reference: str) -> str: + return self.data[reference] + + +def _tool_message(result: ToolCallResultContentT, *, error: bool = False) -> ChatMessage: + return ChatMessage.from_tool( + tool_result=result, origin=ToolCall(tool_name="a", arguments={}, id="1"), error=error, meta={"source": "test"} + ) + + +class TestOffloadableContentBlocks: + @pytest.mark.parametrize( + ("result", "expected"), + [ + pytest.param("", None, id="empty_string"), + pytest.param([], None, id="empty_list"), + pytest.param([TextContent(text="")], None, id="single_empty_text_block"), + pytest.param([TextContent(text=""), TextContent(text="")], None, id="several_empty_text_blocks"), + pytest.param("text", [TextContent(text="text")], id="string"), + pytest.param( + [TextContent(text=""), IMAGE_BLOCK], [TextContent(text=""), IMAGE_BLOCK], id="empty_text_with_image" + ), + ], + ) + def test_offloadable_content_blocks(self, tmp_path, result, expected): + tool_result = _tool_message(result).tool_call_result + assert tool_result is not None + + assert _offloadable_content_blocks(result=tool_result, store=FileSystemToolResultStore(root=tmp_path)) == ( + expected + ) + + def test_text_only_store_rejects_binary_content(self, caplog): + tool_result = _tool_message([TextContent("caption"), FILE_BLOCK]).tool_call_result + assert tool_result is not None + + assert _offloadable_content_blocks(result=tool_result, store=TextOnlyToolResultStore()) is None + assert "does not support binary content" in caplog.text + + +class TestOffloadedMessage: + def test_text_result(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + offloaded = _offloaded_message( + message=_tool_message("ABCDEFGH", error=True), + content_blocks=[TextContent("ABCDEFGH")], + store=store, + key_prefix="key", + preview_chars=5, + additional_meta={"extra": 1}, + ) + + reference = offloaded.meta["tool_result_offloaded"][0] + assert Path(reference).name == "key.txt" + assert store.read(reference) == "ABCDEFGH" + assert offloaded.tool_call_result is not None + assert offloaded.tool_call_result.result == ( + f"Tool result offloaded to text (8 characters) at '{reference}'. Preview: ABCDE..." + ) + assert offloaded.tool_call_result.origin == ToolCall(tool_name="a", arguments={}, id="1") + assert offloaded.tool_call_result.error is True + assert offloaded.meta == {"source": "test", "extra": 1, "tool_result_offloaded": [reference]} + + def test_writes_every_block_to_its_own_entry(self, tmp_path): + store = FileSystemToolResultStore(root=tmp_path) + content: list[TextContent | ImageContent | FileContent] = [TextContent("caption"), IMAGE_BLOCK, FILE_BLOCK] + offloaded = _offloaded_message( + message=_tool_message(content), content_blocks=content, store=store, key_prefix="key", preview_chars=4 + ) + + references = offloaded.meta["tool_result_offloaded"] + assert [Path(reference).name for reference in references] == ["key_0.txt", "key_1.png", "key_2.pdf"] + assert store.read(references[0]) == "caption" + assert store.read(references[1]) == PNG_BYTES + assert store.read(references[2]) == PDF_BYTES + assert offloaded.tool_call_result is not None + assert offloaded.tool_call_result.result == "\n".join( + [ + "Tool result offloaded to 3 files:", + f"1. text (7 characters) at '{references[0]}'. Preview: capt...", + f"2. image/png ({len(PNG_BYTES)} bytes) at '{references[1]}'", + f"3. application/pdf named 'report.pdf' ({len(PDF_BYTES)} bytes) at '{references[2]}'", + ] + ) + + @pytest.mark.parametrize( + ("block", "expected_extension"), + [ + pytest.param(FileContent(base64_data="aGk=", mime_type="text/csv", filename=None), ".csv", id="mime_type"), + pytest.param( + FileContent(base64_data="aGk=", mime_type="application/pdf", filename="notes.md"), ".md", id="filename" + ), + pytest.param(FileContent(base64_data="aGk=", mime_type=None, filename=None), ".bin", id="fallback"), + ], + ) + def test_binary_block_extension(self, tmp_path, block, expected_extension): + offloaded = _offloaded_message( + message=_tool_message([block]), + content_blocks=[block], + store=FileSystemToolResultStore(root=tmp_path), + key_prefix="key", + preview_chars=0, + ) + + assert Path(offloaded.meta["tool_result_offloaded"][0]).suffix == expected_extension + + def test_rejects_a_message_without_a_tool_result(self, tmp_path): + with pytest.raises(ValueError, match="Only tool-result messages can be offloaded."): + _offloaded_message( + message=ChatMessage.from_user("hi"), + content_blocks=[TextContent("hi")], + store=FileSystemToolResultStore(root=tmp_path), + key_prefix="key", + preview_chars=0, + )