save code

metagpt/actions/write_code.py

@@ -11,6 +11,7 @@ from metagpt.const import WORKSPACE_ROOT
 from metagpt.logs import logger
 from metagpt.schema import Message
 from metagpt.utils.common import CodeParser
+from tenacity import retry, stop_after_attempt, wait_fixed
 
 PROMPT_TEMPLATE = """
 # Context
@@ -19,12 +20,11 @@ PROMPT_TEMPLATE = """
 NOTICE
 1. Role: You are an engineer; the main goal is to write PEP8 compliant, elegant, modular, easy to read and maintain Python 3.9 code (but you can also use other programming language)
 2. Requirement: Based on the context, implement one following code file, note to return only in code form, your code will be part of the entire project, so please implement complete, reliable, reusable code snippets
-3. Attention1: Use '##' to split sections, not '#', and '## <SECTION_NAME>' SHOULD WRITE BEFORE the code.
-4. Attention2: If there is any setting, ALWAYS SET A DEFAULT VALUE, ALWAYS USE STRONG TYPE AND EXPLICIT VARIABLE.
-5. Attention3: YOU MUST FOLLOW "Data structures and interface definitions". DONT CHANGE ANY DESIGN.
-6. Think before writing: What should be implemented and provided in this document?
-7. CAREFULLY CHECK THAT YOU DONT MISS ANY NECESSARY CLASS/FUNCTION IN THIS FILE.
-Attention: Use '##' to split sections, not '#', and '## <SECTION_NAME>' SHOULD WRITE BEFORE the code and triple quote.
+3. Attention1: If there is any setting, ALWAYS SET A DEFAULT VALUE, ALWAYS USE STRONG TYPE AND EXPLICIT VARIABLE.
+4. Attention2: YOU MUST FOLLOW "Data structures and interface definitions". DONT CHANGE ANY DESIGN.
+5. Think before writing: What should be implemented and provided in this document?
+6. CAREFULLY CHECK THAT YOU DONT MISS ANY NECESSARY CLASS/FUNCTION IN THIS FILE.
+7. Do not use public member functions that do not exist in your design.
 
 ## {filename}: Write code with triple quoto. Do your best to implement THIS ONLY ONE FILE. ONLY USE EXISTING API. IF NO API, IMPLEMENT IT.
 
@@ -43,7 +43,7 @@ class WriteCode(Action):
     def _is_invalid(self, filename):
         return any(i in filename for i in ["mp3", "wav"])
 
-    def _save(self, context, filename, code_rsp):
+    def _save(self, context, filename, code):
         # logger.info(filename)
         # logger.info(code_rsp)
         if self._is_invalid(filename):
@@ -57,16 +57,19 @@ class WriteCode(Action):
             ws_path = ws_path / ws_name
         code_path = ws_path / filename
         code_path.parent.mkdir(parents=True, exist_ok=True)
-        code = CodeParser.parse_code(block="", text=code_rsp)
         code_path.write_text(code)
         logger.info(f"Saving Code to {code_path}")
 
-    async def run(self, **kwargs):
-        prompt = PROMPT_TEMPLATE.format(**kwargs)
-        filename = kwargs['filename']
-        context = kwargs['context']
-        logger.info(f'Writing {filename}..')
+    @retry(stop=stop_after_attempt(2), wait=wait_fixed(1))
+    async def write_code(self, prompt):
         code_rsp = await self._aask(prompt)
+        code = CodeParser.parse_code(block="", text=code_rsp)
+        return code
+
+    async def run(self, context, filename):
+        prompt = PROMPT_TEMPLATE.format(context=context, filename=filename)
+        logger.info(f'Writing {filename}..')
+        code = await self.write_code(prompt)
         # code_rsp = await self._aask_v1(prompt, "code_rsp", OUTPUT_MAPPING)
-        self._save(context, filename, code_rsp)
-        return code_rsp
+        # self._save(context, filename, code)
+        return code

metagpt/actions/write_code_review.py

@@ -7,38 +7,52 @@
 """
 
 from metagpt.actions.action import Action
+from metagpt.logs import logger
+from metagpt.schema import Message
+from metagpt.utils.common import CodeParser
+from tenacity import retry, stop_after_attempt, wait_fixed
 
 PROMPT_TEMPLATE = """
