1  """
  2  Filename: MetaGPT/examples/build_customized_multi_agents.py
  3  Created Date: Wednesday, November 15th 2023, 7:12:39 pm
  4  Author: garylin2099
  5  """
  6  import re
  7  
  8  import fire
  9  
 10  from metagpt.actions import Action, UserRequirement
 11  from metagpt.logs import logger
 12  from metagpt.roles import Role
 13  from metagpt.schema import Message
 14  from metagpt.team import Team
 15  
 16  
 17  def parse_code(rsp):
 18      pattern = r"```python(.*)```"
 19      match = re.search(pattern, rsp, re.DOTALL)
 20      code_text = match.group(1) if match else rsp
 21      return code_text
 22  
 23  
 24  class SimpleWriteCode(Action):
 25      PROMPT_TEMPLATE: str = """
 26      Write a python function that can {instruction}.
 27      Return ```python your_code_here ``` with NO other texts,
 28      your code:
 29      """
 30      name: str = "SimpleWriteCode"
 31  
 32      async def run(self, instruction: str):
 33          prompt = self.PROMPT_TEMPLATE.format(instruction=instruction)
 34  
 35          rsp = await self._aask(prompt)
 36  
 37          code_text = parse_code(rsp)
 38  
 39          return code_text
 40  
 41  
 42  class SimpleCoder(Role):
 43      name: str = "Alice"
 44      profile: str = "SimpleCoder"
 45  
 46      def __init__(self, **kwargs):
 47          super().__init__(**kwargs)
 48          self._watch([UserRequirement])
 49          self.set_actions([SimpleWriteCode])
 50  
 51  
 52  class SimpleWriteTest(Action):
 53      PROMPT_TEMPLATE: str = """
 54      Context: {context}
 55      Write {k} unit tests using pytest for the given function, assuming you have imported it.
 56      Return ```python your_code_here ``` with NO other texts,
 57      your code:
 58      """
 59  
 60      name: str = "SimpleWriteTest"
 61  
 62      async def run(self, context: str, k: int = 3):
 63          prompt = self.PROMPT_TEMPLATE.format(context=context, k=k)
 64  
 65          rsp = await self._aask(prompt)
 66  
 67          code_text = parse_code(rsp)
 68  
 69          return code_text
 70  
 71  
 72  class SimpleTester(Role):
 73      name: str = "Bob"
 74      profile: str = "SimpleTester"
 75  
 76      def __init__(self, **kwargs):
 77          super().__init__(**kwargs)
 78          self.set_actions([SimpleWriteTest])
 79          # self._watch([SimpleWriteCode])
 80          self._watch([SimpleWriteCode, SimpleWriteReview])  # feel free to try this too
 81  
 82      async def _act(self) -> Message:
 83          logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
 84          todo = self.rc.todo
 85  
 86          # context = self.get_memories(k=1)[0].content # use the most recent memory as context
 87          context = self.get_memories()  # use all memories as context
 88  
 89          code_text = await todo.run(context, k=5)  # specify arguments
 90          msg = Message(content=code_text, role=self.profile, cause_by=type(todo))
 91  
 92          return msg
 93  
 94  
 95  class SimpleWriteReview(Action):
 96      PROMPT_TEMPLATE: str = """
 97      Context: {context}
 98      Review the test cases and provide one critical comments:
 99      """
100  
101      name: str = "SimpleWriteReview"
102  
103      async def run(self, context: str):
104          prompt = self.PROMPT_TEMPLATE.format(context=context)
105  
106          rsp = await self._aask(prompt)
107  
108          return rsp
109  
110  
111  class SimpleReviewer(Role):
112      name: str = "Charlie"
113      profile: str = "SimpleReviewer"
114  
115      def __init__(self, **kwargs):
116          super().__init__(**kwargs)
117          self.set_actions([SimpleWriteReview])
118          self._watch([SimpleWriteTest])
119  
120  
121  async def main(
122      idea: str = "write a function that calculates the product of a list",
123      investment: float = 3.0,
124      n_round: int = 5,
125      add_human: bool = False,
126  ):
127      logger.info(idea)
128  
129      team = Team()
130      team.hire(
131          [
132              SimpleCoder(),
133              SimpleTester(),
134              SimpleReviewer(is_human=add_human),
135          ]
136      )
137  
138      team.invest(investment=investment)
139      team.run_project(idea)
140      await team.run(n_round=n_round)
141  
142  
143  if __name__ == "__main__":
144      fire.Fire(main)