Coverage for haystack/hooks/tool_result_offloading/policies.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

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

8from haystack.core.serialization import default_to_dict 

9from haystack.hooks.tool_result_offloading.types import OffloadPolicy 

10 

11 

12class AlwaysOffload(OffloadPolicy): 

13 """Offload every result of the tool it is assigned to.""" 

14 

15 def should_offload(self, tool_name: str, result: str, state: State) -> bool: # noqa: ARG002 

16 """ 

17 Decide whether to offload the given tool result. 

18 

19 :param tool_name: The name of the tool that produced the result (unused; this policy always offloads). 

20 :param result: The tool result string (unused; this policy always offloads). 

21 :param state: The Agent's live `State` (unused; this policy always offloads). 

22 :returns: Always True. 

23 """ 

24 return True 

25 

26 

27class NeverOffload(OffloadPolicy): 

28 """Never offload; keep the tool's full result in context. Use to opt a tool out of a wildcard default.""" 

29 

30 def should_offload(self, tool_name: str, result: str, state: State) -> bool: # noqa: ARG002 

31 """ 

32 Decide whether to offload the given tool result. 

33 

34 :param tool_name: The name of the tool that produced the result (unused; this policy never offloads). 

35 :param result: The tool result string (unused; this policy never offloads). 

36 :param state: The Agent's live `State` (unused; this policy never offloads). 

37 :returns: Always False. 

38 """ 

39 return False 

40 

41 

42class OffloadOverChars(OffloadPolicy): 

43 """Offload a result only when its string length exceeds `threshold` characters.""" 

44 

45 def __init__(self, threshold: int) -> None: 

46 """ 

47 Initialize the policy with its character threshold. 

48 

49 :param threshold: Offload the result when its length in characters is strictly greater than this value. 

50 """ 

51 self.threshold = threshold 

52 

53 def should_offload(self, tool_name: str, result: str, state: State) -> bool: # noqa: ARG002 

54 """ 

55 Decide whether to offload the given tool result based on its length. 

56 

57 :param tool_name: The name of the tool that produced the result (unused; only length is considered). 

58 :param result: The tool result string whose length is compared against the threshold. 

59 :param state: The Agent's live `State` (unused; only length is considered). 

60 :returns: True when `result` is longer than `threshold` characters, otherwise False. 

61 """ 

62 return len(result) > self.threshold 

63 

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

65 """ 

66 Serialize the policy, including its threshold. 

67 

68 :returns: A dictionary representation of the policy. 

69 """ 

70 return default_to_dict(self, threshold=self.threshold)