merge mle & interpreter

metagpt/actions/mi/execute_nb_code.py

@@ -201,7 +201,7 @@ def truncate(result: str, keep_len: int = 2000, is_success: bool = True):
     if is_success:
         desc = f"Executed code successfully. Truncated to show only first {keep_len} characters\n"
     else:
-        desc = f"Executed code failed, please reflect the cause of bug and then debug. Truncated to show only last {keep_len} characters\n"
+        desc = f"Executed code failed, please reflect on the cause of bug and then debug. Truncated to show only last {keep_len} characters\n"
 
     if result.strip().startswith("<coroutine object"):
         result = "Executed code failed, you need use key word 'await' to run a async code."

metagpt/actions/mi/write_analysis_code.py

@@ -11,14 +11,19 @@
 
 from metagpt.actions import Action
 from metagpt.logs import logger
+from metagpt.prompts.mi.ml_action import MODEL_TRAIN_EXAMPLE, USE_ML_TOOLS_EXAMPLE
 from metagpt.prompts.mi.write_analysis_code import (
+    CHECK_DATA_PROMPT,
+    DEBUG_REFLECTION_EXAMPLE,
+    REFLECTION_PROMPT,
+    STRUCTUAL_PROMPT,
     TOOL_RECOMMENDATION_PROMPT,
-    TOOL_USAGE_PROMPT,
 )
 from metagpt.schema import Message, Plan, SystemMessage
 from metagpt.tools import TOOL_REGISTRY
 from metagpt.tools.tool_registry import validate_tool_names
-from metagpt.utils.common import CodeParser
+from metagpt.tools.tool_type import ToolType
+from metagpt.utils.common import CodeParser, process_message, remove_comments
 
 
 class WriteCodeWithTools(Action):
@@ -27,7 +32,6 @@ class WriteCodeWithTools(Action):
     use_tools: bool = True
     # selected tools to choose from, listed by their names. An empty list means selection from all tools.
     selected_tools: list[str] = []
-    DEFAULT_SYSTEM_MSG: str = """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. Must reuse variables in the lastest other code directly, dont creat it again, it is very import for you. Use !pip install in a standalone block to install missing packages.Usually the libraries you need are already installed.Dont check if packages already imported.**"""  # prompt reference: https://github.com/KillianLucas/open-interpreter/blob/v0.1.4/interpreter/system_message.txt
 
     def _insert_system_message(self, context: list[Message], system_msg: str = None):
         system_msg = system_msg or self.DEFAULT_SYSTEM_MSG
@@ -98,34 +102,89 @@ async def _prepare_tools(self, plan: Plan) -> Tuple[dict, str]:
             TOOL_REGISTRY.get_tool_type(tool_type).usage_prompt if TOOL_REGISTRY.has_tool_type(tool_type) else ""
         )
 
+        # ML-specific tool usage examples
+        examples = ""
+        if plan.current_task.task_type in [
+            ToolType.DATA_PREPROCESS.type_name,
+            ToolType.FEATURE_ENGINEERING.type_name,
+        ]:
+            examples = USE_ML_TOOLS_EXAMPLE
+        elif plan.current_task.task_type in [ToolType.MODEL_TRAIN.type_name]:
+            examples = MODEL_TRAIN_EXAMPLE
+
         # prepare schemas of available tools
         tool_schemas = {}
         available_tools = self._get_tools_by_type(tool_type)
         if available_tools:
             available_tools = {tool_name: tool.schemas["description"] for tool_name, tool in available_tools.items()}
             tool_schemas = await self._recommend_tool(plan.current_task.instruction, available_tools)
 
-        return tool_schemas, tool_type_usage_prompt
+        return tool_schemas, tool_type_usage_prompt, examples
+
+    async def _debug_with_reflection(self, context: list[Message], working_memory: list[Message]):
+        reflection_prompt = REFLECTION_PROMPT.format(
+            debug_example=DEBUG_REFLECTION_EXAMPLE,
+            context=context,
+            previous_impl=working_memory,
+        )
+        # print(reflection_prompt)
+        system_prompt = "You are an AI Python assistant. You will be given your previous implementation code of a task, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
+
+        rsp = await self._aask(reflection_prompt, system_msgs=[system_prompt])
+        reflection = json.loads(CodeParser.parse_code(block=None, text=rsp))
+
+        return reflection["improved_impl"]
 
     async def run(
         self,
-        context: list[Message],
         plan: Plan,
+        working_memory: list[Message] = [],
+        use_reflection: bool = False,
         **kwargs,
     ) -> str:
