refine code

This commit is contained in:
lidanyang 2024-07-29 11:01:43 +08:00
parent 226f9d9729
commit 8d6a2801a0
2 changed files with 12 additions and 14 deletions

View file

@ -1,7 +1,5 @@
from __future__ import annotations
import re
from pydantic import Field, model_validator
from metagpt.actions.di.execute_nb_code import ExecuteNbCode
@ -56,13 +54,11 @@ class DataAnalyst(RoleZero):
await self.execute_code.init_code()
# plan info
try:
plan_status = self.planner.get_plan_status()
plan_status = re.sub(
r"### execution result.*?(?=## Current Task|## Task Guidance)", "", plan_status, flags=re.DOTALL
)
except AttributeError as e:
return "All tasks have been completed. Please check to end or append task first before writing code."
if self.planner.current_task:
# clear task result from plan to save token, since it has been in memory
plan_status = self.planner.get_plan_status(exclude=["task_result"])
else:
return "No current_task found now. Please use command Plan.append_task to add a task first."
# tool info
if self.custom_tool_recommender:

View file

@ -1,6 +1,7 @@
from __future__ import annotations
import json
from typing import List
from pydantic import BaseModel, Field
@ -165,8 +166,9 @@ class Planner(BaseModel):
return context_msg + self.working_memory.get()
def get_plan_status(self) -> str:
def get_plan_status(self, exclude: List[str] = None) -> str:
# prepare components of a plan status
exclude = exclude or []
finished_tasks = self.plan.get_finished_tasks()
code_written = [remove_comments(task.code) for task in finished_tasks]
code_written = "\n\n".join(code_written)
@ -178,11 +180,11 @@ class Planner(BaseModel):
# combine components in a prompt
prompt = PLAN_STATUS.format(
code_written=code_written,
task_results=task_results,
code_written=code_written if "code" not in exclude else "omit here",
task_results=task_results if "task_result" not in exclude else "omit here",
current_task=self.current_task.instruction,
current_task_code=self.current_task.code if self.current_task.code else "",
current_task_result=self.current_task.result if self.current_task.result else "",
current_task_code=self.current_task.code if "code" not in exclude else "omit here",
current_task_result=self.current_task.result if "task_result" not in exclude else "omit here",
guidance=guidance,
)