diff --git a/metagpt/actions/write_task_guide.py b/metagpt/actions/write_task_guide.py new file mode 100644 index 000000000..eff53feef --- /dev/null +++ b/metagpt/actions/write_task_guide.py @@ -0,0 +1,82 @@ + +import json +from typing import Dict, List, Union + +from metagpt.actions import Action +from metagpt.schema import Message, Task, Plan + + +TASK_GUIDE_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. +A planning process can reduce the difficulty and improve the quality of coding. +You may be given some code plans for the tasks ahead, but you don't have to follow the existing plan when planning the current task. +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. +""" + +STRUCTURAL_CONTEXT = """ +## User Requirement +{user_requirement} +## Current Plan +{tasks} +## Current Task +{current_task} +""" + + +class WriteTaskGuide(Action): + + async def run(self, plan: Plan) -> str: + """Run of a task guide writing action, used in ml engineer + + Args: + plan (plan): task plan + useful_memories (list): useful_memories + Returns: + str: The dataset_descriptions string. + """ + + context = self.get_context(plan) + task_guide_prompt = TASK_GUIDE_PROMPT_TEMPLATE.format( + context=context, + ) + task_guide = await self._aask(task_guide_prompt) + return task_guide + + def get_context(self, plan: Plan): + user_requirement = plan.goal + task_rename_map = { + 'task_id': 'task_id', + 'instruction': 'instruction', + 'is_finished': 'is_finished', + # 'task_guide': 'code_plan' + } + + def process_task(task): + task_dict = task.dict() + ptask = {task_rename_map[k]: task_dict[k] for k in task_dict if k in task_rename_map} + return ptask + tasks = json.dumps( + [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 {} + context = STRUCTURAL_CONTEXT.format( + user_requirement=user_requirement, tasks=tasks, current_task=current_task + ) + # print(context) + return context + diff --git a/metagpt/roles/ml_engineer.py b/metagpt/roles/ml_engineer.py index 65583638e..d905b7bfd 100644 --- a/metagpt/roles/ml_engineer.py +++ b/metagpt/roles/ml_engineer.py @@ -12,6 +12,7 @@ from metagpt.logs import logger from metagpt.actions.write_plan import WritePlan from metagpt.actions.write_analysis_code import WriteCodeByGenerate, WriteCodeWithTools from metagpt.actions.execute_code import ExecutePyCode +from metagpt.actions.write_task_guide import WriteTaskGuide STRUCTURAL_CONTEXT = """ ## User Requirement @@ -66,11 +67,6 @@ class AskReview(Action): return rsp, confirmed -class WriteTaskGuide(Action): - async def run(self, task_instruction: str, data_desc: str = "") -> str: - return "" - - class MLEngineer(Role): def __init__( self, name="ABC", profile="MLEngineer", goal="", auto_run: bool = False @@ -79,7 +75,7 @@ class MLEngineer(Role): self._set_react_mode(react_mode="plan_and_act") self.plan = Plan(goal=goal) self.use_tools = False - self.use_task_guide = False + self.use_task_guide = True self.execute_code = ExecutePyCode() self.auto_run = auto_run @@ -92,7 +88,7 @@ class MLEngineer(Role): logger.info(f"ready to take on task {task}") # take on current task - code, result, success = await self._write_and_exec_code() + code, result, success, task_guide = await self._write_and_exec_code() # ask for acceptance, users can other refuse and change tasks in the plan task_result_confirmed = await self._ask_review() @@ -101,6 +97,7 @@ class MLEngineer(Role): # tick off this task and record progress task.code = code task.result = result + task.task_guide = task_guide self.plan.finish_current_task() self.working_memory.clear() @@ -110,7 +107,7 @@ class MLEngineer(Role): async def _write_and_exec_code(self, max_retry: int = 3): task_guide = ( - await WriteTaskGuide().run(self.plan.current_task.instruction) + await WriteTaskGuide().run(self.plan) if self.use_task_guide else "" ) @@ -156,7 +153,7 @@ class MLEngineer(Role): counter += 1 - return code, result, success + return code, result, success, task_guide async def _ask_review(self): if not self.auto_run: @@ -185,7 +182,7 @@ class MLEngineer(Role): def get_useful_memories(self) -> List[Message]: """find useful memories only to reduce context length and improve performance""" - + # TODO dataset description , code steps user_requirement = self.plan.goal tasks = json.dumps( [task.dict() for task in self.plan.tasks], indent=4, ensure_ascii=False @@ -204,9 +201,9 @@ class MLEngineer(Role): if __name__ == "__main__": - requirement = "Run data analysis on sklearn Iris dataset, include a plot" + # requirement = "Run data analysis on sklearn Iris dataset, include a plot" # 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 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" diff --git a/metagpt/schema.py b/metagpt/schema.py index e39f54a0c..db6861280 100644 --- a/metagpt/schema.py +++ b/metagpt/schema.py @@ -81,6 +81,7 @@ class Task(BaseModel): code: str = "" result: str = "" is_finished: bool = False + task_guide: str = "" class Plan(BaseModel):