Skip to content
This repository was archived by the owner on Apr 12, 2026. It is now read-only.

Commit 20e5ba8

Browse files
SoulterNekyuuYa
authored andcommitted
refactor: improve astrbot builtin tool management (AstrBotDevs#7418)
* feat: Refactor astrbot builtin tool management - Introduced a new registry for builtin tools to streamline their management. - Added `SendMessageToUserTool`, `KnowledgeBaseQueryTool`, and various web search tools as builtin tools. - Updated `FunctionToolManager` to cache and retrieve builtin tools efficiently. - Modified `CronJobManager` to utilize the new `SendMessageToUserTool`. - Enhanced the dashboard to display readonly status for builtin tools and prevent toggling their state. - Added tests for builtin tool injection and retrieval to ensure proper functionality. * fix: escape file path in shell command to prevent injection vulnerabilities
1 parent d1fce99 commit 20e5ba8

20 files changed

Lines changed: 648 additions & 446 deletions

File tree

astrbot/core/astr_agent_tool_exec.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
LOCAL_EXECUTE_SHELL_TOOL,
2626
LOCAL_PYTHON_TOOL,
2727
PYTHON_TOOL,
28-
SEND_MESSAGE_TO_USER_TOOL,
2928
)
3029
from astrbot.core.cron.events import CronMessageEvent
3130
from astrbot.core.message.components import Image
@@ -37,6 +36,7 @@
3736
from astrbot.core.platform.message_session import MessageSession
3837
from astrbot.core.provider.entites import ProviderRequest
3938
from astrbot.core.provider.register import llm_tools
39+
from astrbot.core.tools.message_tools import SendMessageToUserTool
4040
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
4141
from astrbot.core.utils.history_saver import persist_agent_history
4242
from astrbot.core.utils.image_ref_utils import is_supported_image_ref
@@ -515,7 +515,9 @@ async def _wake_main_agent_for_background_result(
515515
)
516516
if not req.func_tool:
517517
req.func_tool = ToolSet()
518-
req.func_tool.add_tool(SEND_MESSAGE_TO_USER_TOOL)
518+
req.func_tool.add_tool(
519+
ctx.get_llm_tool_manager().get_builtin_tool(SendMessageToUserTool)
520+
)
519521

