Coverage for haystack/hooks/from_function.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import inspect 

6from collections.abc import Awaitable, Callable 

7from typing import Any, cast 

8 

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

10from haystack.core.serialization import default_from_dict, default_to_dict 

11from haystack.core.type_utils import _resolve_parameter_types 

12from haystack.utils.callable_serialization import deserialize_callable, serialize_callable 

13 

14 

15def _takes_single_state_argument(function: Callable) -> bool: 

16 """ 

17 Return whether `function` takes exactly one parameter annotated with `State`. 

18 

19 :param function: The callable to inspect. 

20 :returns: True if the signature is a single `State`-annotated parameter, False otherwise. 

21 """ 

22 param_types = list(_resolve_parameter_types(function).values()) 

23 return len(param_types) == 1 and param_types[0] is State 

24 

25 

26class FunctionHook: 

27 """ 

28 Wraps a function (or a sync/async pair) into a serializable `Hook`. 

29 

30 Produced by the `@hook` decorator for the single-function case. To give a hook both an optimized sync and async 

31 path, construct it directly with both `function` and `async_function` set. 

32 """ 

33 

34 def __init__( 

35 self, 

36 function: Callable[[State], None] | None = None, 

37 async_function: Callable[[State], Awaitable[None]] | None = None, 

38 ) -> None: 

39 """ 

40 Initialize the hook with a synchronous function, an async function, or both. 

41 

42 :param function: The synchronous function invoked by `run`. Must be a regular function — coroutine functions 

43 should be passed to `async_function` instead. Either `function` or `async_function` (or both) must be set. 

44 :param async_function: Optional coroutine function awaited by `run_async`. When only `async_function` is set, 

45 `run` raises a `RuntimeError`. When only `function` is set, `run_async` calls `function`. 

46 :raises ValueError: If neither is set, if `function` is a coroutine function, if `async_function` is not, or 

47 if a provided function does not declare a `State`-typed parameter. 

48 """ 

49 if function is None and async_function is None: 

50 raise ValueError("A FunctionHook requires at least one of `function` or `async_function` to be set.") 

51 if function is not None and inspect.iscoroutinefunction(function): 

52 raise ValueError( 

53 f"`function` must be a synchronous function. '{function.__name__}' is a coroutine function. " 

54 "Pass it as `async_function` instead." 

55 ) 

56 if async_function is not None and not inspect.iscoroutinefunction(async_function): 

57 raise ValueError( 

58 f"`async_function` must be a coroutine function defined with `async def`. " 

59 f"Got '{getattr(async_function, '__name__', repr(async_function))}'." 

60 ) 

61 for func in (function, async_function): 

62 if func is not None and not _takes_single_state_argument(func): 

63 raise ValueError( 

64 f"Hook function '{func.__name__}' must take a single parameter annotated with `State` " 

65 "(e.g. `def my_hook(state: State) -> None`)." 

66 ) 

67 self.function = function 

68 self.async_function = async_function 

69 

70 def run(self, state: State) -> None: 

71 """ 

72 Run the synchronous function against the live `State`. 

73 

74 :param state: The Agent's live `State`, mutated in place by the wrapped function. 

75 :raises RuntimeError: If the hook only has an `async_function`; use the Agent's async run methods instead. 

76 """ 

77 if self.function is None: 

78 raise RuntimeError( 

79 "This FunctionHook only has an `async_function` and cannot run in a synchronous Agent run. " 

80 "Use the Agent's async run methods, or provide a synchronous `function`." 

81 ) 

82 self.function(state) 

83 

84 async def run_async(self, state: State) -> None: 

85 """ 

86 Await the async function if set, otherwise call the synchronous function. 

87 

88 :param state: The Agent's live `State`, mutated in place by the wrapped function. 

89 """ 

90 if self.async_function is not None: 

91 await self.async_function(state) 

92 else: 

93 self.function(state) # type: ignore[misc] # guaranteed non-None: at least one is always set 

94 

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

96 """ 

97 Serialize the hook, storing each wrapped function as an importable reference. 

98 

99 :returns: A dictionary with the hook's type and the import paths of its sync/async functions. 

100 """ 

101 return default_to_dict( 

102 self, 

103 function=serialize_callable(self.function) if self.function is not None else None, 

104 async_function=serialize_callable(self.async_function) if self.async_function is not None else None, 

105 ) 

106 

107 @classmethod 

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

109 """ 

110 Deserialize the hook, resolving each function from its importable reference. 

111 

112 :param data: The serialized hook dictionary produced by `to_dict`. 

113 :returns: The reconstructed `FunctionHook`. 

114 """ 

115 init_params = data.get("init_parameters", {}) 

116 if init_params.get("function") is not None: 

117 init_params["function"] = deserialize_callable(init_params["function"]) 

118 if init_params.get("async_function") is not None: 

119 init_params["async_function"] = deserialize_callable(init_params["async_function"]) 

120 return default_from_dict(cls, data) 

121 

122 

123def hook(function: Callable[[State], None | Awaitable[None]]) -> FunctionHook: 

124 """ 

125 Wrap a function into a `Hook` the Agent can invoke during its run loop. 

126 

127 The decorated function receives the Agent's `State` and influences the run by mutating it in place. A coroutine 

128 function is wrapped as the hook's async path; a regular function as its sync path. To give a single hook both 

129 paths, construct a `FunctionHook` directly with both `function` and `async_function`. 

130 

131 ### Usage example 

132 

133 ```python 

134 from haystack.components.agents import Agent 

135 from haystack.components.generators.chat import OpenAIChatGenerator 

136 from haystack.hooks import hook 

137 from haystack.components.agents.state import State 

138 from haystack.dataclasses import ChatMessage 

139 from haystack.tools import tool 

140 

141 @tool 

142 def weather_tool(city: str) -> str: 

143 '''Get the current weather for a given city.''' 

144 return f"The weather in {city} is sunny." 

145 

146 @tool 

147 def save(content: str) -> str: 

148 '''Save content to durable storage.''' 

149 return "Saved." 

150 

151 @hook 

152 def require_save(state: State) -> None: 

153 if state.get("tool_call_counts", {}).get("save", 0) == 0: 

154 state.set("messages", [ChatMessage.from_system("You must call `save` before finishing.")]) 

155 state.set("continue_run", True) 

156 

157 agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[weather_tool, save], hooks={"on_exit": [require_save]}) 

158 ``` 

159 

160 :param function: A callable taking the Agent's `State` and returning `None` (sync or async). 

161 :returns: A `FunctionHook` wrapping the function. 

162 """ 

163 # `iscoroutinefunction` narrows which slot the callable belongs in; cast to satisfy the typed slots. 

164 if inspect.iscoroutinefunction(function): 

165 return FunctionHook(async_function=cast("Callable[[State], Awaitable[None]]", function)) 

166 return FunctionHook(function=cast("Callable[[State], None]", function))