Merge branch 'dev_tool_selection' of https://gitlab.deepwisdomai.com/agents/data_agents_opt into dev_tool_selection

metagpt/actions/debug_code.py

@@ -3,7 +3,7 @@ from typing import Dict, List, Union, Tuple, Optional, Any
 from metagpt.actions import Action
 from metagpt.logs import logger
 from metagpt.schema import Message, Plan
-from metagpt.utils.common import CodeParser
+from metagpt.utils.common import CodeParser, create_func_config
 from metagpt.actions.write_analysis_code import BaseWriteAnalysisCode
 
 DEBUG_REFLECTION_EXAMPLE = '''Example 1:
@@ -39,25 +39,39 @@ DEBUG_REFLECTION_EXAMPLE = '''Example 1:
 REFLECTION_PROMPT = """
                        Here is an example for you.
                        {debug_example}
-                       [requirement]
-                       {goal}
-                       [finished code]
-                       finished code are executable, and you should based on the code to continue your current code debug
-                       {finished_code}
-                       
-                       try to reuse the code here to understand the coding task.
+                       [context]
+                       {context}
                        
                        [previous impl]
                        {code}
                        [runtime Error]
                        {runtime_result}
 
-                        Analysis the error step by step, provide me improve method. Do not repeat [previous impl]
+                        Analysis the error step by step, provide me improve method and code. Remember to follow [context] requirement.
                         [reflection on previous impl]:
                         xxx
 
                        """
 
+CODE_REFLECTION = {
+    "name": "execute_reflection_code",
+    "description": "Execute reflection code.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "reflection": {
+                "type": "string",
+                "description": "Reflection on previous impl.",
+            },
+            "improved_impl": {
+                "type": "string",
+                "description": "Refined code after reflection.",
+            },
+        },
+        "required": ["reflection", "improved_impl"],
+    },
+}
+
 
 def message_to_str(message: Message) -> str:
     return f"{message.role}: {message.content}"
@@ -75,52 +89,68 @@ class DebugCode(BaseWriteAnalysisCode):
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
     
-    async def run_reflection(self, goal, finished_code, finished_code_result, code, runtime_result) -> str:
+    async def run_reflection(
+        self,
+        # goal,
+        # finished_code,
+        # finished_code_result,
+        context: List[Message],
+        code,
+        runtime_result,
+    ) -> dict:
         info = []
-        finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
+        # finished_code_and_result = finished_code + "\n [finished results]\n\n" + finished_code_result
         reflection_prompt = REFLECTION_PROMPT.format(debug_example=DEBUG_REFLECTION_EXAMPLE,
-                                                     goal=goal,
-                                                     finished_code=finished_code_and_result,
+                                                     context=context,
+                                                     # goal=goal,
+                                                     # finished_code=finished_code_and_result,
                                                      code=code,
                                                      runtime_result=runtime_result
                                                      )
-        system_prompt = "You are an AI Python assistant. You will be given your previous implementation of a function, runtime error results, and a hint to change the implementation appropriately. Write your full implementation "
+        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 "
         info.append(Message(role="system", content=system_prompt))
-        info.append(Message(role="assistant", content=reflection_prompt))
+        info.append(Message(role="user", content=reflection_prompt))
         
-        msg = messages_to_str(info)
-        resp = await self.llm.aask(msg=msg)
+        # msg = messages_to_str(info)
+        # resp = await self.llm.aask(msg=msg)
+        resp = await self.llm.aask_code(messages=info, **create_func_config(CODE_REFLECTION))
         logger.info(f"reflection is {resp}")
         return resp
     
-    async def rewrite_code(self, reflection: str = "", code_context: str = "") -> str:
-        """
-        根据reflection重写代码
-        """
-        info = []
-        info.append(Message(role="assistant", content=f"[code context]:{code_context}"
-                                                      f"finished code are executable, and you should based on the code to continue your current code debug and improvement"
-                                                      f"[reflection]: \n {reflection}"))
-        info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
-        msg = messages_to_str(info)
-        resp = await self.llm.aask(msg=msg)
-        logger.info(f"improve code is {resp}")
-        improv_code = CodeParser.parse_code(block=None, text=resp)
-        return improv_code
+    # async def rewrite_code(self, reflection: str = "", context: List[Message] = None) -> str:
+    #     """
+    #     根据reflection重写代码
+    #     """
+    #     info = context
+    #     # info.append(Message(role="assistant", content=f"[code context]:{code_context}"
+    #     #                                               f"finished code are executable, and you should based on the code to continue your current code debug and improvement"
+    #     #                                               f"[reflection]: \n {reflection}"))
+    #     info.append(Message(role="assistant", content=f"[reflection]: \n {reflection}"))
+    #     info.append(Message(role="user", content=f"[improved impl]:\n Return in Python block"))
+    #     msg = messages_to_str(info)
+    #     resp = await self.llm.aask(msg=msg)
+    #     improv_code = CodeParser.parse_code(block=None, text=resp)
+    #     return improv_code
     
     async def run(self,
+                  context: List[Message] = None,
                   plan: str = "",
-                  finished_code: str = "",
-                  finished_code_result: str = "",
+                  # finished_code: str = "",
+                  # finished_code_result: str = "",
                   code: str = "",
                   runtime_result: str = "") -> str:
         """
         根据当前运行代码和报错信息进行reflection和纠错
         """
