【AI Agent系列】【MetaGPT多智能体学习】4. 基于MetaGPT的Team组件开发你的第一个智能体团队

class Team(BaseModel):
    """
    Team: Possesses one or more roles (agents), SOP (Standard Operating Procedures), and a env for instant messaging,
    dedicated to env any multi-agent activity, such as collaboratively writing executable code.
    团队：拥有一个或多个角色（代理人）、标准操作流程（SOP）和一个用于即时通讯的环境，致力于开展任何多代理活动，如协作编写可执行代码。
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)
    env: Optional[Environment] = None
    investment: float = Field(default=10.0)
    idea: str = Field(default="")

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

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

def run_project(self, idea, send_to: str = ""):
    """Run a project from publishing user requirement."""
    self.idea = idea
    # Human requirement.
    self.env.publish_message(
        Message(role="Human", content=idea, cause_by=UserRequirement, send_to=send_to or MESSAGE_ROUTE_TO_ALL),
        peekable=False,
    )

@serialize_decorator
async def run(self, n_round=3, idea="", send_to="", auto_archive=True):
    """Run company until target round or no money"""
    if idea:
        self.run_project(idea=idea, send_to=send_to)
    while n_round > 0:
        # self._save()
        n_round -= 1
        logger.debug(f"max {n_round=} left.")
        self._check_balance()
        await self.env.run()
    self.env.archive(auto_archive)
    return self.env.history

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

async def main(topic: str, n_round=3):
  ## 类比 Team 的 hire 函数添加 Roles
    classroom.add_roles([Student(), Teacher()])
  ## 类比 Team 的 run_project 函数往环境中写入用户消息
    classroom.publish_message(
        Message(role="Human", content=topic, cause_by=UserRequirement,
                send_to='' or MESSAGE_ROUTE_TO_ALL),
        peekable=False,
    )
  ## 类比 Team 的 run 函数控制循环次数
    while n_round > 0:
        # self._save()
        n_round -= 1
        logger.debug(f"max {n_round=} left.")
        await classroom.run()
    return classroom.history

def parse_code(rsp):
    pattern = r"```python(.*)```"
    match = re.search(pattern, rsp, re.DOTALL)
    code_text = match.group(1) if match else rsp
    return code_text
class SimpleWriteCode(Action):
    PROMPT_TEMPLATE: str = """
    Write a python function that can {instruction}.
    Return ```python your_code_here ```with NO other texts,
    your code:
    """
    name: str = "SimpleWriteCode"
    async def run(self, instruction: str):
        prompt = self.PROMPT_TEMPLATE.format(instruction=instruction)
        rsp = await self._aask(prompt)
        code_text = parse_code(rsp)
        return code_text
class SimpleCoder(Role):
    name: str = "Alice"
    profile: str = "SimpleCoder"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([UserRequirement])
        self.set_actions([SimpleWriteCode])

class SimpleWriteTest(Action):
    PROMPT_TEMPLATE: str = """
    Context: {context}
    Write {k} unit tests using pytest for the given function, assuming you have imported it.
    Return ```python your_code_here ```with NO other texts,
    your code:
    """
    name: str = "SimpleWriteTest"
    async def run(self, context: str, k: int = 3):
        prompt = self.PROMPT_TEMPLATE.format(context=context, k=k)
        rsp = await self._aask(prompt)
        code_text = parse_code(rsp)
        return code_text
class SimpleTester(Role):
    name: str = "Bob"
    profile: str = "SimpleTester"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([SimpleWriteTest])
        # self._watch([SimpleWriteCode])
        self._watch([SimpleWriteCode, SimpleWriteReview])  # feel free to try this too
        
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo
        # context = self.get_memories(k=1)[0].content # use the most recent memory as context
        context = self.get_memories()  # use all memories as context
        code_text = await todo.run(context, k=5)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))
        return msg

class SimpleWriteReview(Action):
    PROMPT_TEMPLATE: str = """
    Context: {context}
    Review the test cases and provide one critical comments:
    """
    name: str = "SimpleWriteReview"
    async def run(self, context: str):
        prompt = self.PROMPT_TEMPLATE.format(context=context)
        rsp = await self._aask(prompt)
        return rsp
class SimpleReviewer(Role):
    name: str = "Charlie"
    profile: str = "SimpleReviewer"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([SimpleWriteReview])
        self._watch([SimpleWriteTest])

async def main(
    idea: str = "write a function that calculates the product of a list",
    investment: float = 3.0,
    n_round: int = 5,
    add_human: bool = False,
):
    logger.info(idea)
    team = Team()
    team.hire(
        [
            SimpleCoder(),
            SimpleTester(),
            SimpleReviewer(is_human=add_human),
        ]
    )
    team.invest(investment=investment)
    team.run_project(idea)
    await team.run(n_round=n_round)
if __name__ == "__main__":
    fire.Fire(main)

"""
Filename: MetaGPT/examples/build_customized_multi_agents.py
Created Date: Wednesday, November 15th 2023, 7:12:39 pm
Author: garylin2099
"""
import re
import fire
from metagpt.actions import Action, UserRequirement
from metagpt.logs import logger
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.team import Team
def parse_code(rsp):
    pattern = r"```python(.*)```"
    match = re.search(pattern, rsp, re.DOTALL)
    code_text = match.group(1) if match else rsp
    return code_text
class SimpleWriteCode(Action):
    PROMPT_TEMPLATE: str = """
    Write a python function that can {instruction}.
    Return ```python your_code_here ```with NO other texts,
    your code:
    """
    name: str = "SimpleWriteCode"
    async def run(self, instruction: str):
        prompt = self.PROMPT_TEMPLATE.format(instruction=instruction)
        rsp = await self._aask(prompt)
        code_text = parse_code(rsp)
        return code_text
class SimpleCoder(Role):
    name: str = "Alice"
    profile: str = "SimpleCoder"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([UserRequirement])
        self.set_actions([SimpleWriteCode])
class SimpleWriteTest(Action):
    PROMPT_TEMPLATE: str = """
    Context: {context}
    Write {k} unit tests using pytest for the given function, assuming you have imported it.
    Return ```python your_code_here ```with NO other texts,
    your code:
    """
    name: str = "SimpleWriteTest"
    async def run(self, context: str, k: int = 3):
        prompt = self.PROMPT_TEMPLATE.format(context=context, k=k)
        rsp = await self._aask(prompt)
        code_text = parse_code(rsp)
        return code_text
class SimpleTester(Role):
    name: str = "Bob"
    profile: str = "SimpleTester"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([SimpleWriteTest])
        # self._watch([SimpleWriteCode])
        self._watch([SimpleWriteCode, SimpleWriteReview])  # feel free to try this too
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo
        # context = self.get_memories(k=1)[0].content # use the most recent memory as context
        context = self.get_memories()  # use all memories as context
        code_text = await todo.run(context, k=5)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))
        return msg
class SimpleWriteReview(Action):
    PROMPT_TEMPLATE: str = """
    Context: {context}
    Review the test cases and provide one critical comments:
    """
    name: str = "SimpleWriteReview"
    async def run(self, context: str):
        prompt = self.PROMPT_TEMPLATE.format(context=context)
        rsp = await self._aask(prompt)
        return rsp
class SimpleReviewer(Role):
    name: str = "Charlie"
    profile: str = "SimpleReviewer"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([SimpleWriteReview])
        self._watch([SimpleWriteTest])
async def main(
    idea: str = "write a function that calculates the product of a list",
    investment: float = 3.0,
    n_round: int = 5,
    add_human: bool = False,
):
    logger.info(idea)
    team = Team()
    team.hire(
        [
            SimpleCoder(),
            SimpleTester(),
            SimpleReviewer(is_human=add_human),
        ]
    )
    team.invest(investment=investment)
    team.run_project(idea)
    await team.run(n_round=n_round)
    # 最后这两句可以合成一句：await team.run(n_round=n_round, idea=idea)
if __name__ == "__main__":
    fire.Fire(main)