mirror of
https://github.com/countbot-ai/CountBot.git
synced 2026-09-14 20:46:47 +08:00
w
229 lines
8.8 KiB
Python
229 lines
8.8 KiB
Python
"""工具注册统一配置模块"""
|
||
|
||
from typing import List, Optional
|
||
from pathlib import Path
|
||
from loguru import logger
|
||
|
||
from backend.modules.tools.registry import ToolRegistry
|
||
|
||
|
||
def register_all_tools(
|
||
workspace: Path,
|
||
command_timeout: int = 180,
|
||
max_output_length: int = 10000,
|
||
allow_dangerous: bool = False,
|
||
restrict_to_workspace: bool = True,
|
||
custom_deny_patterns: Optional[List[str]] = None,
|
||
custom_allow_patterns: Optional[List[str]] = None,
|
||
audit_log_enabled: bool = True,
|
||
subagent_manager=None,
|
||
skills_loader=None,
|
||
session_id: Optional[str] = None,
|
||
channel_manager=None,
|
||
session_manager=None,
|
||
memory_store=None,
|
||
config_loader=None,
|
||
) -> ToolRegistry:
|
||
"""
|
||
注册所有可用工具
|
||
|
||
Args:
|
||
workspace: 工作空间路径
|
||
command_timeout: 命令超时时间（秒）
|
||
max_output_length: 最大输出长度（字符）
|
||
allow_dangerous: 是否允许危险命令
|
||
restrict_to_workspace: 是否限制在工作空间内
|
||
custom_deny_patterns: 自定义拒绝模式列表
|
||
custom_allow_patterns: 自定义允许模式列表
|
||
audit_log_enabled: 是否启用审计日志
|
||
subagent_manager: SubagentManager 实例（可选）
|
||
skills_loader: SkillsLoader 实例（可选，用于检查禁用的技能）
|
||
session_id: 会话 ID（可选，用于审计日志）
|
||
channel_manager: ChannelManager 实例（可选）
|
||
session_manager: SessionManager 实例（可选）
|
||
memory_store: MemoryStore 实例（可选，用于记忆工具）
|
||
|
||
Returns:
|
||
ToolRegistry: 已注册所有工具的注册表
|
||
"""
|
||
tools = ToolRegistry()
|
||
|
||
# 配置审计日志
|
||
tools.set_audit_enabled(audit_log_enabled)
|
||
if session_id:
|
||
tools.set_session_id(session_id)
|
||
|
||
# 1. 注册文件系统工具（AI 可用，但前端隐藏）
|
||
from backend.modules.tools.filesystem import (
|
||
ReadFileTool,
|
||
WriteFileTool,
|
||
EditFileTool,
|
||
ListDirTool,
|
||
)
|
||
|
||
tools.register(ReadFileTool(workspace, skills_loader=skills_loader, restrict_to_workspace=restrict_to_workspace))
|
||
tools.register(WriteFileTool(workspace, restrict_to_workspace=restrict_to_workspace))
|
||
tools.register(EditFileTool(workspace, restrict_to_workspace=restrict_to_workspace))
|
||
tools.register(ListDirTool(workspace, restrict_to_workspace=restrict_to_workspace))
|
||
logger.debug("Registered filesystem tools")
|
||
|
||
# 2. 注册 Shell 工具（AI 可用，但前端隐藏）
|
||
from backend.modules.tools.shell import ExecTool
|
||
|
||
# 合并自定义拒绝模式
|
||
deny_patterns = None
|
||
if custom_deny_patterns:
|
||
from backend.modules.tools.shell import DANGEROUS_PATTERNS
|
||
deny_patterns = list(DANGEROUS_PATTERNS) + custom_deny_patterns
|
||
|
||
tools.register(
|
||
ExecTool(
|
||
workspace=workspace,
|
||
timeout=command_timeout,
|
||
max_output_length=max_output_length,
|
||
allow_dangerous=allow_dangerous,
|
||
deny_patterns=deny_patterns,
|
||
allow_patterns=custom_allow_patterns,
|
||
restrict_to_workspace=restrict_to_workspace,
|
||
)
|
||
)
|
||
logger.debug(
|
||
f"Registered shell tools (dangerous_blocked={not allow_dangerous}, "
|
||
f"workspace_restricted={restrict_to_workspace})"
|
||
)
|
||
|
||
# 2b. 注册外部编程代理工具
|
||
try:
|
||
from backend.modules.tools.external_coding_agent import ExternalCodingAgentTool
|
||
|
||
external_coding_tool = ExternalCodingAgentTool(
|
||
workspace=workspace,
|
||
default_timeout=command_timeout,
|
||
max_output_length=max_output_length,
|
||
)
|
||
enabled_profiles = external_coding_tool.registry.enabled_profile_names()
|
||
if enabled_profiles:
|
||
if session_id:
|
||
external_coding_tool.set_session_id(session_id)
|
||
tools.register(external_coding_tool)
|
||
logger.debug(
|
||
"Registered external coding agent tool with enabled profiles: {}",
|
||
", ".join(enabled_profiles),
|
||
)
|
||
else:
|
||
logger.info("Skipped external coding agent tool registration: no enabled profiles")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register external coding agent tool: {e}")
|
||
|
||
# 3. 注册 Web 工具
|
||
try:
|
||
from backend.modules.tools.web import WebFetchTool
|
||
|
||
tools.register(WebFetchTool())
|
||
logger.debug("Registered web fetch tool")
|
||
except ImportError:
|
||
logger.warning("Web tools not available")
|
||
|
||
# 4. 注册 Spawn 工具（如果提供了 SubagentManager）
|
||
if subagent_manager is not None:
|
||
try:
|
||
from backend.modules.tools.spawn import SpawnTool
|
||
|
||
spawn_tool = SpawnTool(subagent_manager, config_loader=config_loader)
|
||
tools.register(spawn_tool)
|
||
logger.debug("Registered spawn tool")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register spawn tool: {e}")
|
||
|
||
# 4b. 注册 Workflow 工具（如果提供了 SubagentManager）
|
||
if subagent_manager is not None:
|
||
try:
|
||
from backend.modules.tools.workflow_tool import WorkflowTool
|
||
|
||
workflow_tool = WorkflowTool(subagent_manager, skills=skills_loader)
|
||
if session_id:
|
||
workflow_tool.set_session_id(session_id)
|
||
tools.register(workflow_tool)
|
||
logger.debug("Registered workflow_run tool")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register workflow_run tool: {e}")
|
||
|
||
# 5. 注册发送媒体工具（如果提供了 ChannelManager）
|
||
if channel_manager is not None:
|
||
try:
|
||
from backend.modules.tools.send_media import SendMediaTool
|
||
|
||
send_media_tool = SendMediaTool(
|
||
channel_manager=channel_manager,
|
||
session_manager=session_manager
|
||
)
|
||
# 设置当前会话 ID
|
||
if session_id:
|
||
send_media_tool.set_session_id(session_id)
|
||
tools.register(send_media_tool)
|
||
logger.debug("Registered send_media tool")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register send_media tool: {e}")
|
||
|
||
# 6. 注册截图工具
|
||
try:
|
||
from backend.modules.tools.screenshot import ScreenshotTool
|
||
|
||
screenshot_tool = ScreenshotTool(workspace=workspace)
|
||
tools.register(screenshot_tool)
|
||
logger.debug("Registered screenshot tool")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register screenshot tool: {e}")
|
||
|
||
# 7. 注册文件搜索工具
|
||
try:
|
||
from backend.modules.tools.file_search import FileSearchTool
|
||
|
||
file_search_tool = FileSearchTool(
|
||
workspace=workspace,
|
||
default_max_results=20,
|
||
restrict_to_workspace=restrict_to_workspace,
|
||
)
|
||
tools.register(file_search_tool)
|
||
logger.debug("Registered file search tool")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register file search tool: {e}")
|
||
|
||
# 8. 注册记忆工具（合并为单一工具，减少 token 消耗）
|
||
if memory_store is not None:
|
||
try:
|
||
from backend.modules.tools.memory_tool import MemoryTool
|
||
|
||
tools.register(MemoryTool(memory_store))
|
||
logger.debug("Registered memory tool (unified)")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register memory tool: {e}")
|
||
|
||
# 8b. 注册 Wiki 知识库工具
|
||
try:
|
||
from backend.modules.wiki.tool import WikiTool
|
||
|
||
wiki_dir = workspace / "wiki"
|
||
wiki_dir.mkdir(parents=True, exist_ok=True)
|
||
wiki_tool = WikiTool(wiki_dir)
|
||
tools.register(wiki_tool)
|
||
logger.debug("Registered wiki tool")
|
||
except Exception as e:
|
||
logger.error(f"Failed to register wiki tool: {e}")
|
||
|
||
# 9. 注册小智AI send_message 工具（仅在小智频道启用且对话模式开启时注册）
|
||
try:
|
||
from backend.modules.config.loader import config_loader as _cl
|
||
xiaozhi_cfg = getattr(getattr(_cl.config, "channels", None), "xiaozhi", None)
|
||
if xiaozhi_cfg and xiaozhi_cfg.enabled and getattr(xiaozhi_cfg, "enable_conversation", False):
|
||
from backend.modules.tools.xiaozhi_message import XiaozhiMessageTool
|
||
tools.register(XiaozhiMessageTool())
|
||
logger.debug("✓ Registered xiaozhi send_message tool (conversation mode enabled)")
|
||
else:
|
||
logger.debug(f"Xiaozhi send_message tool not registered - enabled={xiaozhi_cfg.enabled if xiaozhi_cfg else False}, conversation={getattr(xiaozhi_cfg, 'enable_conversation', False) if xiaozhi_cfg else False}")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to register xiaozhi send_message tool: {e}")
|
||
|
||
logger.debug(f"Registered {len(tools.get_definitions())} tools")
|
||
return tools
|