From 63a41ba81ded62d306fe9b406dfaaefdc635a6d7 Mon Sep 17 00:00:00 2001 From: hongjiongteng Date: Tue, 18 Jun 2024 21:35:46 +0800 Subject: [PATCH] add code_review tool to engineer2 to enhance code effectiveness --- metagpt/actions/di/rewrite_code.py | 80 +++++++++++++++++++++++++ metagpt/prompts/di/engineer2.py | 9 +++ metagpt/roles/di/engineer2.py | 13 +++- tests/metagpt/roles/di/run_engineer2.py | 8 +-- 4 files changed, 105 insertions(+), 5 deletions(-) create mode 100644 metagpt/actions/di/rewrite_code.py diff --git a/metagpt/actions/di/rewrite_code.py b/metagpt/actions/di/rewrite_code.py new file mode 100644 index 000000000..a8586d1a6 --- /dev/null +++ b/metagpt/actions/di/rewrite_code.py @@ -0,0 +1,80 @@ +from tenacity import retry, stop_after_attempt, wait_random_exponential + +from metagpt.actions.action import Action +from metagpt.actions.write_code_review import ( + EXAMPLE_AND_INSTRUCTION, + FORMAT_EXAMPLE, + PROMPT_TEMPLATE, + REWRITE_CODE_TEMPLATE, +) +from metagpt.tools.tool_registry import register_tool +from metagpt.utils.common import CodeParser, aread, awrite + + +@register_tool(tags=["RewriteCode"], include_functions=["run"]) +class RewriteCode(Action): + """Accordding design doc and task doc to review the code, to make the complete and correct code.""" + + name: str = "RewriteCode" + + async def run(self, code_path: str, design_doc: str = "", task_doc: str = "", code_review_k_times: int = 2) -> str: + """Reviews the provided code based on the accompanying design and task documentation, return the complete and correct code. + + Read the code from `code_path`, and write the final code to `code_path`. + + Args: + code_path (str): The file path of the code snippet to be reviewed. This should be a string containing the path to the source code file. + design_doc (str): The design document associated with the code. This should describe the system architecture, used in the code. It helps provide context for the review process. + task_doc (str): The task document describing what the code is intended to accomplish. This should outline the functional requirements or objectives of the code. + code_review_k_times (int, optional): The number of iterations for reviewing and potentially rewriting the code. Defaults to 2. + + Returns: + str: The potentially corrected or approved code after review. + + Example Usage: + # Example of how to call the run method with a code snippet and documentation + await WriteCodeReview().run( + code_path="/tmp/game.js", + design_doc='{"Implementation approach":"We will implement the 2048 game using plain JavaScript and HTML, ensuring no frameworks are used. The game logic will handle tile movements, merging, and game state updates. The UI will be simple and clean, with a responsive design to fit different screen sizes. We will use CSS for styling and ensure the game is playable with keyboard arrow keys. The game will display the current score, have a restart button, and show a game over message when no more moves are possible.","File list":["index.html","style.css","script.js"],"Data structures and interfaces":"\nclassDiagram\n class Game {\n -grid: int[][]\n -score: int\n +init(): void\n +move(direction: str): void\n +merge(direction: str): void\n +isGameOver(): bool\n +restart(): void\n }\n class UI {\n -game: Game\n +init(): void\n +update(): void\n +showGameOver(): void\n +bindEvents(): void\n }\n Game --> UI\n","Program call flow":"\nsequenceDiagram\n participant U as UI\n participant G as Game\n U->>G: init()\n G-->>U: return\n U->>U: bindEvents()\n U->>G: move(direction)\n G->>G: merge(direction)\n G->>U: update()\n U->>U: update()\n U->>G: isGameOver()\n G-->>U: return bool\n alt Game Over\n U->>U: showGameOver()\n end\n U->>G: restart()\n G-->>U: return\n","Anything UNCLEAR":"Clarify if there are any specific design preferences or additional features required beyond the basic 2048 game functionality."}', + task_doc='{"Required packages":["No third-party dependencies required"],"Required Other language third-party packages":["No third-party dependencies required"],"Logic Analysis":[["script.js","Contains Game and UI classes, and their methods: init, move, merge, isGameOver, restart, update, showGameOver, bindEvents"],["index.html","Contains the HTML structure for the game UI"],["style.css","Contains the CSS styles for the game UI"]],"Task list":["index.html","style.css","script.js"],"Full API spec":"","Shared Knowledge":"The `script.js` file will contain the core game logic and UI handling. The `index.html` file will provide the structure for the game, and `style.css` will handle the styling.","Anything UNCLEAR":"Clarify if there are any specific design preferences or additional features required beyond the basic 2048 game functionality."}' + ) + """ + code = await aread(code_path) + + context = "\n".join( + [ + "## System Design\n" + design_doc + "\n", + "## Task\n" + task_doc + "\n", + ] + ) + + for _ in range(code_review_k_times): + context_prompt = PROMPT_TEMPLATE.format(context=context, code=code, filename=code_path) + cr_prompt = EXAMPLE_AND_INSTRUCTION.format( + format_example=FORMAT_EXAMPLE.format(filename=code_path), + ) + result, rewrited_code = await self.write_code_review_and_rewrite( + context_prompt, cr_prompt, filename=code_path + ) + + if "LBTM" in result: + code = rewrited_code + elif "LGTM" in result: + break + + await awrite(filename=code_path, data=code) + + return code + + @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6)) + async def write_code_review_and_rewrite(self, context_prompt: str, cr_prompt: str, filename: str): + cr_rsp = await self._aask(context_prompt + cr_prompt) + result = CodeParser.parse_block("Code Review Result", cr_rsp) + if "LGTM" in result: + return result, None + + # if LBTM, rewrite code + rewrite_prompt = f"{context_prompt}\n{cr_rsp}\n{REWRITE_CODE_TEMPLATE.format(filename=filename)}" + code_rsp = await self._aask(rewrite_prompt) + code = CodeParser.parse_code(block="", text=code_rsp) + return result, code diff --git a/metagpt/prompts/di/engineer2.py b/metagpt/prompts/di/engineer2.py index 4fd52e320..8e2722b49 100644 --- a/metagpt/prompts/di/engineer2.py +++ b/metagpt/prompts/di/engineer2.py @@ -5,5 +5,14 @@ EXTRA_INSTRUCTION = """ 5. Take on ONE task and write ONE code file in each response. DON'T attempt all tasks in one response. 6. When not specified, you should write files in a folder named "src". If you know the project path, then write in a "src" folder under the project path. 7. When provided system design or project schedule, read them first, then adhere to them in your implementation. +8. Write at most one file per task, do your best to implement THE ONLY ONE FILE. CAREFULLY CHECK THAT YOU DONT MISS ANY NECESSARY CLASS/FUNCTION IN THIS FILE. +9. COMPLETE CODE: Your code will be part of the entire project, so please implement complete, reliable, reusable code snippets. +10. When provided system design, YOU MUST FOLLOW "Data structures and interfaces". DONT CHANGE ANY DESIGN. Do not use public member functions that do not exist in your design. +11. Write out EVERY CODE DETAIL, DON'T LEAVE TODO. +12. Do not use Editor to find start_line and end_line, just rewrite the file with the all complete code. +13. Revise task is to use RewriteCode.run to correct code. +14. At the end of the plan, add a Revise task for each file; for example, if there are three files, add three Revise tasks. """ + + ENGINEER2_INSTRUCTION = ROLE_INSTRUCTION + EXTRA_INSTRUCTION.strip() diff --git a/metagpt/roles/di/engineer2.py b/metagpt/roles/di/engineer2.py index e013ef09e..2720364ef 100644 --- a/metagpt/roles/di/engineer2.py +++ b/metagpt/roles/di/engineer2.py @@ -1,5 +1,6 @@ from __future__ import annotations +from metagpt.actions.di.rewrite_code import RewriteCode from metagpt.prompts.di.engineer2 import ENGINEER2_INSTRUCTION from metagpt.roles.di.role_zero import RoleZero @@ -10,4 +11,14 @@ class Engineer2(RoleZero): goal: str = "Take on game, app, and web development" instruction: str = ENGINEER2_INSTRUCTION - tools: str = ["Plan", "Editor:write,read,write_content", "RoleZero"] + tools: str = ["Plan", "Editor:write,read,write_content", "RoleZero", "RewriteCode"] + + def _update_tool_execution(self): + rewrite_code = RewriteCode() + + self.tool_execution_map.update( + { + "RewriteCode.run": rewrite_code.run, + "RewriteCode": rewrite_code.run, + } + ) diff --git a/tests/metagpt/roles/di/run_engineer2.py b/tests/metagpt/roles/di/run_engineer2.py index 4e948bad7..e5ae74485 100644 --- a/tests/metagpt/roles/di/run_engineer2.py +++ b/tests/metagpt/roles/di/run_engineer2.py @@ -67,18 +67,18 @@ Create a 2048 game, follow the design doc and task doc. Write your code under /U After writing all codes, write a code review for the codes, make improvement or adjustment based on the review. Notice: You MUST implement the full code, don't leave comment without implementation! Design doc: -{TASK_DOC_2048} -Task doc: {DESIGN_DOC_2048} +Task doc: +{TASK_DOC_2048} """ GAME_REQ_SNAKE = f""" Create a snake game, follow the design doc and task doc. Write your code under /Users/gary/Files/temp/workspace/snake_game/src. After writing all codes, write a code review for the codes, make improvement or adjustment based on the review. Notice: You MUST implement the full code, don't leave comment without implementation! Design doc: -{TASK_DOC_SNAKE} -Task doc: {DESIGN_DOC_SNAKE} +Task doc: +{TASK_DOC_SNAKE} """ GAME_REQ_2048_NO_DOC = """ Create a 2048 game with pygame. Write your code under /Users/gary/Files/temp/workspace/2048_game/src.