Coverage for haystack/tools/serde_utils.py: 93%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import TYPE_CHECKING, Any 

6 

7from haystack.core.errors import DeserializationError 

8from haystack.core.serialization import import_class_by_name 

9from haystack.tools.tool import Tool 

10from haystack.tools.toolset import Toolset 

11 

12if TYPE_CHECKING: 

13 from haystack.tools import ToolsType 

14 

15 

16def serialize_tools_or_toolset(tools: "ToolsType | None") -> dict[str, Any] | list[dict[str, Any]] | None: 

17 """ 

18 Serialize tools or toolsets to dictionaries. 

19 

20 :param tools: A Toolset, a list of Tools and/or Toolsets, or None 

21 :returns: Serialized representation preserving Tool/Toolset boundaries when provided 

22 """ 

23 if tools is None: 

24 return None 

25 if isinstance(tools, Toolset): 

26 return tools.to_dict() 

27 if isinstance(tools, list): 

28 serialized: list[dict[str, Any]] = [] 

29 for item in tools: 

30 if isinstance(item, (Toolset, Tool)): 

31 serialized.append(item.to_dict()) 

32 else: 

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

34 return serialized 

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

36 

37 

38def deserialize_tools_or_toolset_inplace(data: dict[str, Any], key: str = "tools") -> None: 

39 """ 

40 Deserialize a list of Tools and/or Toolsets, or a single Toolset in a dictionary inplace. 

41 

42 :param data: 

43 The dictionary with the serialized data. 

44 :param key: 

45 The key in the dictionary where the list of Tools and/or Toolsets, or single Toolset is stored. 

46 """ 

47 if key in data: 

48 serialized_tools = data[key] 

49 

50 if serialized_tools is None: 

51 return 

52 

53 # Check if it's a serialized Toolset (a dict with "type" and "data" keys) 

54 if isinstance(serialized_tools, dict) and all(k in serialized_tools for k in ["type", "data"]): 

55 toolset_class_name = serialized_tools.get("type") 

56 if not toolset_class_name: 

57 raise DeserializationError("The 'type' key is missing or None in the serialized toolset data") 

58 

59 toolset_class = import_class_by_name(toolset_class_name) 

60 

61 if not issubclass(toolset_class, Toolset): 

62 raise TypeError(f"Class '{toolset_class}' is not a subclass of Toolset") 

63 

64 data[key] = toolset_class.from_dict(serialized_tools) 

65 return 

66 

67 if not isinstance(serialized_tools, list): 

68 raise TypeError(f"The value of '{key}' is not a list or a dictionary") 

69 

70 deserialized_tools: list[Tool | Toolset] = [] 

71 for tool in serialized_tools: 

72 if not isinstance(tool, dict): 

73 raise TypeError(f"Serialized tool '{tool}' is not a dictionary") 

74 

75 # different classes are allowed: Tool, ComponentTool, Toolset, etc. 

76 tool_class = import_class_by_name(tool["type"]) 

77 if issubclass(tool_class, (Tool, Toolset)): 

78 deserialized_tools.append(tool_class.from_dict(tool)) 

79 else: 

80 raise TypeError(f"Class '{tool_class}' is neither Tool nor Toolset") 

81 

82 data[key] = deserialized_tools