-        if self.use_tools:
-            # prepare tool schemas and tool-type-specific instruction
-            tool_schemas, tool_type_usage_prompt = await self._prepare_tools(plan=plan)
+        # prepare tool schemas and tool-type-specific instruction
+        tool_schemas, tool_type_usage_prompt, examples = await self._prepare_tools(plan=plan)
+
+        # necessary components to be used in prompt
+        finished_tasks = plan.get_finished_tasks()
+        code_written = [remove_comments(task.code) for task in finished_tasks]
+        code_written = "\n\n".join(code_written)
+        task_results = [task.result for task in finished_tasks]
+        task_results = "\n\n".join(task_results)
+
+        # structure prompt
+        structual_prompt = STRUCTUAL_PROMPT.format(
+            user_requirement=plan.goal,
+            code_written=code_written,
+            task_results=task_results,
+            current_task=plan.current_task.instruction,
+            tool_type_usage_prompt=tool_type_usage_prompt,
+            tool_schemas=tool_schemas,
+            examples=examples,
+        )
+        context = [Message(content=structual_prompt, role="user")] + working_memory
+        context = process_message(context)
+
+        # temp = context + working_memory
+        # print(*temp, sep="***\n\n***")
 
-            # form a complete tool usage instruction and include it as a message in context
-            tools_instruction = TOOL_USAGE_PROMPT.format(
-                tool_schemas=tool_schemas, tool_type_usage_prompt=tool_type_usage_prompt
-            )
-            context.append(Message(content=tools_instruction, role="user"))
+        # LLM call
+        if not use_reflection:
+            rsp = await self.llm.aask(context, **kwargs)
+            code = CodeParser.parse_code(block=None, text=rsp)
 
-        # prepare prompt & LLM call
-        prompt = self._insert_system_message(context)
+        else:
+            code = await self._debug_with_reflection(context=context, working_memory=working_memory)
 
-        rsp = await self.llm.aask_code(prompt)
+        return code
 
-        return rsp
+
+class CheckData(Action):
+    async def run(self, plan: Plan = None) -> dict:
+        finished_tasks = plan.get_finished_tasks()
+        code_written = [remove_comments(task.code) for task in finished_tasks]
+        code_written = "\n\n".join(code_written)
+        prompt = CHECK_DATA_PROMPT.format(code_written=code_written)
+        rsp = await self._aask(prompt)
+        code = CodeParser.parse_code(block=None, text=rsp)
+        return code

metagpt/prompts/mi/ml_action.py

@@ -4,62 +4,7 @@
 # @Author  : lidanyang
 # @File    : ml_action
 # @Desc    :
-UPDATE_DATA_COLUMNS = """
-# Background
-Keep dataset column information updated before model train.
-## Tasks Done
-```python
-{history_code}
-```end
-
-# Task
-Print the the latest column information after 'Tasks Done' code. Use the following code:
-```python
-from metagpt.tools.libs.data_preprocess import get_column_info
-
-column_info = get_column_info(df)
-print("column_info")
-print(column_info)
-```end
-
-# Constraints:
-- Use the DataFrame variable from 'Tasks Done' in place of df.
-- Your code is to be added to a new cell in jupyter.
-"""
-
-ML_PROMPT = """
-# Background
-As a data scientist, you need to help user to achieve their goal [{user_requirement}] step-by-step in an continuous Jupyter notebook.
-
-## Done Tasks
-```python
-{history_code}
-```end
-
-## Current Task
-{current_task}
-
-# Latest Data Info
-Latest data info after previous tasks:
-{column_info}
-
-# Task
-Write complete code for 'Current Task'. And avoid duplicating code from 'Done Tasks', such as repeated import of packages, reading data, etc.
-Specifically, {tool_type_usage_prompt}
-
-# Capabilities
-- You can utilize pre-defined tools in any code lines from 'Available Tools' in the form of Python Class.
-- You can freely combine the use of any other public packages, like sklearn, numpy, pandas, etc..
-
-# Available Tools:
-Each Class tool is described in JSON format. When you call a tool, import the tool from its path first.
-{tool_schemas}
-
-{examples}
-"""
-
-USE_NO_TOOLS_EXAMPLE = """
-# Output Example:
+MODEL_TRAIN_EXAMPLE = """
 when current task is "train a lightgbm model on training data", the code can be like:
 ```python
 # Step 1: check data type and convert to numeric
@@ -80,8 +25,7 @@
 - Ensure the output new code is executable in the same Jupyter notebook with previous tasks code have been executed.
 """
 
-USE_TOOLS_EXAMPLE = """
-# Output Example:
+USE_ML_TOOLS_EXAMPLE = """
 when current task is "do data preprocess, like fill missing value, handle outliers, etc.", the code can be like:
 ```python
 # Step 1: fill missing value