diff --git a/metagpt/roles/di/data_analyst.py b/metagpt/roles/di/data_analyst.py index f13fc78fb..b3144a100 100644 --- a/metagpt/roles/di/data_analyst.py +++ b/metagpt/roles/di/data_analyst.py @@ -1,12 +1,13 @@ from __future__ import annotations -from pydantic import Field +from pydantic import Field, model_validator from metagpt.actions.di.execute_nb_code import ExecuteNbCode from metagpt.actions.di.write_analysis_code import WriteAnalysisCode from metagpt.logs import logger from metagpt.roles.di.role_zero import RoleZero from metagpt.schema import TaskResult, Message +from metagpt.tools.tool_recommend import BM25ToolRecommender, ToolRecommender from metagpt.tools.tool_registry import register_tool @@ -18,21 +19,22 @@ class DataAnalyst(RoleZero): tools: list[str] = ["Plan", "DataAnalyst", "RoleZero"] custom_tools: list[str] = ["machine learning", "web scraping", "Terminal"] + custom_tool_recommender: ToolRecommender = None use_reflection: bool = True write_code: WriteAnalysisCode = Field(default_factory=WriteAnalysisCode, exclude=True) execute_code: ExecuteNbCode = Field(default_factory=ExecuteNbCode, exclude=True) task_result: TaskResult = None + @model_validator(mode="after") + def set_custom_tool(self): + if self.custom_tools and not self.custom_tool_recommender: + self.custom_tool_recommender = BM25ToolRecommender(tools=self.custom_tools) + def _update_tool_execution(self): - self.tool_execution_map = { - "Plan.append_task": self.planner.plan.append_task, - "Plan.reset_task": self.planner.plan.reset_task, - "Plan.replace_task": self.planner.plan.replace_task, + self.tool_execution_map.update({ "DataAnalyst.write_and_exec_code": self.write_and_exec_code, - "RoleZero.ask_human": self.ask_human, - "RoleZero.reply_to_human": self.reply_to_human, - } + }) async def write_and_exec_code(self): """Write a code block for current task and execute it in an interactive notebook environment.""" diff --git a/metagpt/roles/di/role_zero.py b/metagpt/roles/di/role_zero.py index 4225d1f76..93abe8c02 100644 --- a/metagpt/roles/di/role_zero.py +++ b/metagpt/roles/di/role_zero.py @@ -41,10 +41,8 @@ class RoleZero(Role): max_react_loop: int = 20 # used for react mode # Tools - tools: list[str] = [] + tools: list[str] = [] # Use special symbol [""] to indicate use of all registered tools tool_recommender: ToolRecommender = None - custom_tools: list[str] = [] - custom_tool_recommender: ToolRecommender = None tool_execution_map: dict[str, Callable] = {} special_tool_commands: list[str] = ["Plan.finish_current_task", "end"] # Equipped with three basic tools by default for optional use @@ -70,8 +68,6 @@ class RoleZero(Role): self._set_react_mode(react_mode=self.react_mode, max_react_loop=self.max_react_loop) if self.tools and not self.tool_recommender: self.tool_recommender = BM25ToolRecommender(tools=self.tools, force=True) - if self.custom_tools and not self.custom_tool_recommender: - self.custom_tool_recommender = BM25ToolRecommender(tools=self.custom_tools) self.set_actions([RunCommand]) # HACK: Init Planner, control it through dynamic thinking; Consider formalizing as a react mode