# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 from pathlib import Path from typing import Annotated, Any from unittest.mock import AsyncMock, MagicMock import pytest from haystack import component from haystack.components.agents import Agent from haystack.components.agents.state.state import State from haystack.dataclasses import ChatMessage, ImageContent, TextContent, ToolCall from haystack.hooks.tool_result_offloading import ( RESULT_STORE_CONTEXT_KEY, AlwaysOffload, FileSystemToolResultStore, NeverOffload, OffloadOverChars, ToolResultOffloadHook, ) from haystack.tools import Tool, Toolset, tool @component class MockChatGenerator: @component.output_types(replies=list[ChatMessage]) def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]: return {"replies": [ChatMessage.from_assistant("done")]} @component.output_types(replies=list[ChatMessage]) async def run_async( self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs ) -> dict[str, Any]: return {"replies": [ChatMessage.from_assistant("done")]} @tool def big_tool(query: Annotated[str, "the query"]) -> str: """Return a large result.""" return "R" * 500 def _state_with_messages(messages: list[ChatMessage]) -> State: return State(schema={"messages": {"type": list[ChatMessage]}}, data={"messages": messages, "step_count": 1}) def _tool_message(tool_name: str, result: str, *, error: bool = False, call_id: str = "c1") -> ChatMessage: return ChatMessage.from_tool( tool_result=result, origin=ToolCall(tool_name=tool_name, arguments={}, id=call_id), error=error ) class TestToolResultOffloadHookRouting: def test_exact_tuple_and_wildcard_keys(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={ "a": AlwaysOffload(), ("b", "c"): AlwaysOffload(), "d": NeverOffload(), "*": AlwaysOffload(), }, ) state = _state_with_messages( [ _tool_message("a", "A" * 50, call_id="1"), # exact -> offload _tool_message("b", "B" * 50, call_id="2"), # tuple -> offload _tool_message("d", "D" * 50, call_id="3"), # exact NeverOffload -> keep _tool_message("z", "Z" * 50, call_id="4"), # wildcard -> offload ] ) hook.run(state) results = [m.tool_call_result.result for m in state.data["messages"]] assert results[0].startswith("Tool result offloaded") assert results[1].startswith("Tool result offloaded") assert results[2] == "D" * 50 assert results[3].startswith("Tool result offloaded") def test_tool_without_matching_key_is_not_offloaded(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"a": AlwaysOffload()} ) state = _state_with_messages([_tool_message("b", "B" * 50)]) hook.run(state) assert state.data["messages"][0].tool_call_result.result == "B" * 50 def test_over_chars_threshold(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": OffloadOverChars(10)} ) state = _state_with_messages( [_tool_message("a", "x" * 10, call_id="1"), _tool_message("a", "x" * 11, call_id="2")] ) hook.run(state) results = [m.tool_call_result.result for m in state.data["messages"]] assert results[0] == "x" * 10 assert results[1].startswith("Tool result offloaded") class TestToolResultOffloadHookBehavior: def test_error_results_are_never_offloaded(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} ) state = _state_with_messages([_tool_message("a", "boom", error=True)]) hook.run(state) assert state.data["messages"][0].tool_call_result.result == "boom" def test_offloads_sequence_of_text_content(self, tmp_path): store = FileSystemToolResultStore(root=tmp_path) hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) message = ChatMessage.from_tool( tool_result=[TextContent("A" * 30), TextContent("B" * 30)], origin=ToolCall(tool_name="a", arguments={}, id="1"), ) state = _state_with_messages([message]) hook.run(state) offloaded = state.data["messages"][0] assert offloaded.tool_call_result.result.startswith("Tool result offloaded") assert "60 characters" in offloaded.tool_call_result.result assert store.read(offloaded.meta["tool_result_offloaded"]) == "A" * 30 + "B" * 30 def test_result_with_image_content_is_not_offloaded(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} ) content = [TextContent("caption"), ImageContent(base64_image="aGVsbG8=", mime_type="image/png")] message = ChatMessage.from_tool(tool_result=content, origin=ToolCall(tool_name="a", arguments={}, id="1")) state = _state_with_messages([message]) hook.run(state) assert state.data["messages"][0].tool_call_result.result == content assert not list(Path(tmp_path).iterdir()) def test_id_less_parallel_calls_do_not_collide(self, tmp_path): store = FileSystemToolResultStore(root=tmp_path) hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) first = ChatMessage.from_tool(tool_result="FIRST" * 20, origin=ToolCall(tool_name="a", arguments={}, id=None)) second = ChatMessage.from_tool(tool_result="SECOND" * 20, origin=ToolCall(tool_name="a", arguments={}, id=None)) state = _state_with_messages([first, second]) hook.run(state) refs = [m.meta["tool_result_offloaded"] for m in state.data["messages"]] assert refs[0] != refs[1] assert store.read(refs[0]) == "FIRST" * 20 assert store.read(refs[1]) == "SECOND" * 20 assert len(list(Path(tmp_path).iterdir())) == 2 def test_only_trailing_tool_results_are_offloaded(self, tmp_path): store = FileSystemToolResultStore(root=tmp_path) hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) # A tool result from a prior turn, then an assistant message, then this step's fresh tool result. history = _tool_message("old", "H" * 50, call_id="old1") assistant = ChatMessage.from_assistant(tool_calls=[ToolCall("a", {}, id="c1")]) fresh = _tool_message("a", "F" * 50, call_id="c1") state = _state_with_messages([history, assistant, fresh]) hook.run(state) out = state.data["messages"] assert out[0].tool_call_result.result == "H" * 50 assert out[2].tool_call_result.result.startswith("Tool result offloaded") assert len(list(Path(tmp_path).iterdir())) == 1 def test_second_offload_hook_does_not_reoffload_pointer(self, tmp_path): # Two offload hooks under `after_tool` run in sequence on the same state. The first offloads the result and # marks it; the `_OFFLOADED_META_KEY` marker stops the second from offloading the pointer text again. store = FileSystemToolResultStore(root=tmp_path) first_hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) second_hook = ToolResultOffloadHook(store=store, offload_strategies={"*": AlwaysOffload()}) state = _state_with_messages([_tool_message("a", "A" * 50)]) first_hook.run(state) pointer = state.data["messages"][0].tool_call_result.result second_hook.run(state) assert state.data["messages"][0].tool_call_result.result == pointer assert len(list(Path(tmp_path).iterdir())) == 1 def test_pointer_contains_reference_size_and_preview(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()}, preview_chars=5 ) state = _state_with_messages([_tool_message("a", "ABCDEFGH")]) hook.run(state) message = state.data["messages"][0] reference = message.meta["tool_result_offloaded"] pointer = message.tool_call_result.result assert reference in pointer assert "8 characters" in pointer assert "ABCDE..." in pointer def test_concurrent_runs_are_isolated_by_per_run_store(self, tmp_path): # One shared hook instance, two runs each supplying its own store via hook_context. Even with identical store # keys (same tool, id and step), each run writes to and reads from its own store — no cross-run collision. hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path / "shared"), offload_strategies={"*": AlwaysOffload()} ) store_a = FileSystemToolResultStore(root=tmp_path / "a") store_b = FileSystemToolResultStore(root=tmp_path / "b") state_a = _state_with_messages([_tool_message("t", "AAA" * 20, call_id="x")]) state_a.data["hook_context"] = {RESULT_STORE_CONTEXT_KEY: store_a} state_b = _state_with_messages([_tool_message("t", "BBB" * 20, call_id="x")]) state_b.data["hook_context"] = {RESULT_STORE_CONTEXT_KEY: store_b} hook.run(state_a) hook.run(state_b) ref_a = state_a.data["messages"][0].meta["tool_result_offloaded"] ref_b = state_b.data["messages"][0].meta["tool_result_offloaded"] assert store_a.read(ref_a) == "AAA" * 20 assert store_b.read(ref_b) == "BBB" * 20 assert not (tmp_path / "shared").exists() def test_hook_context_store_overrides_constructor_store(self, tmp_path): default_store = FileSystemToolResultStore(root=tmp_path / "default") request_store = FileSystemToolResultStore(root=tmp_path / "request") hook = ToolResultOffloadHook(store=default_store, offload_strategies={"*": AlwaysOffload()}) state = _state_with_messages([_tool_message("a", "A" * 50)]) state.data["hook_context"] = {RESULT_STORE_CONTEXT_KEY: request_store} hook.run(state) assert (tmp_path / "request").exists() assert not (tmp_path / "default").exists() class TestToolResultOffloadHookSerde: def test_to_dict_from_dict_roundtrip(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"a": AlwaysOffload(), ("b", "c"): OffloadOverChars(100), "*": NeverOffload()}, preview_chars=42, ) restored = ToolResultOffloadHook.from_dict(hook.to_dict()) assert restored.preview_chars == 42 assert isinstance(restored.store, FileSystemToolResultStore) assert set(restored.offload_strategies) == {"a", ("b", "c"), "*"} assert isinstance(restored.offload_strategies[("b", "c")], OffloadOverChars) assert restored.offload_strategies[("b", "c")].threshold == 100 class TestToolResultOffloadHookInAgent: def test_offloads_tool_result_seen_by_next_llm_call(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} ) agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]}) agent.warm_up() agent.chat_generator.run = MagicMock( side_effect=[ {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("big_tool", {"query": "x"})])]}, {"replies": [ChatMessage.from_assistant("done")]}, ] ) agent.run(messages=[ChatMessage.from_user("hi")]) second_call_messages = agent.chat_generator.run.call_args_list[1].kwargs["messages"] offloaded = [m for m in second_call_messages if m.tool_call_result is not None] assert offloaded[0].tool_call_result.result.startswith("Tool result offloaded") assert len(list(Path(tmp_path).iterdir())) == 1 class TestToolResultOffloadHookInAgentAsync: @pytest.mark.asyncio async def test_offloads_tool_result_async(self, tmp_path): hook = ToolResultOffloadHook( store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()} ) agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]}) agent.warm_up() agent.chat_generator.run_async = AsyncMock( side_effect=[ {"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("big_tool", {"query": "x"})])]}, {"replies": [ChatMessage.from_assistant("done")]}, ] ) result = await agent.run_async(messages=[ChatMessage.from_user("hi")]) offloaded = [m for m in result["messages"] if m.tool_call_result is not None] assert offloaded[0].tool_call_result.result.startswith("Tool result offloaded")