-        reflection = await self.run_reflection(plan, finished_code=finished_code,
-                                               finished_code_result=finished_code_result,
-                                               code=code,
-                                               runtime_result=runtime_result)
+        reflection = await self.run_reflection(
+            # plan,
+            # finished_code=finished_code,
+            # finished_code_result=finished_code_result,
+            code=code,
+            context=context,
+            runtime_result=runtime_result,
+        )
         # 根据reflection结果重写代码
-        improv_code = await self.rewrite_code(reflection, code_context=finished_code)
+        # improv_code = await self.rewrite_code(reflection, context=context)
+        improv_code = reflection['improved_impl']
         return improv_code

metagpt/actions/write_analysis_code.py

@@ -4,7 +4,9 @@
 @Author  :   orange-crow
 @File    :   write_code_v2.py
 """
-from typing import Dict, List, Union, Tuple, Optional, Any
+from typing import Dict, List, Union, Tuple
+
+import yaml
 
 from metagpt.actions import Action
 from metagpt.logs import logger
@@ -15,11 +17,9 @@ from metagpt.prompts.ml_engineer import (
     TOOL_USAGE_PROMPT,
     ML_SPECIFIC_PROMPT,
     ML_MODULE_MAP,
-    TOOL_OUTPUT_DESC, DATA_PROCESS_PROMPT,
-    GENERATE_CODE_PROMPT
+    GENERATE_CODE_PROMPT,
 )
 from metagpt.schema import Message, Plan
-from metagpt.tools.functions import registry
 from metagpt.utils.common import create_func_config, remove_comments
 
 
@@ -100,40 +100,55 @@ class WriteCodeByGenerate(BaseWriteAnalysisCode):
 class WriteCodeWithTools(BaseWriteAnalysisCode):
     """Write code with help of local available tools. Choose tools first, then generate code to use the tools"""
 
-    @staticmethod
-    def _parse_recommend_tools(module: str, recommend_tools: list) -> List[Dict]:
+    def __init__(self, name: str = "", context=None, llm=None, schema_path=None):
+        super().__init__(name, context, llm)
+        self.schema_path = schema_path
+        self.available_tools = {}
+
+        if self.schema_path is not None:
+            self._load_tools(schema_path)
+
+    def _load_tools(self, schema_path):
+        """Load tools from yaml file"""
+        yml_files = schema_path.glob("*.yml")
+        for yml_file in yml_files:
+            module = yml_file.stem
+            with open(yml_file, "r", encoding="utf-8") as f:
+                self.available_tools[module] = yaml.safe_load(f)
+
+    def _parse_recommend_tools(self, module: str, recommend_tools: list) -> dict:
         """
         Parses and validates a list of recommended tools, and retrieves their schema from registry.
 
         Args:
             module (str): The module name for querying tools in the registry.
-            recommend_tools (list): A list of lists of recommended tools for each step.
+            recommend_tools (list): A list of recommended tools.
 
         Returns:
-            List[Dict]: A list of dicts of valid tool schemas.
+            dict: A dict of valid tool schemas.
         """
         valid_tools = []
-        available_tools = registry.get_all_by_module(module).keys()
+        available_tools = self.available_tools[module].keys()
         for tool in recommend_tools:
             if tool in available_tools:
                 valid_tools.append(tool)
 
-        tool_catalog = registry.get_schemas(module, valid_tools)
+        tool_catalog = {tool: self.available_tools[module][tool] for tool in valid_tools}
         return tool_catalog
 
     async def _tool_recommendation(
-            self,
-            task: str,
-            code_steps: str,
-            available_tools: list
+        self,
+        task: str,
+        code_steps: str,
+        available_tools: dict,
     ) -> list:
         """
         Recommend tools for the specified task.
 
         Args:
-            context (List[Message]): Action output history, source action denoted by Message.cause_by
+            task (str): the task to recommend tools for
             code_steps (str): the code steps to generate the full code for the task
-            available_tools (list): the available tools for the task
+            available_tools (dict): the available tools description
 
         Returns:
             list: recommended tools for the specified task
