| from typing import Callable, Dict, List, Optional, Union | |
| from lagent.schema import AgentMessage | |
| class Memory: | |
| def __init__(self, recent_n=None) -> None: | |
| self.memory: List[AgentMessage] = [] | |
| self.recent_n = recent_n | |
| def get_memory( | |
| self, | |
| recent_n: Optional[int] = None, | |
| filter_func: Optional[Callable[[int, dict], bool]] = None, | |
| ) -> list: | |
| recent_n = recent_n or self.recent_n | |
| if recent_n is not None: | |
| memory = self.memory[-recent_n:] | |
| else: | |
| memory = self.memory | |
| if filter_func is not None: | |
| memory = [m for i, m in enumerate(memory) if filter_func(i, m)] | |
| return memory | |
| def add(self, memories: Union[List[Dict], Dict, None]) -> None: | |
| for memory in memories if isinstance(memories, | |
| (list, tuple)) else [memories]: | |
| if isinstance(memory, str): | |
| memory = AgentMessage(sender='user', content=memory) | |
| if isinstance(memory, AgentMessage): | |
| self.memory.append(memory) | |
| def delete(self, index: Union[List, int]) -> None: | |
| if isinstance(index, int): | |
| del self.memory[index] | |
| else: | |
| for i in index: | |
| del self.memory[i] | |
| def load( | |
| self, | |
| memories: Union[str, Dict, List], | |
| overwrite: bool = True, | |
| ) -> None: | |
| if overwrite: | |
| self.memory = [] | |
| if isinstance(memories, dict): | |
| self.memory.append(AgentMessage(**memories)) | |
| elif isinstance(memories, list): | |
| for m in memories: | |
| self.memory.append(AgentMessage(**m)) | |
| else: | |
| raise TypeError(f'{type(memories)} is not supported') | |
| def save(self) -> List[dict]: | |
| memory = [] | |
| for m in self.memory: | |
| memory.append(m.model_dump()) | |
| return memory | |