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

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

2

3 from abc import ABC, abstractmethod

4 from typing import Any

5

6 from dataknobs_bots.bot.context import BotContext

7

8

9 class 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 ...