@@ -149,27 +164,23 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
         return recommend_tools
 
     async def run(
-            self,
-            context: List[Message],
-            plan: Plan = None,
-            code_steps: str = "",
-            column_info: str = "",
-            **kwargs,
-    ) -> str:
+        self,
+        context: List[Message],
+        plan: Plan = None,
+        code_steps: str = "",
+        column_info: str = "",
+        **kwargs,
+    ) -> Tuple[List[Message], str]:
         task_type = plan.current_task.task_type
-        available_tools = registry.get_all_schema_by_module(task_type)
+        available_tools = self.available_tools.get(task_type, {})
         special_prompt = ML_SPECIFIC_PROMPT.get(task_type, "")
 
-        column_names = kwargs.get("column_names", {})
         finished_tasks = plan.get_finished_tasks()
         code_context = [remove_comments(task.code) for task in finished_tasks]
         code_context = "\n\n".join(code_context)
 
         if len(available_tools) > 0:
-            available_tools = [
-                {k: tool[k] for k in ["name", "description"] if k in tool}
-                for tool in available_tools
-            ]
+            available_tools = {k: v["description"] for k, v in available_tools.items()}
 
             recommend_tools = await self._tool_recommendation(
                 plan.current_task.instruction,
@@ -180,46 +191,27 @@ class WriteCodeWithTools(BaseWriteAnalysisCode):
             logger.info(f"Recommended tools: \n{recommend_tools}")
 
             module_name = ML_MODULE_MAP[task_type]
-            output_desc = TOOL_OUTPUT_DESC.get(task_type, "")
-            new_code = ""
-
-            for idx, tool in enumerate(recommend_tools):
-                hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n "
-
-                prompt = TOOL_USAGE_PROMPT.format(
-                    goal=plan.current_task.instruction,
-                    context=hist_info,
-                    code_steps=code_steps,
-                    column_names=column_names,
-                    special_prompt=special_prompt,
-                    module_name=module_name,
-                    output_desc=output_desc,
-                    function_catalog=tool_catalog[idx],
-                )
-
-                tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
-
-                rsp = await self.llm.aask_code(prompt, **tool_config)
-                logger.info(f"rsp is: {rsp}")
-                # final_code = final_code + "\n\n" + rsp["code"]
-                # final_code[key] = rsp["code"]
-                new_code = new_code + "\n\n" + rsp["code"]
-                code_context = code_context + "\n\n" + rsp["code"]
-            return new_code
-
-        else:
-            hist_info = f"Previous finished code is \n\n ```Python {code_context} ``` \n\n "
-
-            prompt = GENERATE_CODE_PROMPT.format(
-                goal=plan.current_task.instruction,
-                context=hist_info,
-                code_steps=code_steps,
+            prompt = TOOL_USAGE_PROMPT.format(
+                user_requirement=plan.goal,
+                history_code=code_context,
+                current_task=plan.current_task.instruction,
+                column_info=column_info,
                 special_prompt=special_prompt,
-                # column_names=column_names
+                code_steps=code_steps,
+                module_name=module_name,
+                tool_catalog=tool_catalog,
+            )
+        else:
+            prompt = GENERATE_CODE_PROMPT.format(
+                user_requirement=plan.goal,
+                history_code=code_context,
+                current_task=plan.current_task.instruction,
+                column_info=column_info,
+                special_prompt=special_prompt,
+                code_steps=code_steps,
             )
 
-            tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
-            logger.info(f"prompt is: {prompt}")
-            rsp = await self.llm.aask_code(prompt, **tool_config)
-            logger.info(f"rsp is: {rsp}")
-            return rsp["code"]
+        tool_config = create_func_config(CODE_GENERATOR_WITH_TOOLS)
+        rsp = await self.llm.aask_code(prompt, **tool_config)
+        context = [Message(content=prompt, role="user")]
+        return context, rsp["code"]

metagpt/actions/write_code_steps.py

@@ -63,8 +63,8 @@ class WriteCodeSteps(Action):
 
     def get_context(self, plan: Plan):
         user_requirement = plan.goal
-        # select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
-        select_task_keys = ['task_id','code']
+        select_task_keys = ['task_id', 'instruction', 'is_finished', 'code']
+        # select_task_keys = ['task_id','code']
         
         def process_task(task):
             task_dict = task.dict()
@@ -72,7 +72,7 @@ class WriteCodeSteps(Action):
             return ptask
         
         tasks = json.dumps(
-            [process_task(task) for task in plan.tasks if task.is_finished==True], indent=4, ensure_ascii=False
+            [process_task(task) for task in plan.tasks], indent=4, ensure_ascii=False
         )
         
         current_task = json.dumps(process_task(plan.current_task)) if plan.current_task else {}