add mg ser&deser

This commit is contained in:
better629 2023-11-28 09:29:00 +08:00
parent 9e97e6d3d0
commit d99b4c62e3
16 changed files with 641 additions and 31 deletions

View file

@ -11,3 +11,20 @@ from metagpt.actions import Action, WritePRD, WriteTest
def test_action_repr():
actions = [Action(), WriteTest(), WritePRD()]
assert "WriteTest" in str(actions)
def test_action_serdes():
action_info = WriteTest.ser_class()
assert action_info["action_class"] == "WriteTest"
action_class = Action.deser_class(action_info)
assert action_class == WriteTest
def test_action_class_serdes():
name = "write test"
action_info = WriteTest(name=name).serialize()
assert action_info["name"] == name
action = Action.deserialize(action_info)
assert action.name == name

View file

@ -0,0 +1,42 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Desc : unittest of memory
from pathlib import Path
from metagpt.schema import Message
from metagpt.memory.memory import Memory
from metagpt.actions.action_output import ActionOutput
from metagpt.actions.design_api import WriteDesign
from metagpt.actions.add_requirement import BossRequirement
serdes_path = Path(__file__).absolute().parent.joinpath("../../data/serdes_storage")
def test_memory_serdes():
msg1 = Message(role="User",
content="write a 2048 game",
cause_by=BossRequirement)
out_mapping = {"field1": (list[str], ...)}
out_data = {"field1": ["field1 value1", "field1 value2"]}
ic_obj = ActionOutput.create_model_class("system_design", out_mapping)
msg2 = Message(role="Architect",
instruct_content=ic_obj(**out_data),
content="system design content",
cause_by=WriteDesign)
memory = Memory()
memory.add_batch([msg1, msg2])
stg_path = serdes_path.joinpath("team/environment")
memory.serialize(stg_path)
assert stg_path.joinpath("memory.json").exists()
new_memory = Memory.deserialize(stg_path)
assert new_memory.count() == 2
new_msg2 = new_memory.get(1)[0]
assert new_msg2.instruct_content.field1 == ["field1 value1", "field1 value2"]
assert new_msg2.cause_by == WriteDesign
stg_path.joinpath("memory.json").unlink()

View file

@ -0,0 +1,85 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Desc : unittest of Role
from pathlib import Path
import shutil
import pytest
from metagpt.roles.role import Role, RoleReactMode
from metagpt.actions.action import Action
from metagpt.schema import Message
from metagpt.actions.add_requirement import BossRequirement
from metagpt.roles.product_manager import ProductManager
serdes_path = Path(__file__).absolute().parent.joinpath("../../data/serdes_storage")
def test_role_serdes():
stg_path_prefix = serdes_path.joinpath("team/environment/roles/")
shutil.rmtree(serdes_path.joinpath("team"), ignore_errors=True)
pm = ProductManager()
role_tag = f"{pm.__class__.__name__}_{pm.name}"
stg_path = stg_path_prefix.joinpath(role_tag)
pm.serialize(stg_path)
assert stg_path.joinpath("actions/actions_info.json").exists()
new_pm = Role.deserialize(stg_path)
assert new_pm.name == pm.name
assert len(new_pm.get_memories(1)) == 0
class ActionOK(Action):
async def run(self, messages: list["Message"]):
return "ok"
class ActionRaise(Action):
async def run(self, messages: list["Message"]):
raise RuntimeError("parse error")
class RoleA(Role):
def __init__(self,
name: str = "RoleA",
profile: str = "Role A",
goal: str = "",
constraints: str = ""):
super(RoleA, self).__init__(name=name, profile=profile, goal=goal, constraints=constraints)
self._init_actions([ActionOK, ActionRaise])
self._watch([BossRequirement])
self._rc.react_mode = RoleReactMode.BY_ORDER
async def run(self, message: "Message" = None, stg_path: str = None):
try:
await super(RoleA, self).run(message)
except Exception as exp:
print("exp ", exp)
self.serialize(stg_path)
@pytest.mark.asyncio
async def test_role_serdes_interrupt():
role_a = RoleA()
shutil.rmtree(serdes_path.joinpath("team"), ignore_errors=True)
stg_path = serdes_path.joinpath(f"team/environment/roles/{role_a.__class__.__name__}_{role_a.name}")
await role_a.run(
message=Message(content="demo", cause_by=BossRequirement),
stg_path=stg_path
)
assert role_a._rc.memory.count() == 2
assert stg_path.joinpath("actions/todo.json").exists()
new_role_a: Role = Role.deserialize(stg_path)
assert new_role_a._rc.state == 1
await role_a.run(
message=Message(content="demo", cause_by=BossRequirement),
stg_path=stg_path
)

