MetaGPT

1
#!/usr/bin/env python
2
# -*- coding: utf-8 -*-
3
"""
4
@Time    : 2023/5/12 00:30
5
@Author  : alexanderwu
6
@File    : team.py
7
@Modified By: mashenquan, 2023/11/27. Add an archiving operation after completing the project, as specified in
8
        Section 2.2.3.3 of RFC 135.
9
"""
10

11
import warnings
12
from pathlib import Path
13
from typing import Any, Optional
14

15
from pydantic import BaseModel, ConfigDict, Field
16

17
from metagpt.actions import UserRequirement
18
from metagpt.const import MESSAGE_ROUTE_TO_ALL, SERDESER_PATH
19
from metagpt.context import Context
20
from metagpt.environment import Environment
21
from metagpt.logs import logger
22
from metagpt.roles import Role
23
from metagpt.schema import Message
24
from metagpt.utils.common import (
25
    NoMoneyException,
26
    read_json_file,
27
    serialize_decorator,
28
    write_json_file,
29
)
30

31

32
class Team(BaseModel):
33
    """
34
    Team: Possesses one or more roles (agents), SOP (Standard Operating Procedures), and a env for instant messaging,
35
    dedicated to env any multi-agent activity, such as collaboratively writing executable code.
36
    """
37

38
    model_config = ConfigDict(arbitrary_types_allowed=True)
39

40
    env: Optional[Environment] = None
41
    investment: float = Field(default=10.0)
42
    idea: str = Field(default="")
43

44
    def __init__(self, context: Context = None, **data: Any):
45
        super(Team, self).__init__(**data)
46
        ctx = context or Context()
47
        if not self.env:
48
            self.env = Environment(context=ctx)
49
        else:
50
            self.env.context = ctx  # The `env` object is allocated by deserialization
51
        if "roles" in data:
52
            self.hire(data["roles"])
53
        if "env_desc" in data:
54
            self.env.desc = data["env_desc"]
55

56
    def serialize(self, stg_path: Path = None):
57
        stg_path = SERDESER_PATH.joinpath("team") if stg_path is None else stg_path
58
        team_info_path = stg_path.joinpath("team.json")
59
        serialized_data = self.model_dump()
60
        serialized_data["context"] = self.env.context.serialize()
61

62
        write_json_file(team_info_path, serialized_data)
63

64
    @classmethod
65
    def deserialize(cls, stg_path: Path, context: Context = None) -> "Team":
66
        """stg_path = ./storage/team"""
67
        # recover team_info
68
        team_info_path = stg_path.joinpath("team.json")
69
        if not team_info_path.exists():
70
            raise FileNotFoundError(
71
                "recover storage meta file `team.json` not exist, " "not to recover and please start a new project."
72
            )
73

74
        team_info: dict = read_json_file(team_info_path)
75
        ctx = context or Context()
76
        ctx.deserialize(team_info.pop("context", None))
77
        team = Team(**team_info, context=ctx)
78
        return team
79

80
    def hire(self, roles: list[Role]):
81
        """Hire roles to cooperate"""
82
        self.env.add_roles(roles)
83

84
    @property
85
    def cost_manager(self):
86
        """Get cost manager"""
87
        return self.env.context.cost_manager
88

89
    def invest(self, investment: float):
90
        """Invest company. raise NoMoneyException when exceed max_budget."""
91
        self.investment = investment
92
        self.cost_manager.max_budget = investment
93
        logger.info(f"Investment: ${investment}.")
94

95
    def _check_balance(self):
96
        if self.cost_manager.total_cost >= self.cost_manager.max_budget:
97
            raise NoMoneyException(self.cost_manager.total_cost, f"Insufficient funds: {self.cost_manager.max_budget}")
98

99
    def run_project(self, idea, send_to: str = ""):
100
        """Run a project from publishing user requirement."""
101
        self.idea = idea
102

103
        # Human requirement.
104
        self.env.publish_message(
105
            Message(role="Human", content=idea, cause_by=UserRequirement, send_to=send_to or MESSAGE_ROUTE_TO_ALL),
106
            peekable=False,
107
        )
108

109
    def start_project(self, idea, send_to: str = ""):
110
        """
111
        Deprecated: This method will be removed in the future.
112
        Please use the `run_project` method instead.
113
        """
114
        warnings.warn(
115
            "The 'start_project' method is deprecated and will be removed in the future. "
116
            "Please use the 'run_project' method instead.",
117
            DeprecationWarning,
118
            stacklevel=2,
119
        )
120
        return self.run_project(idea=idea, send_to=send_to)
121

122
    @serialize_decorator
123
    async def run(self, n_round=3, idea="", send_to="", auto_archive=True):
124
        """Run company until target round or no money"""
125
        if idea:
126
            self.run_project(idea=idea, send_to=send_to)
127

128
        while n_round > 0:
129
            n_round -= 1
130
            self._check_balance()
131
            await self.env.run()
132

133
            logger.debug(f"max {n_round=} left.")
134
        self.env.archive(auto_archive)
135
        return self.env.history
136