Coverage for haystack/hooks/tool_result_offloading/types/protocol.py: 89%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any, Protocol 

6 

7from haystack.components.agents.state.state import State 

8from haystack.core.serialization import default_from_dict, default_to_dict 

9 

10 

11class ToolResultStore(Protocol): 

12 """ 

13 A place a `ToolResultOffloadHook` writes offloaded tool results to, and reads them back from. 

14 

15 Implementations decide where and how the content lives (local disk, an isolated sandbox filesystem, object 

16 storage, ...). `write` returns an opaque reference string that the Agent puts in the conversation in place of the 

17 full result; `read` resolves that reference back to the original content. 

18 

19 Implement both `to_dict` and `from_dict` to make a custom store serializable; the default implementations below 

20 cover stores whose constructor takes no arguments. 

21 """ 

22 

23 def write(self, *, key: str, content: str) -> str: 

24 """ 

25 Persist `content` under `key` and return an opaque reference to it. 

26 

27 :param key: A stable, per-result identifier the hook derives from the tool call (e.g. a file name). 

28 :param content: The tool result to persist. 

29 :returns: A reference string (e.g. a path or URI) that `read` can later resolve. 

30 """ 

31 ... 

32 

33 def read(self, reference: str) -> str: 

34 """Return the content previously stored under `reference`.""" 

35 ... 

36 

37 def to_dict(self) -> dict[str, Any]: 

38 """Serialize the store to a dictionary.""" 

39 return default_to_dict(self) 

40 

41 @classmethod 

42 def from_dict(cls, data: dict[str, Any]) -> "ToolResultStore": 

43 """Deserialize the store from a dictionary.""" 

44 return default_from_dict(cls, data) 

45 

46 

47class OffloadPolicy(Protocol): 

48 """ 

49 Decides, per tool result, whether the `ToolResultOffloadHook` offloads it to the store or leaves it in context. 

50 

51 A `ToolResultOffloadHook` maps tool names to policies, so different tools can offload under different conditions 

52 (always, never, or a custom rule such as a size threshold). 

53 

54 Implement both `to_dict` and `from_dict` to make a custom policy serializable; the default implementations below 

55 cover policies whose constructor takes no arguments. 

56 """ 

57 

58 def should_offload(self, tool_name: str, result: str, state: State) -> bool: 

59 """ 

60 Return whether the given tool result should be offloaded. 

61 

62 :param tool_name: The name of the tool that produced the result. 

63 :param result: The tool result as a string (the content that would otherwise stay in the conversation). 

64 :param state: The Agent's live `State`, for policies that decide based on run context. 

65 :returns: True to offload the result to the store, False to leave it in context. 

66 """ 

67 ... 

68 

69 def to_dict(self) -> dict[str, Any]: 

70 """Serialize the policy to a dictionary.""" 

71 return default_to_dict(self) 

72 

73 @classmethod 

74 def from_dict(cls, data: dict[str, Any]) -> "OffloadPolicy": 

75 """Deserialize the policy from a dictionary.""" 

76 return default_from_dict(cls, data)