View file

@ -7,13 +7,18 @@
"""
import pytest
from pathlib import Path
import shutil
from metagpt.actions import BossRequirement
from metagpt.environment import Environment
from metagpt.logs import logger
from metagpt.manager import Manager
from metagpt.roles import Architect, ProductManager, Role
from metagpt.schema import Message
from tests.metagpt.roles.test_role import RoleA
serdes_path = Path(__file__).absolute().parent.joinpath("../data/serdes_storage")
@pytest.fixture
@ -36,21 +41,29 @@ def test_get_roles(env: Environment):
assert roles == {role1.profile: role1, role2.profile: role2}
def test_set_manager(env: Environment):
manager = Manager()
env.set_manager(manager)
assert env.manager == manager
@pytest.mark.asyncio
async def test_publish_and_process_message(env: Environment):
product_manager = ProductManager("Alice", "Product Manager", "做AI Native产品", "资源有限")
architect = Architect("Bob", "Architect", "设计一个可用、高效、较低成本的系统,包括数据结构与接口", "资源有限,需要节省成本")
env.add_roles([product_manager, architect])
env.set_manager(Manager())
env.publish_message(Message(role="BOSS", content="需要一个基于LLM做总结的搜索引擎", cause_by=BossRequirement))
await env.run(k=2)
logger.info(f"{env.history=}")
assert len(env.history) > 10
def test_environment_serdes():
environment = Environment()
role_a = RoleA()
shutil.rmtree(serdes_path.joinpath("team"), ignore_errors=True)
stg_path = serdes_path.joinpath("team/environment")
environment.add_role(role_a)
environment.serialize(stg_path)
new_env: Environment = Environment()
new_env.deserialize(stg_path)
assert len(new_env.roles) == 1

View file

@ -5,7 +5,11 @@
@Author : alexanderwu
@File : test_schema.py
"""
from metagpt.schema import AIMessage, Message, SystemMessage, UserMessage
from metagpt.actions.action_output import ActionOutput
from metagpt.actions.write_code import WriteCode
from metagpt.utils.serialize import serialize_general_message, deserialize_general_message
def test_messages():
@ -19,3 +23,41 @@ def test_messages():
text = str(msgs)
roles = ['user', 'system', 'assistant', 'QA']
assert all([i in text for i in roles])
def test_message_serdes():
out_mapping = {"field3": (str, ...), "field4": (list[str], ...)}
out_data = {"field3": "field3 value3", "field4": ["field4 value1", "field4 value2"]}
ic_obj = ActionOutput.create_model_class("code", out_mapping)
message = Message(
content="code",
instruct_content=ic_obj(**out_data),
role="engineer",
cause_by=WriteCode
)
message_dict = serialize_general_message(message)
assert message_dict["cause_by"] == {"action_class": "WriteCode"}
assert message_dict["instruct_content"] == {
"class": "code",
"mapping": {
"field3": "(<class 'str'>, Ellipsis)",
"field4": "(list[str], Ellipsis)"
},
"value": {
"field3": "field3 value3",
"field4": ["field4 value1", "field4 value2"]
}
}
new_message = deserialize_general_message(message_dict)
assert new_message.content == message.content
assert new_message.instruct_content == message.instruct_content
assert new_message.cause_by == message.cause_by
assert new_message.instruct_content.field3 == out_data["field3"]
message = Message(content="code")
message_dict = serialize_general_message(message)
new_message = deserialize_general_message(message_dict)
assert new_message.instruct_content is None
assert new_message.cause_by == ""

View file

@ -0,0 +1,27 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Desc : unittest of team
from pathlib import Path
import shutil
from metagpt.team import Team
from tests.metagpt.roles.test_role import RoleA
serdes_path = Path(__file__).absolute().parent.joinpath("../data/serdes_storage")
def test_team_serdes():
company = Team()
company.hire([RoleA()])
stg_path = serdes_path.joinpath("team")
shutil.rmtree(stg_path, ignore_errors=True)
company.serialize(stg_path=stg_path)
new_company = Team()
new_company.deserialize(stg_path)
assert len(new_company.environment.roles) == 1