-Please review the following code:
+# Context
+{context}
+
+# Code{filename}
+```
 {code}
+```
+-----
+NOTICE
+1. Role: You are a professional software engineer, and your main task is to review the code. You need to ensure that the code conforms to the PEP8 standards, is elegantly designed and modularized, easy to read and maintain, and is written in Python 3.9 (or in another programming language).
+2. Task 1: Based on the following context and code, conduct a code review and provide improvement suggestions.
+2. Task 2: Rewrite the code based on the improvement suggestions, ensure the code is complete and do not omit anything.
+3. Check 0: Is the code implemented as per the requirements?
+4. Check 1: Are there any issues with the code logic?
+5. Check 2: Does the existing code follow the "data structure and interface definition"?
+6. Check 3: Is the existing code complete and functional?
+7. Check 4: Does the code have unnecessary dependencies?
 
-The main aspects you need to focus on include but are not limited to the code structure, coding standards, possible errors, and improvement suggestions.
+## Code Review: Provide key, clear, concise, and specific code modification suggestions, up to 5.
+
+## {filename}: Write code with triple quotes. Do your utmost to optimize THIS SINGLE FILE. ONLY USE EXISTING API. IF NO API, IMPLEMENT IT.
+Ensure that the functionality of the rewritten code is consistent with the source code. 
 
-Please write your code review:
 """
 
 
 class WriteCodeReview(Action):
-    def __init__(self, name, context=None, llm=None):
+    def __init__(self, name="WriteCodeReview", context: list[Message] = None, llm=None):
         super().__init__(name, context, llm)
 
-    async def run(self, code):
-        """
-        Generate a code review for the given code.
+    @retry(stop=stop_after_attempt(2), wait=wait_fixed(1))
+    async def write_code(self, prompt):
+        code_rsp = await self._aask(prompt)
+        code = CodeParser.parse_code(block="", text=code_rsp)
+        return code
 
