feat: merge geekan:main

metagpt/environment.py

@@ -12,29 +12,85 @@
     functionality is to be consolidated into the `Environment` class.
 """
 import asyncio
+from pathlib import Path
 from typing import Iterable, Set
 
 from pydantic import BaseModel, Field
 
+from metagpt.config import CONFIG
 from metagpt.logs import logger
-from metagpt.roles import Role
+from metagpt.roles.role import Role, role_subclass_registry
 from metagpt.schema import Message
-from metagpt.utils.common import is_subscribed
+from metagpt.utils.common import is_subscribed, read_json_file, write_json_file
 
 
 class Environment(BaseModel):
     """环境，承载一批角色，角色可以向环境发布消息，可以被其他角色观察到
     Environment, hosting a batch of roles, roles can publish messages to the environment, and can be observed by other roles
-
     """
 
     roles: dict[str, Role] = Field(default_factory=dict)
     members: dict[Role, Set] = Field(default_factory=dict)
-    history: str = Field(default="")  # For debug
+    history: str = ""  # For debug
 
     class Config:
         arbitrary_types_allowed = True
 
+    def __init__(self, **kwargs):
+        roles = []
+        for role_key, role in kwargs.get("roles", {}).items():
+            current_role = kwargs["roles"][role_key]
+            if isinstance(current_role, dict):
+                item_class_name = current_role.get("builtin_class_name", None)
+                for name, subclass in role_subclass_registry.items():
+                    registery_class_name = subclass.__fields__["builtin_class_name"].default
+                    if item_class_name == registery_class_name:
+                        current_role = subclass(**current_role)
+                        break
+                kwargs["roles"][role_key] = current_role
+                roles.append(current_role)
+        super().__init__(**kwargs)
+
+        self.add_roles(roles)  # add_roles again to init the Role.set_env
+
+    def serialize(self, stg_path: Path):
+        roles_path = stg_path.joinpath("roles.json")
+        roles_info = []
+        for role_key, role in self.roles.items():
+            roles_info.append(
+                {
+                    "role_class": role.__class__.__name__,
+                    "module_name": role.__module__,
+                    "role_name": role.name,
+                    "role_sub_tags": list(self.members.get(role)),
+                }
+            )
+            role.serialize(stg_path=stg_path.joinpath(f"roles/{role.__class__.__name__}_{role.name}"))
+        write_json_file(roles_path, roles_info)
+
+        history_path = stg_path.joinpath("history.json")
+        write_json_file(history_path, {"content": self.history})
+
+    @classmethod
+    def deserialize(cls, stg_path: Path) -> "Environment":
+        """stg_path: ./storage/team/environment/"""
+        roles_path = stg_path.joinpath("roles.json")
+        roles_info = read_json_file(roles_path)
+        roles = []
+        for role_info in roles_info:
+            # role stored in ./environment/roles/{role_class}_{role_name}
+            role_path = stg_path.joinpath(f"roles/{role_info.get('role_class')}_{role_info.get('role_name')}")
+            role = Role.deserialize(role_path)
+            roles.append(role)
+
+        history = read_json_file(stg_path.joinpath("history.json"))
+        history = history.get("content")
+
+        environment = Environment(**{"history": history})
+        environment.add_roles(roles)
+
+        return environment
+
     def add_role(self, role: Role):
         """增加一个在当前环境的角色
         Add a role in the current environment
@@ -111,3 +167,8 @@ class Environment(BaseModel):
     def set_subscription(self, obj, tags):
         """Set the labels for message to be consumed by the object"""
         self.members[obj] = tags
+
+    @staticmethod
+    def archive(auto_archive=True):
+        if auto_archive and CONFIG.git_repo:
+            CONFIG.git_repo.archive()