Coverage for haystack/hooks/utils.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.core.serialization import component_to_dict 

8from haystack.hooks.protocol import Hook, HookPoint 

9from haystack.utils.deserialization import deserialize_component_inplace 

10 

11 

12# Hooks are (de)serialized with `component_to_dict` / `deserialize_component_inplace` even though they aren't 

13# Components. Despite the name, those helpers aren't component-specific: they just produce/consume the standard 

14# `{"type", "init_parameters"}` dict, and deserialization enforces the import allowlist (so only trusted modules are 

15# loaded). 

16def _serialize_hooks_dictionary(hooks: dict[HookPoint, list[Hook]]) -> dict[str, list[dict[str, Any]]]: 

17 """ 

18 Serialize a hook-point-keyed dict of hooks to plain dictionaries. 

19 

20 :param hooks: Hooks keyed by hook point; each hook must implement `to_dict`. 

21 :returns: The same mapping with each hook replaced by its serialized dictionary. 

22 """ 

23 return { 

24 hook_point: [component_to_dict(obj=h, name="hook") for h in hook_list] 

25 for hook_point, hook_list in hooks.items() 

26 } 

27 

28 

29def _deserialize_hooks_dictionary(data: dict[str, list[dict[str, Any]]]) -> dict[str, list[Hook]]: 

30 """ 

31 Deserialize a hook-point-keyed dict of hooks from its serialized form. 

32 

33 :param data: Hook-point-keyed lists of serialized hook dictionaries (each with a `type` field). 

34 :returns: The same mapping with each entry rebuilt into a `Hook` instance. 

35 """ 

36 deserialized: dict[str, list[Hook]] = {} 

37 for hook_point, serialized_hooks in data.items(): 

38 hooks: list[Hook] = [] 

39 for serialized_hook in serialized_hooks: 

40 wrapper: dict[str, Any] = {"hook": serialized_hook} 

41 deserialize_component_inplace(wrapper, key="hook") 

42 hooks.append(wrapper["hook"]) 

43 deserialized[hook_point] = hooks 

44 return deserialized 

45 

46 

47def _unique_hooks(hooks: dict[HookPoint, list[Hook]]) -> list[Hook]: 

48 """ 

49 Collect each distinct hook once, preserving first-seen order. 

50 

51 A hook may be registered under several hook points; deduplicating by identity ensures lifecycle methods 

52 (warm up / close) run once per hook object. 

53 

54 :param hooks: Hooks keyed by hook point. 

55 :returns: The distinct hook objects, in the order first encountered. 

56 """ 

57 unique: list[Hook] = [] 

58 seen: set[int] = set() 

59 for hook_list in hooks.values(): 

60 for h in hook_list: 

61 if id(h) not in seen: 

62 seen.add(id(h)) 

63 unique.append(h) 

64 return unique 

65 

66 

67def warm_up_hooks(hooks: dict[HookPoint, list[Hook]]) -> None: 

68 """ 

69 Warm up every hook that defines a `warm_up` method (e.g. to open clients or load credentials). 

70 

71 :param hooks: Hooks keyed by hook point. Each distinct hook is warmed up at most once. 

72 """ 

73 for h in _unique_hooks(hooks): 

74 if hasattr(h, "warm_up"): 

75 h.warm_up() 

76 

77 

78async def warm_up_hooks_async(hooks: dict[HookPoint, list[Hook]]) -> None: 

79 """ 

80 Warm up every hook, awaiting `warm_up_async` when defined and falling back to `warm_up` otherwise. 

81 

82 :param hooks: Hooks keyed by hook point. Each distinct hook is warmed up at most once. 

83 """ 

84 for h in _unique_hooks(hooks): 

85 warm_up_async = getattr(h, "warm_up_async", None) 

86 if warm_up_async is not None: 

87 await warm_up_async() 

88 elif hasattr(h, "warm_up"): 

89 h.warm_up() 

90 

91 

92def close_hooks(hooks: dict[HookPoint, list[Hook]]) -> None: 

93 """ 

94 Release the resources of every hook that defines a `close` method. 

95 

96 :param hooks: Hooks keyed by hook point. Each distinct hook is closed at most once. 

97 """ 

98 for h in _unique_hooks(hooks): 

99 if hasattr(h, "close"): 

100 h.close() 

101 

102 

103async def close_hooks_async(hooks: dict[HookPoint, list[Hook]]) -> None: 

104 """ 

105 Release hook resources, awaiting `close_async` when defined and falling back to `close` otherwise. 

106 

107 :param hooks: Hooks keyed by hook point. Each distinct hook is closed at most once. 

108 """ 

109 for h in _unique_hooks(hooks): 

110 close_async = getattr(h, "close_async", None) 

111 if close_async is not None: 

112 await close_async() 

113 elif hasattr(h, "close"): 

114 h.close()