Coverage for src / dataknobs_bots / memory / buffer.py: 62%

1"""Buffer memory implementation for simple FIFO message storage.""" 

2 

3from collections import deque 

4from typing import Any 

5 

6from .base import Memory 

7 

8 

9class BufferMemory(Memory): 

10 """Simple buffer memory keeping last N messages. 

11 

12 This implementation uses a fixed-size buffer that keeps the most recent 

13 messages in memory. When the buffer is full, the oldest messages are 

14 automatically removed. 

15 

16 Attributes: 

17 max_messages: Maximum number of messages to keep in buffer 

18 messages: Deque containing the messages 

19 """ 

20 

21 def __init__(self, max_messages: int = 10): 

22 """Initialize buffer memory. 

23 

24 Args: 

25 max_messages: Maximum number of messages to keep 

26 """ 

27 self.max_messages = max_messages 

28 self.messages: deque[dict[str, Any]] = deque(maxlen=max_messages) 

29 

30 async def add_message( 

31 self, content: str, role: str, metadata: dict[str, Any] | None = None 

32 ) -> None: 

33 """Add message to buffer. 

34 

35 Args: 

36 content: Message content 

37 role: Message role 

38 metadata: Optional metadata 

39 """ 

40 self.messages.append({"content": content, "role": role, "metadata": metadata or {}}) 

41 

42 async def get_context(self, current_message: str) -> list[dict[str, Any]]: 

43 """Get all messages in buffer. 

44 

45 The current_message parameter is not used in buffer memory since 

46 we simply return all buffered messages in order. 

47 

48 Args: 

49 current_message: Not used in buffer memory 

50 

51 Returns: 

52 List of all buffered messages 

53 """ 

54 return list(self.messages) 

55 

56 async def clear(self) -> None: 

57 """Clear all messages from buffer.""" 

58 self.messages.clear()