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

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

2

3 from collections import deque

4 from typing import Any

5

6 from .base import Memory

7

8

9 class 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()