From 2049f6cd01c66e5ed0402a18bebc20b1a9ceda5d Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 1 Dec 2023 14:29:51 +0800
Subject: [PATCH] only assign task_type when use_tools

---
 metagpt/actions/write_plan.py | 7 +++++--
 metagpt/roles/ml_engineer.py  | 4 +++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/metagpt/actions/write_plan.py b/metagpt/actions/write_plan.py
index f7c096f2c..5145ffd68 100644
--- a/metagpt/actions/write_plan.py
+++ b/metagpt/actions/write_plan.py
@@ -53,7 +53,9 @@ class WritePlan(Action):
             task["task_type"] = task_type
         return json.dumps(tasks)
 
-    async def run(self, context: List[Message], max_tasks: int = 5) -> str:
+    async def run(
+        self, context: List[Message], max_tasks: int = 5, use_tools: bool = False
+    ) -> str:
         prompt = (
             self.PROMPT_TEMPLATE.replace("__context__", "\n".join([str(ct) for ct in context]))
             # .replace("__current_plan__", current_plan)
@@ -61,7 +63,8 @@ class WritePlan(Action):
         )
         rsp = await self._aask(prompt)
         rsp = CodeParser.parse_code(block=None, text=rsp)
-        rsp = await self.assign_task_type(json.loads(rsp))
+        if use_tools:
+            rsp = await self.assign_task_type(json.loads(rsp))
         return rsp
 
     @staticmethod
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 0ea73a045..8e02e093b 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -173,7 +173,9 @@ class MLEngineer(Role):
         plan_confirmed = False
         while not plan_confirmed:
             context = self.get_useful_memories()
-            rsp = await WritePlan().run(context, max_tasks=max_tasks)
+            rsp = await WritePlan().run(
+                context, max_tasks=max_tasks, use_tools=self.use_tools
+            )
             self.working_memory.add(
                 Message(content=rsp, role="assistant", cause_by=WritePlan)
             )