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 Where a `ToolResultOffloadHook` or `ToolResultOffloadCompactor` writes tool results and a read tool fetches them. 

14 

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

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

17 result; `read` resolves that reference back to the original content. Only the store interprets a reference - 

18 callers pass it back to `read` unchanged. 

19 

20 A store that sets `supports_binary_content` takes bytes in `write` and gives them back from `read`. One that 

21 leaves it False is only ever given text, and image and file results stay in the conversation instead. 

22 

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

24 cover stores whose constructor takes no arguments. 

25 """ 

26 

27 # Whether the store handles binary content: `write` accepting bytes and `read` returning them. When False, the 

28 # store is only ever given text, and the image and file content of a tool result stays in the conversation. 

29 supports_binary_content: bool = False 

30 

31 def write(self, *, key: str, content: str | bytes) -> str: 

32 """ 

33 Persist `content` under `key` and return a reference to it. 

34 

35 :param key: A stable, per-result identifier derived from the tool call (e.g. a file name). It carries an 

36 extension matching the content, so a store that maps keys to files can use it as-is. 

37 :param content: The tool result to persist. Text arrives as a string. Image and file content arrives as the 

38 decoded bytes of its base64 payload, and only when the store sets `supports_binary_content` to True - a 

39 text-only store may narrow this parameter to `str`. 

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

41 """ 

42 ... 

43 

44 def read(self, reference: str) -> str | bytes: 

45 """ 

46 Return the content previously stored under `reference`. 

47 

48 :param reference: A reference string returned by `write`. 

49 :returns: The stored content: a string for content written as text, bytes for binary content such as an 

50 offloaded image or file. A store that does not support binary content only ever returns a string. 

51 """ 

52 ... 

53 

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

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

56 return default_to_dict(self) 

57 

58 @classmethod 

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

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

61 return default_from_dict(cls, data) 

62 

63 

64class OffloadPolicy(Protocol): 

65 """ 

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

67 

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

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

70 

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

72 cover policies whose constructor takes no arguments. 

73 """ 

74 

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

76 """ 

77 Return whether the given tool result should be offloaded. 

78 

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

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

81 result carrying image or file blocks, this is the text and base64 payloads of all its blocks joined 

82 together, so its length reflects the context the result actually occupies. 

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

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

85 """ 

86 ... 

87 

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

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

90 return default_to_dict(self) 

91 

92 @classmethod 

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

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

95 return default_from_dict(cls, data)