Coverage for haystack/hooks/invocation.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 import tracing 

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

9from haystack.core.errors import SerializationError 

10from haystack.core.serialization import generate_qualified_class_name 

11from haystack.hooks.from_function import FunctionHook 

12from haystack.hooks.protocol import Hook, HookPoint 

13from haystack.tracing import Span 

14from haystack.utils.async_utils import _execute_component_async 

15from haystack.utils.callable_serialization import serialize_callable 

16 

17 

18def _hook_name(hook: Hook) -> str: 

19 """Return a human-readable identifier for a hook.""" 

20 if isinstance(hook, FunctionHook): 

21 # FunctionHook is a generic wrapper, so identify it by its wrapped callable instead. 

22 function = hook.function or hook.async_function 

23 if function is not None: 

24 try: 

25 try: 

26 return serialize_callable(callable_handle=function) 

27 except SerializationError: 

28 # Nested functions and lambdas are unsupported by serialization, but their qualified names are still 

29 # useful. 

30 return f"{function.__module__}.{function.__qualname__}" 

31 # functools.partial and callable-object instances may not expose __module__, __qualname__, or __name__. 

32 except Exception: 

33 pass 

34 return type(hook).__name__ 

35 

36 

37def _create_hook_span(hook: Hook, hook_point: HookPoint, parent_span: Span | None) -> Any: 

38 """Create a content-free tracing span for one hook invocation.""" 

39 return tracing.tracer.trace( 

40 "haystack.agent.hook", 

41 tags={ 

42 "haystack.agent.hook.point": hook_point, 

43 "haystack.agent.hook.name": _hook_name(hook=hook), 

44 "haystack.agent.hook.type": generate_qualified_class_name(cls=type(hook)), 

45 }, 

46 parent_span=parent_span, 

47 ) 

48 

49 

50def _run_hooks(hooks: dict[HookPoint, list[Hook]], hook_point: HookPoint, state: State) -> None: 

51 """ 

52 Run every hook registered for the given hook point, in list order. 

53 

54 :param hooks: Hooks keyed by hook point. 

55 :param hook_point: The hook point whose hooks to run; hooks registered under other hook points are skipped. 

56 :param state: The Agent's live `State`, passed to each hook and mutated in place. 

57 """ 

58 hooks_to_run = hooks.get(hook_point, []) 

59 if not hooks_to_run: 

60 return 

61 

62 parent_span = tracing.tracer.current_span() 

63 for h in hooks_to_run: 

64 with _create_hook_span(hook=h, hook_point=hook_point, parent_span=parent_span): 

65 h.run(state) 

66 

67 

68async def _run_hooks_async(hooks: dict[HookPoint, list[Hook]], hook_point: HookPoint, state: State) -> None: 

69 """ 

70 Run every hook for the given hook point, preferring `run_async` and offloading sync-only `run` hooks. 

71 

72 :param hooks: Hooks keyed by hook point. 

73 :param hook_point: The hook point whose hooks to run; hooks registered under other hook points are skipped. 

74 :param state: The Agent's live `State`, passed to each hook and mutated in place. 

75 """ 

76 hooks_to_run = hooks.get(hook_point, []) 

77 if not hooks_to_run: 

78 return 

79 

80 parent_span = tracing.tracer.current_span() 

81 for h in hooks_to_run: 

82 with _create_hook_span(hook=h, hook_point=hook_point, parent_span=parent_span): 

83 await _execute_component_async(h, state=state)