【AI Agent教程】【MetaGPT】由易到难，深入源码：看MetaGPT的长短时记忆如何实现

class Memory(BaseModel):
    """The most basic memory: super-memory"""
    storage: list[SerializeAsAny[Message]] = []
    index: DefaultDict[str, list[SerializeAsAny[Message]]] = Field(default_factory=lambda: defaultdict(list))
    ignore_id: bool = False

def add(self, message: Message):
    """Add a new message to storage, while updating the index"""
    if self.ignore_id:
        message.id = IGNORED_MESSAGE_ID
    if message in self.storage:
        return
    self.storage.append(message)
    if message.cause_by:
        self.index[message.cause_by].append(message)

def find_news(self, observed: list[Message], k=0) -> list[Message]:
    """find news (previously unseen messages) from the the most recent k memories, from all memories when k=0"""
    already_observed = self.get(k)
    news: list[Message] = []
    for i in observed:
        if i in already_observed:
            continue
        news.append(i)
    return news

class LongTermMemory(Memory):
    """
    The Long-term memory for Roles
    - recover memory when it staruped
    - update memory when it changed
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)
    memory_storage: MemoryStorage = Field(default_factory=MemoryStorage)
    rc: Optional[RoleContext] = None
    msg_from_recover: bool = False

def add(self, message: Message):
    super().add(message)
    for action in self.rc.watch:
        if message.cause_by == action and not self.msg_from_recover:
            # currently, only add role's watching messages to its memory_storage
            # and ignore adding messages from recover repeatedly
            self.memory_storage.add(message)

def add(self, message: Message) -> bool:
    """add message into memory storage"""
    self.faiss_engine.add_objs([message])
    logger.info(f"Role {self.role_id}'s memory_storage add a message")

def recover_memory(self, role_id: str, rc: RoleContext):
    self.memory_storage.recover_memory(role_id)
    self.rc = rc
    if not self.memory_storage.is_initialized:
        logger.warning(f"It may the first time to run Role {role_id}, the long-term memory is empty")
    else:
        logger.warning(f"Role {role_id} has existing memory storage and has recovered them.")
    self.msg_from_recover = True
    # self.add_batch(messages) # TODO no need
    self.msg_from_recover = False

def recover_memory(self, role_id: str) -> list[Message]:
    self.role_id = role_id
    self.role_mem_path = Path(DATA_PATH / f"role_mem/{self.role_id}/")
    self.role_mem_path.mkdir(parents=True, exist_ok=True)
    self.cache_dir = self.role_mem_path
    if self.role_mem_path.joinpath("default__vector_store.json").exists():
        self.faiss_engine = SimpleEngine.from_index(
            index_config=FAISSIndexConfig(persist_path=self.cache_dir),
            retriever_configs=[FAISSRetrieverConfig()],
            embed_model=self.embedding,
        )
    else:
        self.faiss_engine = SimpleEngine.from_objs(
            objs=[], retriever_configs=[FAISSRetrieverConfig()], embed_model=self.embedding
        )
    self._initialized = True

async def find_news(self, observed: list[Message], k=0) -> list[Message]:
    """
    find news (previously unseen messages) from the the most recent k memories, from all memories when k=0
        1. find the short-term memory(stm) news
        2. furthermore, filter out similar messages based on ltm(long-term memory), get the final news
    """
    stm_news = super().find_news(observed, k=k)  # shot-term memory news
    if not self.memory_storage.is_initialized:
        # memory_storage hasn't initialized, use default `find_news` to get stm_news
        return stm_news
    ltm_news: list[Message] = []
    for mem in stm_news:
        # filter out messages similar to those seen previously in ltm, only keep fresh news
        mem_searched = await self.memory_storage.search_similar(mem)
        if len(mem_searched) == 0:
            ltm_news.append(mem)
    return ltm_news[-k:]

def persist(self):
    if self.faiss_engine:
        self.faiss_engine.retriever._index.storage_context.persist(self.cache_dir)

class BrainMemory(BaseModel):
    history: List[Message] = Field(default_factory=list)
    knowledge: List[Message] = Field(default_factory=list)
    historical_summary: str = ""
    last_history_id: str = ""
    is_dirty: bool = False
    last_talk: Optional[str] = None
    cacheable: bool = True
    llm: Optional[BaseLLM] = Field(default=None, exclude=True)

async def summarize(self, llm, max_words=200, keep_language: bool = False, limit: int = -1, **kwargs):
    if isinstance(llm, MetaGPTLLM):
        return await self._metagpt_summarize(max_words=max_words)
    self.llm = llm
    return await self._openai_summarize(llm=llm, max_words=max_words, keep_language=keep_language, limit=limit)
async def _openai_summarize(self, llm, max_words=200, keep_language: bool = False, limit: int = -1):
    texts = [self.historical_summary]
    for m in self.history:
        texts.append(m.content)
    text = "\n".join(texts)
    text_length = len(text)
    if limit > 0 and text_length < limit:
        return text
    summary = await self._summarize(text=text, max_words=max_words, keep_language=keep_language, limit=limit)
    if summary:
        await self.set_history_summary(history_summary=summary, redis_key=config.redis_key)
        return summary
    raise ValueError(f"text too long:{text_length}")

async def get_title(self, llm, max_words=5, **kwargs) -> str:
    """Generate text title"""
    if isinstance(llm, MetaGPTLLM):
        return self.history[0].content if self.history else "New"
    summary = await self.summarize(llm=llm, max_words=500)
    language = config.language
    command = f"Translate the above summary into a {language} title of less than {max_words} words."
    summaries = [summary, command]
    msg = "\n".join(summaries)
    logger.debug(f"title ask:{msg}")
    response = await llm.aask(msg=msg, system_msgs=[], stream=False)
    logger.debug(f"title rsp: {response}")
    return response