MetaGPT源码 （Memory 类）

目录

• • [MetaGPT源码：Memory 类](#MetaGPT源码：Memory 类)
  • 例子

MetaGPT源码：Memory 类

这段代码定义了一个名为 Memory 的类，用于存储和管理消息(Message)对象。Memory 提供了多种操作消息的功能，包括添加单条或批量消息、按角色或内容筛选消息、删除最新消息、清空存储、按触发动作获取消息等。消息存储在 storage 列表中，同时使用 index 字典对消息的触发动作(cause_by)进行索引，以支持快速检索。通过这些方法，Memory 类能够高效地管理消息流，实现对消息的组织、查询和更新，适用于需要处理复杂消息交互的场景。

from collections import defaultdict
from typing import DefaultDict, Iterable, Set

from pydantic import BaseModel, Field, SerializeAsAny

from metagpt.const import IGNORED_MESSAGE_ID
from metagpt.schema import Message
from metagpt.utils.common import any_to_str, any_to_str_set


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 add_batch(self, messages: Iterable[Message]):
        for message in messages:
            self.add(message)

    def get_by_role(self, role: str) -> list[Message]:
        """Return all messages of a specified role"""
        return [message for message in self.storage if message.role == role]

    def get_by_content(self, content: str) -> list[Message]:
        """Return all messages containing a specified content"""
        return [message for message in self.storage if content in message.content]

    def delete_newest(self) -> "Message":
        """delete the newest message from the storage"""
        if len(self.storage) > 0:
            newest_msg = self.storage.pop()
            if newest_msg.cause_by and newest_msg in self.index[newest_msg.cause_by]:
                self.index[newest_msg.cause_by].remove(newest_msg)
        else:
            newest_msg = None
        return newest_msg

    def delete(self, message: Message):
        """Delete the specified message from storage, while updating the index"""
        if self.ignore_id:
            message.id = IGNORED_MESSAGE_ID
        self.storage.remove(message)
        if message.cause_by and message in self.index[message.cause_by]:
            self.index[message.cause_by].remove(message)

    def clear(self):
        """Clear storage and index"""
        self.storage = []
        self.index = defaultdict(list)

    def count(self) -> int:
        """Return the number of messages in storage"""
        return len(self.storage)

    def try_remember(self, keyword: str) -> list[Message]:
        """Try to recall all messages containing a specified keyword"""
        return [message for message in self.storage if keyword in message.content]

    def get(self, k=0) -> list[Message]:
        """Return the most recent k memories, return all when k=0"""
        return self.storage[-k:]

    def find_news(self, observed: list[Message], k=0) -> list[Message]:
        """find news (previously unseen messages) from 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

    def get_by_action(self, action) -> list[Message]:
        """Return all messages triggered by a specified Action"""
        index = any_to_str(action)
        return self.index[index]

    def get_by_actions(self, actions: Set) -> list[Message]:
        """Return all messages triggered by specified Actions"""
        rsp = []
        indices = any_to_str_set(actions)
        for action in indices:
            if action not in self.index:
                continue
            rsp += self.index[action]
        return rsp

2024-12-11 22:38:16.092 | INFO | metagpt.const:get_metagpt_package_root:21 - Package root set to d:\llm\metagpt

例子

以下是一些例子，帮助你理解 Memory 类及其方法的使用：

1. 创建 Memory 对象并添加消息

from metagpt.schema import Message

# 创建 Memory 实例
memory = Memory()

# 创建一个 Message 实例
message1 = Message(role="user", content="Hello!", cause_by=None, id="1")
message2 = Message(role="assistant", content="Hi there!", cause_by="1", id="2")

# 添加消息
memory.add(message1)
memory.add(message2)

print("Storage:", memory.storage)

Storage: [user: Hello!, assistant: Hi there!]

2. 批量添加消息

messages = [
    Message(role="user", content="How are you?", cause_by=None, id="3"),
    Message(role="assistant", content="I'm fine, thank you!", cause_by="3", id="4"),
]

memory.add_batch(messages)

print("Storage after batch add:", memory.storage)

Storage after batch add: [user: Hello!, assistant: Hi there!, user: How are you?, assistant: I'm fine, thank you!]

3. 按角色获取消息

user_messages = memory.get_by_role("user")
print("Messages from user:", user_messages)

Messages from user: [user: Hello!, user: How are you?]

4. 按内容查找消息

matching_messages = memory.get_by_content("fine")
print("Messages containing 'fine':", matching_messages)

Messages containing 'fine': [assistant: I'm fine, thank you!]

5. 删除最新消息

newest_message = memory.delete_newest()
print("Deleted newest message:", newest_message)
print("Storage after deletion:", memory.storage)

Deleted newest message: assistant: I'm fine, thank you!
Storage after deletion: [user: Hello!, assistant: Hi there!, user: How are you?]

6. 清空存储

memory.clear()
print("Storage after clear:", memory.storage)

Storage after clear: []

7. 按触发动作查找消息

# 添加一些消息
memory.add(Message(role="assistant", content="Task completed!", cause_by="action_1", id="5"))
memory.add(Message(role="assistant", content="Another task completed!", cause_by="action_2", id="6"))

# 根据动作获取消息
action_messages = memory.get_by_action("action_1")
print("Messages triggered by 'action_1':", action_messages)

Messages triggered by 'action_1': [assistant: Task completed!]

如果有任何问题，欢迎在评论区提问。