From 21d97a23bb65b92a0379ff101ecbd497bd6e8537 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Wed, 6 Dec 2023 17:31:51 +0800
Subject: [PATCH] output code_steps to json

---
 metagpt/actions/write_analysis_code.py |  1 -
 metagpt/actions/write_code_steps.py    | 25 ++++++++++++++-----------
 metagpt/roles/ml_engineer.py           |  2 +-
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/metagpt/actions/write_analysis_code.py b/metagpt/actions/write_analysis_code.py
index cfec95deb..71467edd0 100644
--- a/metagpt/actions/write_analysis_code.py
+++ b/metagpt/actions/write_analysis_code.py
@@ -153,7 +153,6 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         context: List[Message],
         plan: Plan = None,
         code_steps: str = "",
-        data_desc: str = "",
     ) -> str:
         task_type = plan.current_task.task_type
         available_tools = registry.get_all_schema_by_module(task_type)
diff --git a/metagpt/actions/write_code_steps.py b/metagpt/actions/write_code_steps.py
index d3f6e5553..0bfb9c225 100644
--- a/metagpt/actions/write_code_steps.py
+++ b/metagpt/actions/write_code_steps.py
@@ -4,18 +4,12 @@ from typing import Dict, List, Union
 
 from metagpt.actions import Action
 from metagpt.schema import Message, Task, Plan
-
+from metagpt.utils.common import CodeParser
 
 CODE_STEPS_PROMPT_TEMPLATE = """
 # Context
 {context}
 
-##  Format example
-1.
-2.
-3.
-...
-
 -----
 Tasks are all code development tasks.
 You are a professional engineer, the main goal is to plan out concise solution steps for Current Task before coding.
@@ -25,7 +19,16 @@ The output plan should following the subsequent principles:
 1.The plan is a rough checklist of steps outlining the entire program's structure.Try to keep the number of steps fewer than 5.
 2.The steps should be written concisely and at a high level, avoiding overly detailed implementation specifics.
 3.The execution of the plan happens sequentially, but the plan can incorporate conditional (if) and looping(loop) keywords for more complex structures.
-4.Output carefully referenced "Format example" in format.
+
+Output the code steps in a JSON format, as shown in this example:
+```json
+{
+    "Step 1": "",
+    "Step 2": "",
+    "Step 3": "",
+    ...
+}
+```
 """
 
 STRUCTURAL_CONTEXT = """
@@ -51,10 +54,11 @@ class WriteCodeSteps(Action):
         """
 
         context = self.get_context(plan)
-        code_steps_prompt = CODE_STEPS_PROMPT_TEMPLATE.format(
-            context=context,
+        code_steps_prompt = CODE_STEPS_PROMPT_TEMPLATE.replace(
+            "{context}", context
         )
         code_steps = await self._aask(code_steps_prompt)
+        code_steps = CodeParser.parse_code(block=None, text=code_steps)
         return code_steps
 
     def get_context(self, plan: Plan):
@@ -74,4 +78,3 @@ class WriteCodeSteps(Action):
         )
         # print(context)
         return context
-
diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py
index 148851e9e..c2841be4c 100644
--- a/metagpt/roles/ml_engineer.py
+++ b/metagpt/roles/ml_engineer.py
@@ -294,7 +294,7 @@ if __name__ == "__main__":
     # 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"
+    # 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"
 
     requirement = "Perform data analysis on the provided data. Train a model to predict the target variable Survived. Include data preprocessing, feature engineering, and modeling in your pipeline. The metric is accuracy."
     data_path = "/data/lidanyang/tabular_data/titanic"