-        :param code: The code to be reviewed.
-        :type code: str
-        :return: The code review.
-        :rtype: str
-        """
-        # Set the context for the llm model
-        self.context = {"code": code}
-
-        # Generate the prompt
-        prompt = PROMPT_TEMPLATE.format(**self.context)
-
-        # Generate the code review
-        self.input_data = prompt
-        self.output_data = await self._aask(prompt)
-
-        return self.output_data
+    async def run(self, context, code, filename):
+        prompt = PROMPT_TEMPLATE.format(context=context, code=code, filename=filename)
+        logger.info(f'Code review {filename}..')
+        code = await self.write_code(prompt)
+        # code_rsp = await self._aask_v1(prompt, "code_rsp", OUTPUT_MAPPING)
+        # self._save(context, filename, code)
+        return code

metagpt/roles/engineer.py

@@ -14,7 +14,7 @@ from pathlib import Path
 from metagpt.const import WORKSPACE_ROOT
 from metagpt.logs import logger
 from metagpt.roles import Role
-from metagpt.actions import WriteCode, RunCode, DebugError, WriteTasks, WriteDesign
+from metagpt.actions import WriteCode, WriteCodeReview, WriteTasks, WriteDesign
 from metagpt.schema import Message
 from metagpt.utils.common import CodeParser
 from collections import OrderedDict
@@ -50,9 +50,12 @@ async def gather_ordered_k(coros, k) -> list:
 class Engineer(Role):
     def __init__(self, name="Alex", profile="Engineer", goal="Write elegant, readable, extensible, efficient code",
                  constraints="The code you write should conform to code standard like PEP8, be modular, easy to read and maintain",
-                 n_borg=1):
+                 n_borg=1, use_code_review=False):
         super().__init__(name, profile, goal, constraints)
         self._init_actions([WriteCode])
+        self.use_code_review = use_code_review
+        if self.use_code_review:
+            self._init_actions([WriteCode, WriteCodeReview])
         self._watch([WriteTasks])
         self.todos = []
         self.n_borg = n_borg
@@ -78,7 +81,8 @@ class Engineer(Role):
         if not msg:
             return WORKSPACE_ROOT / 'src'
         workspace = self.parse_workspace(msg)
-        return WORKSPACE_ROOT / workspace
+        # Codes are written in workspace/{package_name}/{package_name}
+        return WORKSPACE_ROOT / workspace / workspace
 
     def recreate_workspace(self):
         workspace = self.get_workspace()
@@ -132,6 +136,57 @@ class Engineer(Role):
             # logger.info(todo)
             # logger.info(code_rsp)
             # code = self.parse_code(code_rsp)
+            self.write_file(todo, code_rsp)
+            msg = Message(content=code_rsp, role=self.profile, cause_by=type(self._rc.todo))
+            self._rc.memory.add(msg)
+
+        logger.info(f'Done {self.get_workspace()} generating.')
+        msg = Message(content="all done.", role=self.profile, cause_by=type(self._rc.todo))
+        return msg
+
+    async def _act_sp_precision(self) -> Message:
+        """
+        # 从历史信息中挑选必须的信息，以减少prompt长度（人工经验总结）
+        1. ProductManager 分析和需求
+        2. Architect全部
+        3. ProjectManager全部
+        4. 是否需要其他代码？（目标是不需要。在任务拆分清楚后，根据设计思路，
+                            不需要其他代码也能够写清楚单个文件，
+                            如果不能则表示还需要在定义的更清晰，这个是代码能够写长的关键）
+        :return:
+        """
+        context = []
+        for msg in self._rc.history:
+            # todo: 需要再抽象一下
+            if msg.role in ["Architect", "Project Manager"]:
+                context.append(msg.content)
+            if msg.role in ["Product Manager"]:
+                tmp_msg = ""
+                if msg.instruct_content:
+                    temp_str = str(msg.instruct_content.dict().get(
+                        "Requirement Analysis"))
+                    tmp_msg += "## Requirement Analysis\n" + temp_str + "\n"
+                    temp_str = str(msg.instruct_content.dict().get("Requirement Pool"))
+                    tmp_msg += "## Requirement Pool\n" + temp_str + "\n"
+                context.append(tmp_msg)
+        context_str = "\n".join(context)
+        logger.debug(f'context: {context_str}')
+        for todo in self.todos:
+            code_rsp = await WriteCode().run(
+                context=context_str,
+                filename=todo
+            )
+            try:
+                code = await WriteCodeReview().run(
+                    context=context_str,
+                    code=code_rsp,
+                    filename=todo
+                )
+                code_rsp = code
+            except Exception as e:
+                logger.error("code review failed!", e)
+                pass
+            self.write_file(todo, code_rsp)
             msg = Message(content=code_rsp, role=self.profile, cause_by=type(self._rc.todo))
             self._rc.memory.add(msg)
 
@@ -140,4 +195,6 @@ class Engineer(Role):
         return msg
 
     async def _act(self) -> Message:
+        if self.use_code_review:
+            return await self._act_sp_precision()
         return await self._act_sp()

startup.py

@@ -6,23 +6,28 @@ from metagpt.software_company import SoftwareCompany
 from metagpt.roles import ProjectManager, ProductManager, Architect, Engineer
 
 
-async def startup(idea: str, investment: float = 3.0, n_round: int = 5):
+async def startup(idea: str, investment: float = 3.0, n_round: int = 5, code_review: bool = False):
     """Run a startup. Be a boss."""
     company = SoftwareCompany()
-    company.hire([ProductManager(), Architect(), ProjectManager(), Engineer(n_borg=5)])
+    company.hire([ProductManager(),
+                  Architect(),
+                  ProjectManager(),
+                  Engineer(n_borg=5, use_code_review=code_review)])
     company.invest(investment)
     company.start_project(idea)
     await company.run(n_round=n_round)
 
 
-def main(idea: str, investment: float = 3.0, n_round: int = 5):
+def main(idea: str, investment: float = 3.0, n_round: int = 5, code_review: bool = False):
     """
     We are a software startup comprised of AI. By investing in us, you are empowering a future filled with limitless possibilities.
     :param idea: Your innovative idea, such as "Creating a snake game."
     :param investment: As an investor, you have the opportunity to contribute a certain dollar amount to this AI company.
+    :param n_round:
+    :param code_review: Whether to use code review.
     :return:
     """
-    asyncio.run(startup(idea, investment, n_round))
+    asyncio.run(startup(idea, investment, n_round, code_review))
 
 
 if __name__ == '__main__':