Coverage for src / harnessutils / hooks.py: 100%

1"""Message lifecycle hooks for ConversationManager. 

2 

3Hooks let consumers inject behaviour at the two key points in the 

4add_message() lifecycle without subclassing ConversationManager. 

5 

6Pre-hook (on_before_add_message) 

7 Fires before the message is persisted. The return value replaces 

8 the message that will be stored. Raising any exception rejects the 

9 message — it is never written to storage. 

10 

11Post-hook (on_after_add_message) 

12 Fires after the message has been successfully persisted. The return 

13 value is ignored. Exceptions propagate to the caller but the message 

14 has already been stored. 

15 

16See docs/message-hooks.md for usage guide and worked examples. 

17""" 

18 

19from collections.abc import Callable 

20from dataclasses import dataclass, field 

21from typing import TYPE_CHECKING, Any 

22 

23if TYPE_CHECKING: 

24 from harnessutils.models.message import Message 

25 

26 

27@dataclass 

28class MessageHooks: 

29 """Callbacks for the add_message() lifecycle. 

30 

31 Pass an instance to ConversationManager() to intercept messages 

32 before and after they are persisted. 

33 

34 Example — reject messages that contain a forbidden string:: 

35 

36 def guard(conv_id: str, msg: Message) -> Message: 

37 for part in msg.parts: 

38 if part.type == "text" and "forbidden" in part.text: 

39 raise ValueError("Rejected: forbidden content") 

40 return msg 

41 

42 manager = ConversationManager( 

43 message_hooks=MessageHooks(on_before_add_message=guard) 

44 ) 

45 

46 Example — write every stored message to a semantic memory backend:: 

47 

48 def index(conv_id: str, msg: Message) -> None: 

49 import asyncio 

50 asyncio.get_event_loop().run_until_complete( 

51 manager.semantic_memory.add_artifact( 

52 content=str(msg.to_dict()), 

53 kind="message", 

54 namespace=conv_id, 

55 ) 

56 ) 

57 

58 manager = ConversationManager( 

59 semantic_memory=backend, 

60 message_hooks=MessageHooks(on_after_add_message=index), 

61 ) 

62 """ 

63 

64 on_before_add_message: "Callable[[str, Message], Message] | None" = None 

65 """Called before the message is persisted. 

66 

67 Signature:: 

68 

69 def hook(conversation_id: str, message: Message) -> Message 

70 

71 - Return the message unchanged to pass it through. 

72 - Return a *different* Message to replace what gets stored 

73 (e.g. for content redaction or metadata injection). 

74 - Raise any exception to reject the message. The message will 

75 NOT be stored, and the exception propagates to the caller. 

76 on_after_add_message will NOT be called. 

77 """ 

78 

79 on_after_add_message: "Callable[[str, Message], None] | None" = None 

80 """Called after the message has been successfully persisted. 

81 

82 Signature:: 

83 

84 def hook(conversation_id: str, message: Message) -> None 

85 

86 - The message received is the one actually stored (i.e. the value 

87 returned by on_before_add_message, if one was set). 

88 - Return value is ignored. 

89 - Exceptions propagate to the caller. The message has already been 

90 stored at this point — raising here does not undo storage. 

91 - Intended for side effects: audit logging, semantic memory writes, 

92 metrics, etc. 

93 """ 

94 

95 metadata: dict[str, Any] = field(default_factory=dict) 

96 """Arbitrary metadata you can read inside your hook implementations."""