"""
Working Memory - 当前任务上下文（内存级存储）
"""


class WorkingMemory:
    """当前任务上下文，内存级存储"""

    def __init__(self):
        self.current_task = None
        self.recent_messages = []
        self.max_size = 10  # 最大保留 10 轮对话

    def get(self) -> dict:
        """获取当前记忆"""
        return {
            'current_task': self.current_task,
            'recent_messages': self.recent_messages[-self.max_size:]
        }

    def add(self, user_input: str, response: str):
        """添加对话到记忆"""
        self.recent_messages.append({
            'role': 'user',
            'content': user_input
        })
        self.recent_messages.append({
            'role': 'assistant',
            'content': response
        })

        # 保持固定大小
        if len(self.recent_messages) > self.max_size * 2:
            self.recent_messages = self.recent_messages[-self.max_size * 2:]

    def set_current_task(self, task: str):
        """设置当前任务"""
        self.current_task = task

    def clear(self):
        """清空记忆"""
        self.recent_messages = []
        self.current_task = None

    def size(self) -> int:
        """获取对话轮数"""
        return len(self.recent_messages) // 2
