only assign task_type when use_tools

This commit is contained in:
lidanyang 2023-12-01 14:29:51 +08:00
parent e4a17d122c
commit 2049f6cd01
2 changed files with 8 additions and 3 deletions

View file

@ -53,7 +53,9 @@ class WritePlan(Action):
task["task_type"] = task_type
return json.dumps(tasks)
async def run(self, context: List[Message], max_tasks: int = 5) -> str:
async def run(
self, context: List[Message], max_tasks: int = 5, use_tools: bool = False
) -> str:
prompt = (
self.PROMPT_TEMPLATE.replace("__context__", "\n".join([str(ct) for ct in context]))
# .replace("__current_plan__", current_plan)
@ -61,7 +63,8 @@ class WritePlan(Action):
)
rsp = await self._aask(prompt)
rsp = CodeParser.parse_code(block=None, text=rsp)
rsp = await self.assign_task_type(json.loads(rsp))
if use_tools:
rsp = await self.assign_task_type(json.loads(rsp))
return rsp
@staticmethod

View file

@ -173,7 +173,9 @@ class MLEngineer(Role):
plan_confirmed = False
while not plan_confirmed:
context = self.get_useful_memories()
rsp = await WritePlan().run(context, max_tasks=max_tasks)
rsp = await WritePlan().run(
context, max_tasks=max_tasks, use_tools=self.use_tools
)
self.working_memory.add(
Message(content=rsp, role="assistant", cause_by=WritePlan)
)