Class DefaultMemory

ms_agent/memory/default_memory.py:79–736

The memory refine tool

class DefaultMemory(Memory):
    """The memory refine tool"""

    def __init__(self, config: DictConfig):
        super().__init__(config)
        memory_config = config.memory.default_memory
        self.user_id: Optional[str] = getattr(memory_config, 'user_id',
                                              DEFAULT_USER)
        self.agent_id: Optional[str] = getattr(memory_config, 'agent_id', None)
        self.run_id: Optional[str] = getattr(memory_config, 'run_id', None)
        self.compress: Optional[bool] = getattr(config, 'compress', True)
        self.is_retrieve: Optional[bool] = getattr(config, 'is_retrieve', True)
        self.path: Optional[str] = getattr(
            memory_config, 'path',
            os.path.join(DEFAULT_OUTPUT_DIR, '.default_memory'))
        self.history_mode = getattr(memory_config, 'history_mode', 'add')
        self.ignore_roles: List[str] = getattr(memory_config, 'ignore_roles',
                                               ['tool', 'system'])
        self.ignore_fields: List[str] = getattr(memory_config, 'ignore_fields',
                                                ['reasoning_content'])
        self.search_limit: int = getattr(memory_config, 'search_limit',
                                         DEFAULT_SEARCH_LIMIT)
        # Add lock for thread safety in shared usage
        self._lock = asyncio.Lock()
        self.memory = self._init_memory_obj()
        self.load_cache()

    async def init_cache_messages(self):
        if len(self.cache_messages) and not len(self.memory_snapshot):
            for id, messages in self.cache_messages.items():
                await self.add_single(messages, msg_id=id)

    def save_cache(self):
        """
        Save self.max_msg_id, self.cache_messages, and self.memory_snapshot to self.path/cache_messages.json
        """
        cache_file = os.path.join(self.path, 'cache_messages.json')

        # Ensure the directory exists
        os.makedirs(self.path, exist_ok=True)

        data = {
            'max_msg_id': self.max_msg_id,
            'cache_messages': {
                str(k): ([msg.to_dict() for msg in msg_list], _hash)
                for k, (msg_list, _hash) in self.cache_messages.items()
            },
            'memory_snapshot': [mm.to_dict() for mm in self.memory_snapshot]
        }

        with open(cache_file, 'w', encoding='utf-8') as f:
            json5.dump(data, f, indent=2, ensure_ascii=False)

    def load_cache(self):
        """
        Load data from self.path/cache_messages.json into self.max_msg_id, self.cache_messages, and self.memory_snapshot
        """
        cache_file = os.path.join(self.path, 'cache_messages.json')