Coverage for src / dataknobs_bots / middleware / base.py: 100%

1"""Base middleware interface for bot request/response lifecycle.""" 

2 

3from abc import ABC, abstractmethod 

4from typing import Any 

5 

6from dataknobs_bots.bot.context import BotContext 

7 

8 

9class Middleware(ABC): 

10 """Abstract base class for bot middleware. 

11 

12 Middleware provides hooks into the bot request/response lifecycle: 

13 - before_message: Called before processing user message 

14 - after_message: Called after generating bot response (non-streaming) 

15 - post_stream: Called after streaming response completes 

16 - on_error: Called when an error occurs 

17 

18 Example: 

19 ```python 

20 class MyMiddleware(Middleware): 

21 async def before_message(self, message: str, context: BotContext) -> None: 

22 print(f"Processing: {message}") 

23 

24 async def after_message( 

25 self, response: str, context: BotContext, **kwargs: Any 

26 ) -> None: 

27 print(f"Response: {response}") 

28 

29 async def post_stream( 

30 self, message: str, response: str, context: BotContext 

31 ) -> None: 

32 print(f"Streamed response to '{message}': {response}") 

33 

34 async def on_error( 

35 self, error: Exception, message: str, context: BotContext 

36 ) -> None: 

37 print(f"Error: {error}") 

38 ``` 

39 """ 

40 

41 @abstractmethod 

42 async def before_message(self, message: str, context: BotContext) -> None: 

43 """Called before processing user message. 

44 

45 Args: 

46 message: User's input message 

47 context: Bot context with conversation and user info 

48 """ 

49 ... 

50 

51 @abstractmethod 

52 async def after_message( 

53 self, response: str, context: BotContext, **kwargs: Any 

54 ) -> None: 

55 """Called after generating bot response (non-streaming). 

56 

57 Args: 

58 response: Bot's generated response 

59 context: Bot context 

60 **kwargs: Additional data (e.g., tokens_used, response_time_ms, provider, model) 

61 """ 

62 ... 

63 

64 @abstractmethod 

65 async def post_stream( 

66 self, message: str, response: str, context: BotContext 

67 ) -> None: 

68 """Called after streaming response completes. 

69 

70 This hook is called after stream_chat() finishes streaming all chunks. 

71 It provides both the original user message and the complete accumulated 

72 response, useful for logging, analytics, or post-processing. 

73 

74 Args: 

75 message: Original user message that triggered the stream 

76 response: Complete accumulated response from streaming 

77 context: Bot context 

78 """ 

79 ... 

80 

81 @abstractmethod 

82 async def on_error( 

83 self, error: Exception, message: str, context: BotContext 

84 ) -> None: 

85 """Called when an error occurs during message processing. 

86 

87 Args: 

88 error: The exception that occurred 

89 message: User message that caused the error 

90 context: Bot context 

91 """ 

92 ...