Spaces:
Sleeping
Sleeping
File size: 1,925 Bytes
e679d69 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
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
|