520522
result = await build_main_agent(
521523
event=cron_event, plugin_context=ctx, config=config, req=req

astrbot/core/astr_main_agent.py

Lines changed: 35 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232
FILE_UPLOAD_TOOL,
3333
GET_EXECUTION_HISTORY_TOOL,
3434
GET_SKILL_PAYLOAD_TOOL,
35-
KNOWLEDGE_BASE_QUERY_TOOL,
3635
LIST_SKILL_CANDIDATES_TOOL,
3736
LIST_SKILL_RELEASES_TOOL,
3837
LIVE_MODE_SYSTEM_PROMPT,
@@ -44,11 +43,9 @@
4443
ROLLBACK_SKILL_RELEASE_TOOL,
4544
RUN_BROWSER_SKILL_TOOL,
4645
SANDBOX_MODE_PROMPT,
47-
SEND_MESSAGE_TO_USER_TOOL,
4846
SYNC_SKILL_RELEASE_TOOL,
4947
TOOL_CALL_PROMPT,
5048
TOOL_CALL_PROMPT_SKILLS_LIKE_MODE,
51-
retrieve_knowledge_base,
5249
)
5350
from astrbot.core.conversation_mgr import Conversation
5451
from astrbot.core.message.components import File, Image, Record, Reply
@@ -63,16 +60,21 @@
6360
from astrbot.core.star.context import Context
6461
from astrbot.core.star.star_handler import star_map
6562
from astrbot.core.tools.cron_tools import (
66-
CREATE_CRON_JOB_TOOL,
67-
DELETE_CRON_JOB_TOOL,
68-
LIST_CRON_JOBS_TOOL,
63+
CreateActiveCronTool,
64+
DeleteCronJobTool,
65+
ListCronJobsTool,
66+
)
67+
from astrbot.core.tools.knowledge_base_tools import (
68+
KnowledgeBaseQueryTool,
69+
retrieve_knowledge_base,
6970
)
71+
from astrbot.core.tools.message_tools import SendMessageToUserTool
7072
from astrbot.core.tools.web_search_tools import (
71-
TAVILY_EXTRACT_WEB_PAGE_TOOL,
72-
WEB_SEARCH_BAIDU_TOOL,
73-
WEB_SEARCH_BOCHA_TOOL,
74-
WEB_SEARCH_BRAVE_TOOL,
75-
WEB_SEARCH_TAVILY_TOOL,
73+
BaiduWebSearchTool,
74+
BochaWebSearchTool,
75+
BraveWebSearchTool,
76+
TavilyExtractWebPageTool,
77+
TavilyWebSearchTool,
7678
normalize_legacy_web_search_config,
7779
)
7880
from astrbot.core.utils.file_extract import extract_file_moonshotai
@@ -226,7 +228,11 @@ async def _apply_kb(
226228
else:
227229
if req.func_tool is None:
228230
req.func_tool = ToolSet()
229-
req.func_tool.add_tool(KNOWLEDGE_BASE_QUERY_TOOL)
231+
req.func_tool.add_tool(
232+
plugin_context.get_llm_tool_manager().get_builtin_tool(
233+
KnowledgeBaseQueryTool
234+
)
235+
)
230236

231237

232238
async def _apply_file_extract(
@@ -1054,12 +1060,13 @@ def _apply_sandbox_tools(
10541060
req.system_prompt = f"{req.system_prompt or ''}\n{SANDBOX_MODE_PROMPT}\n"
10551061

10561062

1057-
def _proactive_cron_job_tools(req: ProviderRequest) -> None:
1063+
def _proactive_cron_job_tools(req: ProviderRequest, plugin_context: Context) -> None:
10581064
if req.func_tool is None:
10591065
req.func_tool = ToolSet()
1060-
req.func_tool.add_tool(CREATE_CRON_JOB_TOOL)
1061-
req.func_tool.add_tool(DELETE_CRON_JOB_TOOL)
1062-
req.func_tool.add_tool(LIST_CRON_JOBS_TOOL)
1066+
tool_mgr = plugin_context.get_llm_tool_manager()
1067+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(CreateActiveCronTool))
1068+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(DeleteCronJobTool))
1069+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(ListCronJobsTool))
10631070

10641071

10651072
async def _apply_web_search_tools(
@@ -1077,16 +1084,17 @@ async def _apply_web_search_tools(
10771084
if req.func_tool is None:
10781085
req.func_tool = ToolSet()
10791086

1087+
tool_mgr = plugin_context.get_llm_tool_manager()
10801088
provider = prov_settings.get("websearch_provider", "tavily")
10811089
if provider == "tavily":
1082-
req.func_tool.add_tool(WEB_SEARCH_TAVILY_TOOL)
1083-
req.func_tool.add_tool(TAVILY_EXTRACT_WEB_PAGE_TOOL)
1090+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(TavilyWebSearchTool))
1091+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(TavilyExtractWebPageTool))
10841092
elif provider == "bocha":
1085-
req.func_tool.add_tool(WEB_SEARCH_BOCHA_TOOL)
1093+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(BochaWebSearchTool))
10861094
elif provider == "brave":
1087-
req.func_tool.add_tool(WEB_SEARCH_BRAVE_TOOL)
1095+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(BraveWebSearchTool))
10881096
elif provider == "baidu_ai_search":
1089-
req.func_tool.add_tool(WEB_SEARCH_BAIDU_TOOL)
1097+
req.func_tool.add_tool(tool_mgr.get_builtin_tool(BaiduWebSearchTool))
10901098

10911099

10921100
def _get_compress_provider(
@@ -1348,12 +1356,16 @@ async def build_main_agent(
13481356
)
13491357

13501358
if config.add_cron_tools:
1351-
_proactive_cron_job_tools(req)
1359+
_proactive_cron_job_tools(req, plugin_context)
13521360

13531361
if event.platform_meta.support_proactive_message:
13541362
if req.func_tool is None:
13551363
req.func_tool = ToolSet()
1356-
req.func_tool.add_tool(SEND_MESSAGE_TO_USER_TOOL)
1364+
req.func_tool.add_tool(
1365+
plugin_context.get_llm_tool_manager().get_builtin_tool(
1366+
SendMessageToUserTool
1367+
)
1368+
)
13571369

13581370
if provider.provider_config.get("max_context_tokens", 0) <= 0:
13591371
model = provider.get_model()

0 commit comments

Comments
 (0)