From 8afac012b49df5ffb26dc031345c685c748e8797 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=A3=92=E6=A3=92?= <liubangbang@fuzhi.ai>
Date: Wed, 20 Dec 2023 09:57:19 +0800
Subject: [PATCH] set the plan.current_task.task_type to udf when use udfs.

---
 metagpt/roles/ml_engineer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index fa9acadbc..3c1853fd5 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -175,6 +175,8 @@ class MLEngineer(Role):
                     logger.warning("Writing code with user-defined function tools...")
                     logger.info(f"Local user defined function as following:\
                         \n{json.dumps(list(UDFS_YAML.keys()), indent=2, ensure_ascii=False)}")
+                    # set task_type to `udf`
+                    self.plan.current_task.task_type = 'udf'
                     tool_context, code = await WriteCodeWithTools(schema_path=UDFS_YAML).run(
                         context=context,
                         plan=self.plan,
@@ -184,6 +186,7 @@ class MLEngineer(Role):
                     cause_by = WriteCodeWithTools
                 else:
                     logger.info("Write code with pure generation")
+                    # TODO: 添加基于current_task.instruction-code_path的k-v缓存
                     code = await WriteCodeByGenerate().run(
                         context=context, plan=self.plan, temperature=0.0
                     )