autogen/python/samples/common/memory/_buffered.py

from typing import Any, List, Mapping

from agnext.components.memory import ChatMemory
from agnext.components.models import FunctionExecutionResultMessage

from ..types import Message


class BufferedChatMemory(ChatMemory[Message]):
    """A buffered chat memory that keeps a view of the last n messages,
    where n is the buffer size. The buffer size is set at initialization.

    Args:
        buffer_size (int): The size of the buffer.

    """

    def __init__(self, buffer_size: int) -> None:
        self._messages: List[Message] = []
        self._buffer_size = buffer_size

    async def add_message(self, message: Message) -> None:
        """Add a message to the memory."""
        self._messages.append(message)

    async def get_messages(self) -> List[Message]:
        """Get at most `buffer_size` recent messages."""
        messages = self._messages[-self._buffer_size :]
        # Handle the first message is a function call result message.
        if messages and isinstance(messages[0], FunctionExecutionResultMessage):
            # Remove the first message from the list.
            messages = messages[1:]
        return messages

    async def clear(self) -> None:
        """Clear the message memory."""
        self._messages = []

    def save_state(self) -> Mapping[str, Any]:
        return {
            "messages": [message for message in self._messages],
            "buffer_size": self._buffer_size,
        }

    def load_state(self, state: Mapping[str, Any]) -> None:
        self._messages = state["messages"]
        self._buffer_size = state["buffer_size"]