Coverage for haystack/hooks/protocol.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, Literal, Protocol, get_args 

6 

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

8 

9# Points in the Agent's run loop at which hooks can be registered. 

10HookPoint = Literal["before_run", "before_llm", "before_tool", "after_tool", "on_exit", "after_run"] 

11 

12BEFORE_RUN: HookPoint = "before_run" 

13BEFORE_LLM: HookPoint = "before_llm" 

14BEFORE_TOOL: HookPoint = "before_tool" 

15AFTER_TOOL: HookPoint = "after_tool" 

16ON_EXIT: HookPoint = "on_exit" 

17AFTER_RUN: HookPoint = "after_run" 

18VALID_HOOK_POINTS: tuple[HookPoint, ...] = get_args(HookPoint) 

19 

20 

21class Hook(Protocol): 

22 """ 

23 A callable the Agent invokes at a point in its run loop, receiving the live `State`. 

24 

25 A hook influences the run only by mutating `State` in place. At least `messages` (the conversation), 

26 `step_count`, `token_usage` and `tool_call_counts` are available; any additional keys defined in the Agent's 

27 `state_schema` are available too. The same hook object can be registered under multiple hook points. 

28 

29 Implement this protocol directly for stateful hooks (e.g. one wrapping a component), or use the `@hook` decorator to 

30 wrap a plain `(State) -> None` function. 

31 

32 A hook may additionally define `async def run_async(self, state: State) -> None` for true async behavior; when 

33 absent, the Agent calls `run` during async runs. It is left off this protocol on purpose so sync-only hooks 

34 don't have to implement it. 

35 

36 A hook may also implement the optional lifecycle methods `warm_up` / `warm_up_async` and `close` / `close_async`. 

37 The Agent calls them from its own `warm_up` / `warm_up_async` and `close` / `close_async`, so a hook can defer 

38 opening clients or reading credentials until warm-up and release them on close. Because warm-up runs before every 

39 Agent run, hooks should avoid repeating expensive initialization, for example by returning early if a client has 

40 already been initialized. 

41 """ 

42 

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

44 """Run the hook against the live `State`, mutating it in place.""" 

45 ... 

46 

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

48 """Serialize the hook to a dictionary.""" 

49 ... 

50 

51 @classmethod 

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

53 """Deserialize the hook from a dictionary.""" 

54 ...