from langchain.memory import ConversationBufferMemory
from typing import Dict, Any
from pydantic import Field

class LimitedConversationBufferMemory(ConversationBufferMemory):
    """
    可以设置大小的对话Memeory缓存
    """

    max_size: int = Field(default=10, description="缓存大小")

    def __init__(self, max_size: int = 500, **kwargs: Any):
        super().__init__(**kwargs)
        self.max_size = max_size

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        # 先调用父类的保存方法
        super().save_context(inputs, outputs)
        # 当缓冲区超过最大大小时，移除最早的消息
        while len(self.chat_memory.messages) > self.max_size:
            self.chat_memory.messages.pop(0)