Coverage for haystack/tools/utils.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from collections.abc import Sequence 

6from typing import TYPE_CHECKING 

7 

8from haystack.tools.tool import Tool 

9from haystack.tools.toolset import Toolset 

10 

11if TYPE_CHECKING: 

12 from haystack.tools import ToolsType 

13 

14 

15def _as_tool_sequence(tools: "ToolsType | None") -> Sequence[Tool | Toolset]: 

16 if tools is None: 

17 return [] 

18 if isinstance(tools, Toolset): 

19 return [tools] 

20 return tools 

21 

22 

23def warm_up_tools(tools: "ToolsType | None" = None) -> None: 

24 """ 

25 Warm up tools from various formats (Tools, Toolsets, or mixed lists). 

26 

27 For Toolset objects, this delegates to warm_up() if implemented; otherwise, 

28 it calls warm_up() on each tool that implements it. Toolset subclasses can implement warm_up() 

29 to customize initialization behavior (e.g., setting up shared resources). 

30 

31 :param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None. 

32 """ 

33 for item in _as_tool_sequence(tools): 

34 if hasattr(item, "warm_up"): 

35 item.warm_up() 

36 elif isinstance(item, Toolset): 

37 warm_up_tools(tools=item.tools) 

38 

39 

40async def warm_up_tools_async(tools: "ToolsType | None" = None) -> None: 

41 """ 

42 Warm up tools asynchronously from various formats (Tools, Toolsets, or mixed lists). 

43 

44 Call warm_up_async() if implemented, otherwise warm_up() if available. 

45 If a Toolset implements neither method, this function warms up its tools. 

46 

47 :param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None. 

48 """ 

49 for item in _as_tool_sequence(tools): 

50 if hasattr(item, "warm_up_async"): 

51 await item.warm_up_async() 

52 elif hasattr(item, "warm_up"): 

53 item.warm_up() 

54 elif isinstance(item, Toolset): 

55 await warm_up_tools_async(tools=item.tools) 

56 

57 

58def close_tools(tools: "ToolsType | None" = None) -> None: 

59 """ 

60 Close tools from various formats (Tools, Toolsets, or mixed lists). 

61 

62 For Toolset objects, this delegates to close() if implemented; otherwise, 

63 it calls close() on each tool that implements it. 

64 

65 :param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None. 

66 """ 

67 for item in _as_tool_sequence(tools): 

68 if hasattr(item, "close"): 

69 item.close() 

70 elif isinstance(item, Toolset): 

71 close_tools(tools=item.tools) 

72 

73 

74async def close_tools_async(tools: "ToolsType | None" = None) -> None: 

75 """ 

76 Close tools asynchronously from various formats (Tools, Toolsets, or mixed lists). 

77 

78 Call close_async() if implemented, otherwise close() if available. 

79 If a Toolset implements neither method, this function closes its tools. 

80 

81 :param tools: A sequence of Tool and/or Toolset objects, a single Toolset, or None. 

82 """ 

83 for item in _as_tool_sequence(tools): 

84 if hasattr(item, "close_async"): 

85 await item.close_async() 

86 elif hasattr(item, "close"): 

87 item.close() 

88 elif isinstance(item, Toolset): 

89 await close_tools_async(tools=item.tools) 

90 

91 

92def flatten_tools_or_toolsets(tools: "ToolsType | None") -> list[Tool]: 

93 """ 

94 Flatten tools from various formats into a list of Tool instances. 

95 

96 :param tools: Tools in list[Union[Tool, Toolset]], Toolset, or None format. 

97 :returns: A flat list of Tool instances. 

98 """ 

99 if tools is None: 

100 return [] 

101 

102 if isinstance(tools, Toolset): 

103 return list(tools) 

104 

105 if isinstance(tools, list): 

106 flattened: list[Tool] = [] 

107 for item in tools: 

108 if isinstance(item, Toolset): 

109 flattened.extend(list(item)) 

110 elif isinstance(item, Tool): 

111 flattened.append(item) 

112 else: 

113 raise TypeError("Items in the tools list must be Tool or Toolset instances.") 

114 return flattened 

115 

116 raise TypeError("tools must be list[Union[Tool, Toolset]], Toolset, or None")