From 28059f14df0cb4855d64c70a18f626505e93048c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E4=BC=9F=E9=9F=AC?= <huangweitao@fuzhi.ai>
Date: Wed, 31 Jul 2024 14:52:34 +0800
Subject: [PATCH] Add a "finish current task" command before the "end" command
 for the engineer and the data analyst.

---
 metagpt/prompts/di/data_analyst.py |  1 +
 metagpt/prompts/di/team_leader.py  |  1 +
 metagpt/roles/di/data_analyst.py   | 26 ++++++++++++++++----------
 metagpt/roles/di/engineer2.py      | 10 ++++++++++
 4 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/metagpt/prompts/di/data_analyst.py b/metagpt/prompts/di/data_analyst.py
index 56675de9b..3e4dac6f5 100644
--- a/metagpt/prompts/di/data_analyst.py
+++ b/metagpt/prompts/di/data_analyst.py
@@ -10,6 +10,7 @@ EXTRA_INSTRUCTION = """
 7. When you are making plan. It is highly recommend to plan and append all the tasks in first response once time.
 8. Don't finish_current_task multiple times for the same task. 
 9. Finish current task timely, such as when the code is written and executed successfully.
+10. When using the command 'end', add the command 'finish_current_task' before it.
 """
 
 TASK_TYPE_DESC = "\n".join([f"- **{tt.type_name}**: {tt.value.desc}" for tt in TaskType])
diff --git a/metagpt/prompts/di/team_leader.py b/metagpt/prompts/di/team_leader.py
index 86288ec73..4fc03222d 100644
--- a/metagpt/prompts/di/team_leader.py
+++ b/metagpt/prompts/di/team_leader.py
@@ -39,6 +39,7 @@ Fifth, describe if you should terminate, you should use **end** command to termi
  - All tasks are finished and current task is empty
  - You are repetitively replying to human
 Sixth, when planning, describe the requirements as they pertain to software development, data analysis, or other areas. If the requirements is a software development and no specific restrictions are mentioned, you must create a Product Requirements Document (PRD), write a System Design document, develop a project schedule, and then begin coding. List the steps you will undertake. Plan these steps in a single response.
+Seventh, describe the technologies you must use.
 Finally, combine your thoughts, describe what you want to do conscisely in 20 words, including which process you will taked and whether you will end, then follow your thoughts to list the commands, adhering closely to the instructions provided.
 """
 QUICK_THINK_SYSTEM_PROMPT = """
diff --git a/metagpt/roles/di/data_analyst.py b/metagpt/roles/di/data_analyst.py
index 154d1593d..b39053340 100644
--- a/metagpt/roles/di/data_analyst.py
+++ b/metagpt/roles/di/data_analyst.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 from pydantic import Field, model_validator
 
 from metagpt.actions.di.execute_nb_code import ExecuteNbCode
-from metagpt.actions.di.write_analysis_code import WriteAnalysisCode, CheckData
+from metagpt.actions.di.write_analysis_code import CheckData, WriteAnalysisCode
 from metagpt.logs import logger
 from metagpt.prompts.di.data_analyst import (
     CODE_STATUS,
@@ -111,15 +111,11 @@ class DataAnalyst(RoleZero):
         return output
 
     async def _check_data(self):
-        if (
-            not self.planner.plan.get_finished_tasks()
-            or self.planner.plan.current_task.task_type
-            not in [
-                TaskType.DATA_PREPROCESS.type_name,
-                TaskType.FEATURE_ENGINEERING.type_name,
-                TaskType.MODEL_TRAIN.type_name,
-            ]
-        ):
+        if not self.planner.plan.get_finished_tasks() or self.planner.plan.current_task.task_type not in [
+            TaskType.DATA_PREPROCESS.type_name,
+            TaskType.FEATURE_ENGINEERING.type_name,
+            TaskType.MODEL_TRAIN.type_name,
+        ]:
             return
         logger.info("Check updated data")
         code = await CheckData().run(self.planner.plan)
@@ -130,3 +126,13 @@ class DataAnalyst(RoleZero):
             print(result)
             data_info = DATA_INFO.format(info=result)
             self.rc.working_memory.add(Message(content=data_info, role="user", cause_by=CheckData))
+
+    async def _run_special_command(self, cmd) -> str:
+        """command requiring special check or parsing."""
+        # finish current task before end.
+        command_output = ""
+        if cmd["command_name"] == "end" and not self.planner.plan.is_plan_finished():
+            self.planner.plan.finish_current_task()
+            command_output += "Current task is finished. \n"
+        command_output += await super()._run_special_command(cmd)
+        return command_output
diff --git a/metagpt/roles/di/engineer2.py b/metagpt/roles/di/engineer2.py
index 1610a652b..3050071fd 100644
--- a/metagpt/roles/di/engineer2.py
+++ b/metagpt/roles/di/engineer2.py
@@ -26,3 +26,13 @@ class Engineer2(RoleZero):
 
     def _retrieve_experience(self) -> str:
         return ENGINEER_EXAMPLE
+
+    async def _run_special_command(self, cmd) -> str:
+        """command requiring special check or parsing."""
+        # finish current task before end.
+        command_output = ""
+        if cmd["command_name"] == "end" and not self.planner.plan.is_plan_finished():
+            self.planner.plan.finish_current_task()
+            command_output += "Current task is finished. \n"
+        command_output += await super()._run_special_command(cmd)
+        return command_output