reuse code

metagpt/actions/write_analysis_code.py

@@ -25,13 +25,15 @@ class BaseWriteAnalysisCode(Action):
 
 class WriteCodeByGenerate(BaseWriteAnalysisCode):
     """Write code fully by generation"""
+    DEFAULT_SYSTEM_MSG = """You are Code Interpreter, a world-class programmer that can complete any goal by executing code. Strictly follow the plan and generate code step by step. Each step of the code will be executed on the user's machine, and the user will provide the code execution results to you.**Notice: Use !pip install to install missing packages.**"""
+    REUSE_CODE_INSTRUCTION = """ATTENTION: DONT include codes from previous steps in your current code block, include new codes only, DONT repeat codes!"""
 
     def __init__(self, name: str = "", context=None, llm=None) -> str:
         super().__init__(name, context, llm)
 
     def process_msg(self, prompt: Union[str, List[Dict], Message, List[Message]], system_msg: str = None):
         # Reference: https://github.com/KillianLucas/open-interpreter/blob/v0.1.4/interpreter/system_message.txt
-        default_system_msg = """You are Code Interpreter, a world-class programmer that can complete any goal by executing code. Strictly follow the plan and generate code step by step. Each step of the code will be executed on the user's machine, and the user will provide the code execution results to you.**Notice: The code for the next step depends on the code for the previous step. Reuse existing code directly. Use !pip install to install missing packages.**"""
+        default_system_msg = system_msg or self.DEFAULT_SYSTEM_MSG
         # 全部转成list
         if not isinstance(prompt, list):
             prompt = [prompt]
@@ -59,6 +61,7 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
     async def run(
         self, context: [List[Message]], plan: Plan = None, task_guide: str = "", system_msg: str = None, **kwargs
     ) -> str:
+        context.append(Message(content=self.REUSE_CODE_INSTRUCTION, role="user"))
         prompt = self.process_msg(context, system_msg)
         code_content = await self.llm.aask_code(prompt, **kwargs)
         return code_content["code"]

metagpt/roles/ml_engineer.py

@@ -22,7 +22,7 @@ STRUCTURAL_CONTEXT = """
 """
 
 def truncate(result: str, keep_len: int = 1000) -> str:
-    desc = """I truncated the result to only keep the last 1000 characters\n"""
+    desc = "Truncated to show only the last 1000 characters\n"
     if result.startswith(desc):
         result = result[-len(desc):]
 
@@ -38,19 +38,22 @@ class AskReview(Action):
 
     async def run(self, context: List[Message], plan: Plan = None):
         logger.info("Current overall plan:")
-        logger.info("\n".join([f"{task.task_id}: {task.instruction}, is_finished: {task.is_finished}" for task in plan.tasks]))
+        logger.info(
+            "\n".join([f"{task.task_id}: {task.instruction}, is_finished: {task.is_finished}" for task in plan.tasks])
+        )
 
         logger.info("most recent context:")
-        # prompt = "\n".join(
-        #     [f"{msg.cause_by.__name__ if msg.cause_by else 'Main Requirement'}: {msg.content}" for msg in context]
-        # )
-        prompt = ""
         latest_action = context[-1].cause_by.__name__ if context[-1].cause_by else ""
-        prompt += f"\nPlease review output from {latest_action}:\n" \
+        prompt = f"\nPlease review output from {latest_action}:\n" \
             "If you want to change a task in the plan, say 'change task task_id, ... (things to change)'\n" \
-            "If you confirm the output and wish to continue with the current process, type CONFIRM:\n"
+            "If you confirm the output and wish to continue with the current process, type CONFIRM\n" \
+            "If you want to terminate the process, type exit:\n"
         rsp = input(prompt)
-        confirmed = "confirm" in rsp.lower()
+
+        if rsp.lower() in ("exit"):
+            exit()
+
+        confirmed = rsp.lower() in ("confirm", "yes", "y")
 
         return rsp, confirmed
 
@@ -126,7 +129,7 @@ class MLEngineer(Role):
             # print(result)
             self.working_memory.add(Message(content=result, role="user", cause_by=ExecutePyCode))
 
-            if code.startswith("!pip"):
+            if "!pip" in code:
                 success = False
             # if not success:
             #     await self._ask_review()
@@ -139,8 +142,8 @@ class MLEngineer(Role):
         if not self.auto_run:
             context = self.get_useful_memories()
             review, confirmed = await AskReview().run(context=context[-5:], plan=self.plan)
-            if review.lower() not in ("confirm", "y", "yes"):
-                self._rc.memory.add(Message(content=review, role="user", cause_by=AskReview))
+            if not confirmed:
+                self.working_memory.add(Message(content=review, role="user", cause_by=AskReview))
             return confirmed
         return True
 
@@ -172,11 +175,14 @@ class MLEngineer(Role):
         return self._rc.memory
 
 if __name__ == "__main__":
-    # requirement = "create a normal distribution and visualize it"
-    requirement = "run some analysis on iris dataset"
+    requirement = "Run data analysis on sklearn Iris dataset, include a plot"
+    # requirement = "Run data analysis on sklearn Diabetes dataset, include a plot"
+    # requirement = "Run data analysis on sklearn Wine recognition dataset, include a plot, and train a model to predict wine class (20% as validation), and show validation accuracy"
+    # requirement = "Run data analysis on sklearn Wisconsin Breast Cancer dataset, include a plot, train a model to predict targets (20% as validation), and show validation accuracy"
+    # requirement = "Run EDA and visualization on this dataset, train a model to predict survival, report metrics on validation set (20%), dataset: workspace/titanic/train.csv"
 
-    async def main(requirement: str = requirement):
-        role = MLEngineer(goal=requirement)
+    async def main(requirement: str = requirement, auto_run: bool = False):
+        role = MLEngineer(goal=requirement, auto_run=auto_run)
         await role.run(requirement)
 
     fire.Fire(main)