目录
-
- [MetaGPT源码：Memory 类](#MetaGPT源码：Memory 类)
- 例子
MetaGPT源码：Memory 类
这段代码定义了一个名为 Memory 的类，用于存储和管理消息(Message)对象。Memory 提供了多种操作消息的功能，包括添加单条或批量消息、按角色或内容筛选消息、删除最新消息、清空存储、按触发动作获取消息等。消息存储在 storage 列表中，同时使用 index 字典对消息的触发动作(cause_by)进行索引，以支持快速检索。通过这些方法，Memory 类能够高效地管理消息流，实现对消息的组织、查询和更新，适用于需要处理复杂消息交互的场景。
python
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 rsp2024-12-11 22:38:16.092 | INFO | metagpt.const:get_metagpt_package_root:21 - Package root set to d:\llm\metagpt
例子
以下是一些例子，帮助你理解 Memory 类及其方法的使用：
- 创建 Memory 对象并添加消息
python
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!]
- 批量添加消息
python
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!]
- 按角色获取消息
python
user_messages = memory.get_by_role("user")
print("Messages from user:", user_messages)Messages from user: [user: Hello!, user: How are you?]
- 按内容查找消息
python
matching_messages = memory.get_by_content("fine")
print("Messages containing 'fine':", matching_messages)Messages containing 'fine': [assistant: I'm fine, thank you!]
- 删除最新消息
python
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?]
- 清空存储
python
memory.clear()
print("Storage after clear:", memory.storage)Storage after clear: []
- 按触发动作查找消息
python
# 添加一些消息
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!]
如果有任何问题，欢迎在评论区提问。