From f2c0c2a9dee0dc103192cd98d0cde803c40faa7d Mon Sep 17 00:00:00 2001 From: LIghtJUNction Date: Mon, 23 Mar 2026 12:48:34 +0800 Subject: [PATCH] feat: add _internal modules for MCP, Skills, and Tools This commit adds the _internal package structure for AstrBot's standardized MCP & Skills support: astrbot/_internal/mcp/: - MCPClient for MCP server connections - MCPTool wrapper for MCP tools - MCP configuration management astrbot/_internal/skills/: - SkillManager for skill lifecycle - Skill parser and loader - SkillToToolConverter for tool-based skills - Prompt builder for skills astrbot/_internal/tools/: - ToolSchema, FunctionTool, ToolSet base definitions - FunctionToolManager for tool registry - Builtin tools (cron, send_message, kb_query) - Tool providers (internal, plugin, computer) astrbot/api/: - Public API for tools (ToolRegistry, tool decorator) - Public API for MCP (get_mcp_servers, register_mcp_server) - Public API for skills (get_skill_manager, skill_to_tool) --- astrbot/_internal/mcp/__init__.py | 63 ++ astrbot/_internal/mcp/client.py | 411 +++++++++ astrbot/_internal/mcp/config.py | 55 ++ astrbot/_internal/mcp/tool.py | 40 + astrbot/_internal/skills/__init__.py | 37 + astrbot/_internal/skills/loader.py | 79 ++ astrbot/_internal/skills/manager.py | 617 +++++++++++++ astrbot/_internal/skills/parser.py | 81 ++ astrbot/_internal/skills/prompt_builder.py | 158 ++++ astrbot/_internal/skills/to_tool.py | 79 ++ astrbot/_internal/tools/__init__.py | 9 + astrbot/_internal/tools/base.py | 442 ++++++++++ astrbot/_internal/tools/builtin/__init__.py | 52 ++ astrbot/_internal/tools/builtin/cron.py | 201 +++++ astrbot/_internal/tools/builtin/kb_query.py | 139 +++ .../_internal/tools/builtin/send_message.py | 226 +++++ astrbot/_internal/tools/providers/__init__.py | 25 + astrbot/_internal/tools/providers/computer.py | 45 + astrbot/_internal/tools/providers/internal.py | 73 ++ astrbot/_internal/tools/providers/plugin.py | 73 ++ astrbot/_internal/tools/registry.py | 821 ++++++++++++++++++ astrbot/api/__init__.py | 45 + astrbot/api/mcp.py | 93 ++ astrbot/api/skills.py | 52 ++ astrbot/api/tools.py | 110 +++ 25 files changed, 4026 insertions(+) create mode 100644 astrbot/_internal/mcp/__init__.py create mode 100644 astrbot/_internal/mcp/client.py create mode 100644 astrbot/_internal/mcp/config.py create mode 100644 astrbot/_internal/mcp/tool.py create mode 100644 astrbot/_internal/skills/__init__.py create mode 100644 astrbot/_internal/skills/loader.py create mode 100644 astrbot/_internal/skills/manager.py create mode 100644 astrbot/_internal/skills/parser.py create mode 100644 astrbot/_internal/skills/prompt_builder.py create mode 100644 astrbot/_internal/skills/to_tool.py create mode 100644 astrbot/_internal/tools/__init__.py create mode 100644 astrbot/_internal/tools/base.py create mode 100644 astrbot/_internal/tools/builtin/__init__.py create mode 100644 astrbot/_internal/tools/builtin/cron.py create mode 100644 astrbot/_internal/tools/builtin/kb_query.py create mode 100644 astrbot/_internal/tools/builtin/send_message.py create mode 100644 astrbot/_internal/tools/providers/__init__.py create mode 100644 astrbot/_internal/tools/providers/computer.py create mode 100644 astrbot/_internal/tools/providers/internal.py create mode 100644 astrbot/_internal/tools/providers/plugin.py create mode 100644 astrbot/_internal/tools/registry.py create mode 100644 astrbot/api/mcp.py create mode 100644 astrbot/api/skills.py create mode 100644 astrbot/api/tools.py diff --git a/astrbot/_internal/mcp/__init__.py b/astrbot/_internal/mcp/__init__.py new file mode 100644 index 000000000..a496c57e1 --- /dev/null +++ b/astrbot/_internal/mcp/__init__.py @@ -0,0 +1,63 @@ +"""MCP module - Model Context Protocol client and tool implementations. + +This module provides MCP client functionality and MCP tool wrappers. +""" + +import asyncio +from dataclasses import dataclass + +from .client import MCPClient +from .config import ( + DEFAULT_MCP_CONFIG, + get_mcp_config_path, + load_mcp_config, + save_mcp_config, +) +from .tool import MCPTool + + +# Exceptions +class MCPInitError(Exception): + """Base exception for MCP initialization failures.""" + + +class MCPInitTimeoutError(asyncio.TimeoutError, MCPInitError): + """Raised when MCP client initialization exceeds the configured timeout.""" + + +class MCPAllServicesFailedError(MCPInitError): + """Raised when all configured MCP services fail to initialize.""" + + +class MCPShutdownTimeoutError(asyncio.TimeoutError): + """Raised when MCP shutdown exceeds the configured timeout.""" + + def __init__(self, names: list[str], timeout: float) -> None: + self.names = names + self.timeout = timeout + message = f"MCP 服务关闭超时({timeout:g} 秒):{', '.join(names)}" + super().__init__(message) + + +@dataclass +class MCPInitSummary: + """Summary of MCP initialization results.""" + + total: int + success: int + failed: list[str] + + +__all__ = [ + "DEFAULT_MCP_CONFIG", + "MCPAllServicesFailedError", + "MCPClient", + "MCPInitError", + "MCPInitSummary", + "MCPInitTimeoutError", + "MCPShutdownTimeoutError", + "MCPTool", + "get_mcp_config_path", + "load_mcp_config", + "save_mcp_config", +] diff --git a/astrbot/_internal/mcp/client.py b/astrbot/_internal/mcp/client.py new file mode 100644 index 000000000..d4a2db7b2 --- /dev/null +++ b/astrbot/_internal/mcp/client.py @@ -0,0 +1,411 @@ +"""MCP client implementation.""" + +import asyncio +import logging +import os +import sys +from contextlib import AsyncExitStack +from datetime import timedelta + +from tenacity import ( + before_sleep_log, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) + +from astrbot import logger +from astrbot.core.utils.log_pipe import LogPipe + +try: + import anyio + import mcp + from mcp.client.sse import sse_client +except (ModuleNotFoundError, ImportError): + logger.warning( + "Warning: Missing 'mcp' dependency, MCP services will be unavailable." + ) + +try: + from mcp.client.streamable_http import streamablehttp_client +except (ModuleNotFoundError, ImportError): + logger.warning( + "Warning: Missing 'mcp' dependency or MCP library version too old, Streamable HTTP connection unavailable.", + ) + + +def _prepare_config(config: dict) -> dict: + """Prepare configuration, handle nested format.""" + if config.get("mcpServers"): + first_key = next(iter(config["mcpServers"])) + config = config["mcpServers"][first_key] + config.pop("active", None) + return config + + +def _prepare_stdio_env(config: dict) -> dict: + """Preserve Windows executable resolution for stdio subprocesses.""" + if sys.platform != "win32": + return config + + pathext = os.environ.get("PATHEXT") + if not pathext: + return config + + prepared = config.copy() + env = dict(prepared.get("env") or {}) + env.setdefault("PATHEXT", pathext) + prepared["env"] = env + return prepared + + +async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: + """Quick test MCP server connectivity.""" + import aiohttp + + cfg = _prepare_config(config.copy()) + + url = cfg["url"] + headers = cfg.get("headers", {}) + timeout = cfg.get("timeout", 10) + + try: + if "transport" in cfg: + transport_type = cfg["transport"] + elif "type" in cfg: + transport_type = cfg["type"] + else: + raise Exception("MCP connection config missing transport or type field") + + async with aiohttp.ClientSession() as session: + if transport_type == "streamable_http": + test_payload = { + "jsonrpc": "2.0", + "method": "initialize", + "id": 0, + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "test-client", "version": "1.2.3"}, + }, + } + async with session.post( + url, + headers={ + **headers, + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + }, + json=test_payload, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status == 200: + return True, "" + return False, f"HTTP {response.status}: {response.reason}" + else: + async with session.get( + url, + headers={ + **headers, + "Accept": "application/json, text/event-stream", + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status == 200: + return True, "" + return False, f"HTTP {response.status}: {response.reason}" + + except asyncio.TimeoutError: + return False, f"Connection timeout: {timeout} seconds" + except Exception as e: + return False, f"{e!s}" + + +class MCPClient: + def __init__(self) -> None: + # Initialize session and client objects + self.session: mcp.ClientSession | None = None + self.exit_stack = AsyncExitStack() + self._old_exit_stacks: list[AsyncExitStack] = [] # Track old stacks for cleanup + + self.name: str | None = None + self.active: bool = True + self.tools: list[mcp.Tool] = [] + self.server_errlogs: list[str] = [] + self.running_event = asyncio.Event() + self.process_pid: int | None = None + + # Store connection config for reconnection + self._mcp_server_config: dict | None = None + self._server_name: str | None = None + self._reconnect_lock = asyncio.Lock() # Lock for thread-safe reconnection + self._reconnecting: bool = False # For logging and debugging + + @staticmethod + def _extract_stdio_process_pid(streams_context: object) -> int | None: + """Best-effort extraction for stdio subprocess PID used by lease cleanup. + + TODO(refactor): replace this async-generator frame introspection with a + stable MCP library hook once the upstream transport exposes process PID. + """ + generator = getattr(streams_context, "gen", None) + frame = getattr(generator, "ag_frame", None) + if frame is None: + return None + process = frame.f_locals.get("process") + pid = getattr(process, "pid", None) + try: + return int(pid) if pid is not None else None + except (TypeError, ValueError): + return None + + async def connect_to_server(self, mcp_server_config: dict, name: str) -> None: + """Connect to MCP server + + If `url` parameter exists: + 1. When transport is specified as `streamable_http`, use Streamable HTTP connection. + 2. When transport is specified as `sse`, use SSE connection. + 3. If not specified, default to SSE connection to MCP service. + + Args: + mcp_server_config (dict): Configuration for the MCP server. See https://modelcontextprotocol.io/quickstart/server + + """ + # Store config for reconnection + self._mcp_server_config = mcp_server_config + self._server_name = name + self.process_pid = None + + cfg = _prepare_config(mcp_server_config.copy()) + + def logging_callback( + msg: str | mcp.types.LoggingMessageNotificationParams, + ) -> None: + # Handle MCP service error logs + if isinstance(msg, mcp.types.LoggingMessageNotificationParams): + if msg.level in ("warning", "error", "critical", "alert", "emergency"): + log_msg = f"[{msg.level.upper()}] {msg.data!s}" + self.server_errlogs.append(log_msg) + + if "url" in cfg: + success, error_msg = await _quick_test_mcp_connection(cfg) + if not success: + raise Exception(error_msg) + + if "transport" in cfg: + transport_type = cfg["transport"] + elif "type" in cfg: + transport_type = cfg["type"] + else: + raise Exception("MCP connection config missing transport or type field") + + if transport_type != "streamable_http": + # SSE transport method + self._streams_context = sse_client( + url=cfg["url"], + headers=cfg.get("headers", {}), + timeout=cfg.get("timeout", 5), + sse_read_timeout=cfg.get("sse_read_timeout", 60 * 5), + ) + streams = await self.exit_stack.enter_async_context( + self._streams_context, + ) + + # Create a new client session + read_timeout = timedelta(seconds=cfg.get("session_read_timeout", 60)) + self.session = await self.exit_stack.enter_async_context( + mcp.ClientSession( + *streams, + read_timeout_seconds=read_timeout, + logging_callback=logging_callback, # type: ignore + ), + ) + else: + timeout = timedelta(seconds=cfg.get("timeout", 30)) + sse_read_timeout = timedelta( + seconds=cfg.get("sse_read_timeout", 60 * 5), + ) + self._streams_context = streamablehttp_client( + url=cfg["url"], + headers=cfg.get("headers", {}), + timeout=timeout, + sse_read_timeout=sse_read_timeout, + terminate_on_close=cfg.get("terminate_on_close", True), + ) + read_s, write_s, _ = await self.exit_stack.enter_async_context( + self._streams_context, + ) + + # Create a new client session + read_timeout = timedelta(seconds=cfg.get("session_read_timeout", 60)) + self.session = await self.exit_stack.enter_async_context( + mcp.ClientSession( + read_stream=read_s, + write_stream=write_s, + read_timeout_seconds=read_timeout, + logging_callback=logging_callback, # type: ignore + ), + ) + + else: + cfg = _prepare_stdio_env(cfg) + server_params = mcp.StdioServerParameters( + **cfg, + ) + + def callback(msg: str | mcp.types.LoggingMessageNotificationParams) -> None: + # Handle MCP service error logs + if isinstance(msg, mcp.types.LoggingMessageNotificationParams): + if msg.level in ( + "warning", + "error", + "critical", + "alert", + "emergency", + ): + log_msg = f"[{msg.level.upper()}] {msg.data!s}" + self.server_errlogs.append(log_msg) + + stdio_transport = await self.exit_stack.enter_async_context( + mcp.stdio_client( + server_params, + errlog=LogPipe( + level=logging.INFO, + logger=logger, + identifier=f"MCPServer-{name}", + callback=callback, + ), # type: ignore + ), + ) + self.process_pid = self._extract_stdio_process_pid(self._streams_context) + + # Create a new client session + self.session = await self.exit_stack.enter_async_context( + mcp.ClientSession(*stdio_transport), + ) + await self.session.initialize() + + async def list_tools_and_save(self) -> mcp.ListToolsResult: + """List all tools from the server and save them to self.tools""" + if not self.session: + raise Exception("MCP Client is not initialized") + response = await self.session.list_tools() + self.tools = response.tools + return response + + async def _reconnect(self) -> None: + """Reconnect to the MCP server using the stored configuration. + + Uses asyncio.Lock to ensure thread-safe reconnection in concurrent environments. + + Raises: + Exception: raised when reconnection fails + """ + async with self._reconnect_lock: + # Check if already reconnecting (useful for logging) + if self._reconnecting: + logger.debug( + f"MCP Client {self._server_name} is already reconnecting, skipping" + ) + return + + if not self._mcp_server_config or not self._server_name: + raise Exception("Cannot reconnect: missing connection configuration") + + self._reconnecting = True + try: + logger.info( + f"Attempting to reconnect to MCP server {self._server_name}..." + ) + + # Save old exit_stack for later cleanup (don't close it now to avoid cancel scope issues) + if self.exit_stack: + self._old_exit_stacks.append(self.exit_stack) + + # Mark old session as invalid + self.session = None + + # Create new exit stack for new connection + self.exit_stack = AsyncExitStack() + + # Reconnect using stored config + await self.connect_to_server(self._mcp_server_config, self._server_name) + await self.list_tools_and_save() + + logger.info( + f"Successfully reconnected to MCP server {self._server_name}" + ) + except Exception as e: + logger.error( + f"Failed to reconnect to MCP server {self._server_name}: {e}" + ) + raise + finally: + self._reconnecting = False + + async def call_tool_with_reconnect( + self, + tool_name: str, + arguments: dict, + read_timeout_seconds: timedelta, + ) -> mcp.types.CallToolResult: + """Call MCP tool with automatic reconnection on failure, max 2 retries. + + Args: + tool_name: tool name + arguments: tool arguments + read_timeout_seconds: read timeout + + Returns: + MCP tool call result + + Raises: + ValueError: MCP session is not available + anyio.ClosedResourceError: raised after reconnection failure + """ + + @retry( + retry=retry_if_exception_type(anyio.ClosedResourceError), + stop=stop_after_attempt(2), + wait=wait_exponential(multiplier=1, min=1, max=3), + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, + ) + async def _call_with_retry(): + if not self.session: + raise ValueError("MCP session is not available for MCP function tools.") + + try: + return await self.session.call_tool( + name=tool_name, + arguments=arguments, + read_timeout_seconds=read_timeout_seconds, + ) + except anyio.ClosedResourceError: + logger.warning( + f"MCP tool {tool_name} call failed (ClosedResourceError), attempting to reconnect..." + ) + # Attempt to reconnect + await self._reconnect() + # Reraise the exception to trigger tenacity retry + raise + + return await _call_with_retry() + + async def cleanup(self) -> None: + """Clean up resources including old exit stacks from reconnections""" + # Close current exit stack + try: + await self.exit_stack.aclose() + except Exception as e: + logger.debug(f"Error closing current exit stack: {e}") + + # Don't close old exit stacks as they may be in different task contexts + # They will be garbage collected naturally + # Just clear the list to release references + self._old_exit_stacks.clear() + + # Set running_event first to unblock any waiting tasks + self.running_event.set() + self.process_pid = None diff --git a/astrbot/_internal/mcp/config.py b/astrbot/_internal/mcp/config.py new file mode 100644 index 000000000..0ea528948 --- /dev/null +++ b/astrbot/_internal/mcp/config.py @@ -0,0 +1,55 @@ +"""MCP configuration management.""" + +import json +import os + +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + +DEFAULT_MCP_CONFIG = {"mcpServers": {}} + + +def get_mcp_config_path() -> str: + """Get the path to the MCP configuration file.""" + data_dir = get_astrbot_data_path() + return os.path.join(data_dir, "mcp_server.json") + + +def load_mcp_config() -> dict: + """Load MCP configuration from file. + + Returns: + MCP configuration dict. If file doesn't exist, returns default config. + + """ + config_path = get_mcp_config_path() + if not os.path.exists(config_path): + # Create default config if not exists + os.makedirs(os.path.dirname(config_path), exist_ok=True) + with open(config_path, "w", encoding="utf-8") as f: + json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4) + return DEFAULT_MCP_CONFIG + + try: + with open(config_path, encoding="utf-8") as f: + return json.load(f) + except Exception: + return DEFAULT_MCP_CONFIG + + +def save_mcp_config(config: dict) -> bool: + """Save MCP configuration to file. + + Args: + config: MCP configuration dict to save. + + Returns: + True if successful, False otherwise. + + """ + config_path = get_mcp_config_path() + try: + with open(config_path, "w", encoding="utf-8") as f: + json.dump(config, f, ensure_ascii=False, indent=4) + return True + except Exception: + return False diff --git a/astrbot/_internal/mcp/tool.py b/astrbot/_internal/mcp/tool.py new file mode 100644 index 000000000..573784da9 --- /dev/null +++ b/astrbot/_internal/mcp/tool.py @@ -0,0 +1,40 @@ +"""MCP tool wrapper.""" + +from datetime import timedelta +from typing import Generic + +from astrbot._internal.tools.base import FunctionTool +from astrbot.core.agent.run_context import ContextWrapper, TContext + +from .client import MCPClient + +try: + import mcp +except (ModuleNotFoundError, ImportError): + mcp = None # type: ignore + + +class MCPTool(FunctionTool, Generic[TContext]): + """A function tool that calls an MCP service.""" + + def __init__( + self, mcp_tool: mcp.Tool, mcp_client: MCPClient, mcp_server_name: str, **kwargs + ) -> None: + super().__init__( + name=mcp_tool.name, + description=mcp_tool.description or "", + parameters=mcp_tool.inputSchema, + ) + self.mcp_tool = mcp_tool + self.mcp_client = mcp_client + self.mcp_server_name = mcp_server_name + self.source = "mcp" + + async def call( + self, context: ContextWrapper[TContext], **kwargs + ) -> mcp.types.CallToolResult: + return await self.mcp_client.call_tool_with_reconnect( + tool_name=self.mcp_tool.name, + arguments=kwargs, + read_timeout_seconds=timedelta(seconds=context.tool_call_timeout), + ) diff --git a/astrbot/_internal/skills/__init__.py b/astrbot/_internal/skills/__init__.py new file mode 100644 index 000000000..78c77a109 --- /dev/null +++ b/astrbot/_internal/skills/__init__.py @@ -0,0 +1,37 @@ +"""AstrBot internal skills module. + +This module provides the skill management system for AstrBot, including: +- SkillManager: Manages skill lifecycle (install, activate, delete, etc.) +- SkillInfo: Dataclass representing skill metadata +- build_skills_prompt: Builds the system prompt section for skills +- SkillToToolConverter: Converts skills with input_schema to FunctionTool +""" + +from .manager import ( + DEFAULT_SKILLS_CONFIG, + SANDBOX_SKILLS_CACHE_FILENAME, + SANDBOX_SKILLS_ROOT, + SANDBOX_WORKSPACE_ROOT, + SKILLS_CONFIG_FILENAME, + SkillInfo, + SkillManager, + build_skills_prompt, +) +from .parser import parse_frontmatter, parse_skill_markdown +from .to_tool import SkillToToolConverter + +__all__ = [ + # Constants + "DEFAULT_SKILLS_CONFIG", + "SANDBOX_SKILLS_CACHE_FILENAME", + "SANDBOX_SKILLS_ROOT", + "SANDBOX_WORKSPACE_ROOT", + "SKILLS_CONFIG_FILENAME", + "SkillInfo", + "SkillManager", + "SkillToToolConverter", + "build_skills_prompt", + # Parser + "parse_frontmatter", + "parse_skill_markdown", +] diff --git a/astrbot/_internal/skills/loader.py b/astrbot/_internal/skills/loader.py new file mode 100644 index 000000000..ae44d4665 --- /dev/null +++ b/astrbot/_internal/skills/loader.py @@ -0,0 +1,79 @@ +"""Skill loader - loads skills from filesystem.""" + +from __future__ import annotations + +from pathlib import Path + +from .manager import SkillInfo, _normalize_skill_markdown_path, _parse_frontmatter + + +def load_skill(skill_name: str, skills_root: Path) -> SkillInfo | None: + """Load a single skill by name from the skills root directory. + + Args: + skill_name: The name of the skill to load + skills_root: Path to the skills root directory + + Returns: + SkillInfo if the skill exists and has a valid SKILL.md, None otherwise + """ + skill_dir = skills_root / skill_name + if not skill_dir.is_dir(): + return None + + skill_md = _normalize_skill_markdown_path(skill_dir) + if skill_md is None: + return None + + try: + content = skill_md.read_text(encoding="utf-8") + meta = _parse_frontmatter(content) + except Exception: + return None + + description = meta.get("description", "") + if not isinstance(description, str): + description = "" + description = description.strip() + + input_schema = meta.get("input_schema") + output_schema = meta.get("output_schema") + + return SkillInfo( + name=skill_name, + description=description, + path=str(skill_md).replace("\\", "/"), + active=True, + source_type="local_only", + source_label="local", + local_exists=True, + sandbox_exists=False, + input_schema=input_schema, + output_schema=output_schema, + ) + + +def load_all_skills(skills_root: Path) -> list[SkillInfo]: + """Load all skills from the skills root directory. + + Args: + skills_root: Path to the skills root directory + + Returns: + List of SkillInfo objects for all valid skills + """ + skills: list[SkillInfo] = [] + + if not skills_root.is_dir(): + return skills + + for entry in sorted(skills_root.iterdir()): + if not entry.is_dir(): + continue + + skill_name = entry.name + skill_info = load_skill(skill_name, skills_root) + if skill_info is not None: + skills.append(skill_info) + + return skills diff --git a/astrbot/_internal/skills/manager.py b/astrbot/_internal/skills/manager.py new file mode 100644 index 000000000..9e54ced39 --- /dev/null +++ b/astrbot/_internal/skills/manager.py @@ -0,0 +1,617 @@ +"""Skill manager - manages skill lifecycle (install, activate, delete, etc.).""" + +from __future__ import annotations + +import json +import os +import re +import shlex +import shutil +import tempfile +import uuid +import zipfile +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath + +import yaml + +from astrbot.core.utils.astrbot_path import AstrbotPaths, astrbot_paths + +SKILLS_CONFIG_FILENAME = "skills.json" +SANDBOX_SKILLS_CACHE_FILENAME = "sandbox_skills_cache.json" +DEFAULT_SKILLS_CONFIG: dict[str, dict] = {"skills": {}} +SANDBOX_SKILLS_ROOT = "skills" +SANDBOX_WORKSPACE_ROOT = "/workspace" +_SANDBOX_SKILLS_CACHE_VERSION = 1 + +_SKILL_NAME_RE = re.compile(r"^[A-Za-z0-9._-]+$") + + +def _default_sandbox_skill_path(name: str) -> str: + return f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/{name}/SKILL.md" + + +def _normalize_cached_sandbox_skill_path(name: str, path: str) -> str: + normalized = str(path or "").strip().replace("\\", "/") + if not normalized: + return _default_sandbox_skill_path(name) + + pure_path = PurePosixPath(normalized) + if ".." in pure_path.parts: + return _default_sandbox_skill_path(name) + + if pure_path.name != "SKILL.md": + return _default_sandbox_skill_path(name) + + if pure_path.parent.name != name: + return _default_sandbox_skill_path(name) + + return str(pure_path) + + +def _is_ignored_zip_entry(name: str) -> bool: + parts = PurePosixPath(name).parts + if not parts: + return True + return parts[0] == "__MACOSX" + + +def _normalize_skill_markdown_path(skill_dir: Path) -> Path | None: + """Return the canonical `SKILL.md` path for a skill directory. + + If only legacy `skill.md` exists, it is renamed to `SKILL.md` in-place. + """ + canonical = skill_dir / "SKILL.md" + entries = set() + if skill_dir.exists(): + entries = {entry.name for entry in skill_dir.iterdir()} + if "SKILL.md" in entries: + return canonical + legacy = skill_dir / "skill.md" + if "skill.md" not in entries: + return None + try: + tmp = skill_dir / f".{uuid.uuid4().hex}.tmp_skill_md" + legacy.rename(tmp) + tmp.rename(canonical) + except OSError: + return legacy + return canonical + + +@dataclass +class SkillInfo: + name: str + description: str + path: str + active: bool + source_type: str = "local_only" + source_label: str = "local" + local_exists: bool = True + sandbox_exists: bool = False + input_schema: dict | None = None + output_schema: dict | None = None + + +def _parse_frontmatter(text: str) -> dict: + """Extract metadata from YAML frontmatter. + + Expects the standard SKILL.md format used by OpenAI Codex CLI and + Anthropic Claude Skills:: + + --- + name: my-skill + description: What this skill does and when to use it. + input_schema: ... + output_schema: ... + --- + """ + if not text.startswith("---"): + return {} + lines = text.splitlines() + if not lines or lines[0].strip() != "---": + return {} + end_idx = None + for i in range(1, len(lines)): + if lines[i].strip() == "---": + end_idx = i + break + if end_idx is None: + return {} + + frontmatter = "\n".join(lines[1:end_idx]) + try: + payload = yaml.safe_load(frontmatter) or {} + except yaml.YAMLError: + return {} + if not isinstance(payload, dict): + return {} + + return payload + + +# Regex for sanitizing paths used in prompt examples — only allow +# safe path characters to prevent prompt injection via crafted skill paths. +_SAFE_PATH_RE = re.compile(r"[^\w./ ,()'\-]", re.UNICODE) +_WINDOWS_DRIVE_PATH_RE = re.compile(r"^[A-Za-z]:(?:/|\\)") +_WINDOWS_UNC_PATH_RE = re.compile(r"^(//|\\\\)[^/\\]+[/\\][^/\\]+") +_CONTROL_CHARS_RE = re.compile(r"[\x00-\x1F\x7F]") + + +def _is_windows_prompt_path(path: str) -> bool: + if os.name != "nt": + return False + return bool(_WINDOWS_DRIVE_PATH_RE.match(path) or _WINDOWS_UNC_PATH_RE.match(path)) + + +def _sanitize_prompt_path_for_prompt(path: str) -> str: + if not path: + return "" + + if _WINDOWS_DRIVE_PATH_RE.match(path) or _WINDOWS_UNC_PATH_RE.match(path): + path = path.replace("\\", "/") + + drive_prefix = "" + if _WINDOWS_DRIVE_PATH_RE.match(path): + drive_prefix = path[:2] + path = path[2:] + + path = path.replace("`", "") + path = _CONTROL_CHARS_RE.sub("", path) + sanitized = _SAFE_PATH_RE.sub("", path) + return f"{drive_prefix}{sanitized}" + + +def _sanitize_prompt_description(description: str) -> str: + description = description.replace("`", "") + description = _CONTROL_CHARS_RE.sub(" ", description) + description = " ".join(description.split()) + return description + + +def _sanitize_skill_display_name(name: str) -> str: + if _SKILL_NAME_RE.fullmatch(name): + return name + return "" + + +def _build_skill_read_command_example(path: str) -> str: + if path == "//SKILL.md": + return f"cat {path}" + if _is_windows_prompt_path(path): + command = "type" + path_arg = f'"{os.path.normpath(path)}"' + else: + command = "cat" + path_arg = shlex.quote(path) + return f"{command} {path_arg}" + + +def build_skills_prompt(skills: list[SkillInfo]) -> str: + """Build the skills section of the system prompt. + + Generates a markdown-formatted skill inventory for the LLM. Only + ``name`` and ``description`` are shown upfront; the LLM must read + the full ``SKILL.md`` before execution (progressive disclosure). + """ + skills_lines: list[str] = [] + example_path = "" + for skill in skills: + display_name = _sanitize_skill_display_name(skill.name) + + description = skill.description or "No description" + if skill.source_type == "sandbox_only": + description = _sanitize_prompt_description(description) + if not description: + description = "Read SKILL.md for details." + + if skill.source_type == "sandbox_only": + # Prefer the actual path from sandbox cache if available + rendered_path = _sanitize_prompt_path_for_prompt(skill.path) + if not rendered_path: + rendered_path = _default_sandbox_skill_path(skill.name) + else: + rendered_path = _sanitize_prompt_path_for_prompt(skill.path) + if not rendered_path: + rendered_path = "//SKILL.md" + + entry = f"- **{display_name}**: {description}\n File: `{rendered_path}`" + if skill.input_schema: + entry += f"\n Input Schema: {json.dumps(skill.input_schema, ensure_ascii=False)}" + if skill.output_schema: + entry += f"\n Output Schema: {json.dumps(skill.output_schema, ensure_ascii=False)}" + skills_lines.append(entry) + if not example_path: + example_path = rendered_path + skills_block = "\n".join(skills_lines) + # Sanitize example_path — it may originate from sandbox cache (untrusted) + if example_path == "//SKILL.md": + example_path = "//SKILL.md" + else: + example_path = _sanitize_prompt_path_for_prompt(example_path) + example_path = example_path or "//SKILL.md" + example_command = _build_skill_read_command_example(example_path) + + return ( + "## Skills\n\n" + "You have specialized skills — reusable instruction bundles stored " + "in `SKILL.md` files. Each skill has a **name** and a **description** " + "that tells you what it does and when to use it.\n\n" + "### Available skills\n\n" + f"{skills_block}\n\n" + "### Skill rules\n\n" + "1. **Discovery** — The list above is the complete skill inventory " + "for this session. Full instructions are in the referenced " + "`SKILL.md` file.\n" + "2. **When to trigger** — Use a skill if the user names it " + "explicitly, or if the task clearly matches the skill's description. " + "*Never silently skip a matching skill* — either use it or briefly " + "explain why you chose not to.\n" + "3. **Mandatory grounding** — Before executing any skill you MUST " + "first read its `SKILL.md` by running a shell command compatible " + "with the current runtime shell and using the **absolute path** " + f"shown above (e.g. `{example_command}`). " + "Never rely on memory or assumptions about a skill's content.\n" + "4. **Progressive disclosure** — Load only what is directly " + "referenced from `SKILL.md`:\n" + " - If `scripts/` exist, prefer running or patching them over " + "rewriting code from scratch.\n" + " - If `assets/` or templates exist, reuse them.\n" + " - Do NOT bulk-load every file in the skill directory.\n" + "5. **Coordination** — When multiple skills apply, pick the minimal " + "set needed. Announce which skill(s) you are using and why " + "(one short line). Prefer `astrbot_*` tools when running skill " + "scripts.\n" + "6. **Context hygiene** — Avoid deep reference chasing; open only " + "files that are directly linked from `SKILL.md`.\n" + "7. **Failure handling** — If a skill cannot be applied, state the " + "issue clearly and continue with the best alternative.\n" + ) + + +class SkillManager: + def __init__( + self, + skills_root: str | None = None, + astrbot_paths: AstrbotPaths = astrbot_paths, + ) -> None: + self.astrbot_paths = astrbot_paths + self.skills_root = skills_root or str(self.astrbot_paths.skills) + self.config_path = str(self.astrbot_paths.config / SKILLS_CONFIG_FILENAME) + self.sandbox_skills_cache_path = str( + self.astrbot_paths.data / SANDBOX_SKILLS_CACHE_FILENAME + ) + os.makedirs(self.skills_root, exist_ok=True) + + def _load_config(self) -> dict: + if not os.path.exists(self.config_path): + self._save_config(DEFAULT_SKILLS_CONFIG.copy()) + return DEFAULT_SKILLS_CONFIG.copy() + with open(self.config_path, encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict) or "skills" not in data: + return DEFAULT_SKILLS_CONFIG.copy() + return data + + def _save_config(self, config: dict) -> None: + os.makedirs(os.path.dirname(self.config_path), exist_ok=True) + with open(self.config_path, "w", encoding="utf-8") as f: + json.dump(config, f, ensure_ascii=False, indent=4) + + def _load_sandbox_skills_cache(self) -> dict: + if not os.path.exists(self.sandbox_skills_cache_path): + return {"version": _SANDBOX_SKILLS_CACHE_VERSION, "skills": []} + try: + with open(self.sandbox_skills_cache_path, encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict): + return {"version": _SANDBOX_SKILLS_CACHE_VERSION, "skills": []} + skills = data.get("skills", []) + if not isinstance(skills, list): + skills = [] + return { + "version": int(data.get("version", _SANDBOX_SKILLS_CACHE_VERSION)), + "skills": skills, + "updated_at": data.get("updated_at"), + } + except Exception: + return {"version": _SANDBOX_SKILLS_CACHE_VERSION, "skills": []} + + def _save_sandbox_skills_cache(self, cache: dict) -> None: + cache["version"] = _SANDBOX_SKILLS_CACHE_VERSION + cache["updated_at"] = datetime.now(timezone.utc).isoformat() + os.makedirs(os.path.dirname(self.sandbox_skills_cache_path), exist_ok=True) + with open(self.sandbox_skills_cache_path, "w", encoding="utf-8") as f: + json.dump(cache, f, ensure_ascii=False, indent=2) + + def set_sandbox_skills_cache(self, skills: list[dict]) -> None: + """Persist sandbox skill metadata discovered from runtime side.""" + deduped: dict[str, dict[str, str]] = {} + for item in skills: + if not isinstance(item, dict): + continue + name = str(item.get("name", "")).strip() + if not name or not _SKILL_NAME_RE.match(name): + continue + description = str(item.get("description", "") or "") + path = _normalize_cached_sandbox_skill_path( + name, str(item.get("path", "") or "") + ) + deduped[name] = { + "name": name, + "description": description, + "path": path, + } + cache = { + "version": _SANDBOX_SKILLS_CACHE_VERSION, + "skills": [deduped[name] for name in sorted(deduped)], + } + self._save_sandbox_skills_cache(cache) + + def get_sandbox_skills_cache_status(self) -> dict[str, object]: + cache = self._load_sandbox_skills_cache() + skills = cache.get("skills", []) + count = len(skills) if isinstance(skills, list) else 0 + return { + "exists": os.path.exists(self.sandbox_skills_cache_path), + "ready": count > 0, + "count": count, + "updated_at": cache.get("updated_at"), + } + + def list_skills( + self, + *, + active_only: bool = False, + runtime: str = "local", + show_sandbox_path: bool = True, + ) -> list[SkillInfo]: + """List all skills. + + show_sandbox_path: If True and runtime is "sandbox", + return the path as it would appear in the sandbox environment, + otherwise return the local filesystem path. + """ + config = self._load_config() + skill_configs = config.get("skills", {}) + modified = False + skills_by_name: dict[str, SkillInfo] = {} + + sandbox_cached_paths: dict[str, str] = {} + sandbox_cached_descriptions: dict[str, str] = {} + cache_for_paths = self._load_sandbox_skills_cache() + for item in cache_for_paths.get("skills", []): + if not isinstance(item, dict): + continue + name = str(item.get("name", "") or "").strip() + path = _normalize_cached_sandbox_skill_path( + name, str(item.get("path", "") or "") + ) + if not name or not _SKILL_NAME_RE.match(name): + continue + sandbox_cached_descriptions[name] = str(item.get("description", "") or "") + sandbox_cached_paths[name] = path + + for entry in sorted(Path(self.skills_root).iterdir()): + if not entry.is_dir(): + continue + skill_name = entry.name + skill_md = _normalize_skill_markdown_path(entry) + if skill_md is None: + continue + active = skill_configs.get(skill_name, {}).get("active", True) + if skill_name not in skill_configs: + skill_configs[skill_name] = {"active": active} + modified = True + if active_only and not active: + continue + description = "" + input_schema = None + output_schema = None + try: + content = skill_md.read_text(encoding="utf-8") + meta = _parse_frontmatter(content) + description = meta.get("description", "") + if not isinstance(description, str): + description = "" + description = description.strip() + input_schema = meta.get("input_schema") + output_schema = meta.get("output_schema") + except Exception: + description = "" + sandbox_exists = ( + runtime == "sandbox" and skill_name in sandbox_cached_descriptions + ) + source_type = "both" if sandbox_exists else "local_only" + source_label = "synced" if sandbox_exists else "local" + if runtime == "sandbox" and show_sandbox_path: + path_str = sandbox_cached_paths.get( + skill_name + ) or _default_sandbox_skill_path(skill_name) + else: + path_str = str(skill_md) + path_str = path_str.replace("\\", "/") + skills_by_name[skill_name] = SkillInfo( + name=skill_name, + description=description, + path=path_str, + active=active, + source_type=source_type, + source_label=source_label, + local_exists=True, + sandbox_exists=sandbox_exists, + input_schema=input_schema, + output_schema=output_schema, + ) + + if runtime == "sandbox": + cache = self._load_sandbox_skills_cache() + for item in cache.get("skills", []): + if not isinstance(item, dict): + continue + skill_name = str(item.get("name", "")).strip() + if ( + not skill_name + or skill_name in skills_by_name + or not _SKILL_NAME_RE.match(skill_name) + ): + continue + active = skill_configs.get(skill_name, {}).get("active", True) + if skill_name not in skill_configs: + skill_configs[skill_name] = {"active": active} + modified = True + if active_only and not active: + continue + description = sandbox_cached_descriptions.get(skill_name, "") + # For sandbox_only skills, show_sandbox_path is implicitly True + # since there is no local path to show. Always prefer the + # actual path from sandbox cache. + path_str = sandbox_cached_paths.get( + skill_name + ) or _default_sandbox_skill_path(skill_name) + skills_by_name[skill_name] = SkillInfo( + name=skill_name, + description=description, + path=path_str.replace("\\", "/"), + active=active, + source_type="sandbox_only", + source_label="sandbox_preset", + local_exists=False, + sandbox_exists=True, + ) + + if modified: + config["skills"] = skill_configs + self._save_config(config) + + return [skills_by_name[name] for name in sorted(skills_by_name)] + + def is_sandbox_only_skill(self, name: str) -> bool: + skill_dir = Path(self.skills_root) / name + skill_md_exists = _normalize_skill_markdown_path(skill_dir) is not None + if skill_md_exists: + return False + cache = self._load_sandbox_skills_cache() + skills = cache.get("skills", []) + if not isinstance(skills, list): + return False + for item in skills: + if not isinstance(item, dict): + continue + if str(item.get("name", "")).strip() == name: + return True + return False + + def set_skill_active(self, name: str, active: bool) -> None: + if self.is_sandbox_only_skill(name): + raise PermissionError( + "Sandbox preset skill cannot be enabled/disabled from local skill management." + ) + config = self._load_config() + config.setdefault("skills", {}) + config["skills"][name] = {"active": bool(active)} + self._save_config(config) + + def _remove_skill_from_sandbox_cache(self, name: str) -> None: + cache = self._load_sandbox_skills_cache() + skills = cache.get("skills", []) + if not isinstance(skills, list): + return + + filtered = [ + item + for item in skills + if not ( + isinstance(item, dict) and str(item.get("name", "")).strip() == name + ) + ] + + if len(filtered) != len(skills): + cache["skills"] = filtered + self._save_sandbox_skills_cache(cache) + + def delete_skill(self, name: str) -> None: + if self.is_sandbox_only_skill(name): + raise PermissionError( + "Sandbox preset skill cannot be deleted from local skill management." + ) + + skill_dir = Path(self.skills_root) / name + if skill_dir.exists(): + shutil.rmtree(skill_dir) + + # Ensure UI consistency even when there is no active sandbox session + # to refresh cache from runtime side. + self._remove_skill_from_sandbox_cache(name) + + config = self._load_config() + if name in config.get("skills", {}): + config["skills"].pop(name, None) + self._save_config(config) + + def install_skill_from_zip(self, zip_path: str, *, overwrite: bool = True) -> str: + zip_path_obj = Path(zip_path) + if not zip_path_obj.exists(): + raise FileNotFoundError(f"Zip file not found: {zip_path}") + if not zipfile.is_zipfile(zip_path): + raise ValueError("Uploaded file is not a valid zip archive.") + + with zipfile.ZipFile(zip_path) as zf: + names = [ + name + for name in (entry.replace("\\", "/") for entry in zf.namelist()) + if name and not _is_ignored_zip_entry(name) + ] + file_names = [name for name in names if name and not name.endswith("/")] + if not file_names: + raise ValueError("Zip archive is empty.") + + top_dirs = { + PurePosixPath(name).parts[0] for name in file_names if name.strip() + } + + if len(top_dirs) != 1: + raise ValueError("Zip archive must contain a single top-level folder.") + skill_name = next(iter(top_dirs)) + if skill_name in {".", "..", ""} or not _SKILL_NAME_RE.match(skill_name): + raise ValueError("Invalid skill folder name.") + + for name in names: + if not name: + continue + if name.startswith("/") or re.match(r"^[A-Za-z]:", name): + raise ValueError("Zip archive contains absolute paths.") + parts = PurePosixPath(name).parts + if ".." in parts: + raise ValueError("Zip archive contains invalid relative paths.") + if parts and parts[0] != skill_name: + raise ValueError( + "Zip archive contains unexpected top-level entries." + ) + + if ( + f"{skill_name}/SKILL.md" not in file_names + and f"{skill_name}/skill.md" not in file_names + ): + raise ValueError("SKILL.md not found in the skill folder.") + + with tempfile.TemporaryDirectory(dir=str(astrbot_paths.temp)) as tmp_dir: + for member in zf.infolist(): + member_name = member.filename.replace("\\", "/") + if not member_name or _is_ignored_zip_entry(member_name): + continue + zf.extract(member, tmp_dir) + src_dir = Path(tmp_dir) / skill_name + _normalize_skill_markdown_path(src_dir) + if not src_dir.exists(): + raise ValueError("Skill folder not found after extraction.") + dest_dir = Path(self.skills_root) / skill_name + if dest_dir.exists(): + if not overwrite: + raise FileExistsError("Skill already exists.") + shutil.rmtree(dest_dir) + shutil.move(str(src_dir), str(dest_dir)) + + self.set_skill_active(skill_name, True) + return skill_name diff --git a/astrbot/_internal/skills/parser.py b/astrbot/_internal/skills/parser.py new file mode 100644 index 000000000..30c4134f8 --- /dev/null +++ b/astrbot/_internal/skills/parser.py @@ -0,0 +1,81 @@ +"""SKILL.md parser for extracting frontmatter and content.""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + + +def parse_frontmatter(text: str) -> dict: + """Extract metadata from YAML frontmatter. + + Expects the standard SKILL.md format used by OpenAI Codex CLI and + Anthropic Claude Skills:: + + --- + name: my-skill + description: What this skill does and when to use it. + input_schema: ... + output_schema: ... + --- + """ + if not text.startswith("---"): + return {} + lines = text.splitlines() + if not lines or lines[0].strip() != "---": + return {} + end_idx = None + for i in range(1, len(lines)): + if lines[i].strip() == "---": + end_idx = i + break + if end_idx is None: + return {} + + frontmatter = "\n".join(lines[1:end_idx]) + try: + payload = yaml.safe_load(frontmatter) or {} + except yaml.YAMLError: + return {} + if not isinstance(payload, dict): + return {} + + return payload + + +def parse_skill_markdown(path: Path) -> dict: + """Parse a SKILL.md file and return frontmatter + content. + + Args: + path: Path to the SKILL.md file + + Returns: + dict with keys: frontmatter (dict), content (str) + """ + try: + text = path.read_text(encoding="utf-8") + except Exception: + return {"frontmatter": {}, "content": ""} + + frontmatter = parse_frontmatter(text) + + # Extract content after frontmatter + if text.startswith("---"): + lines = text.splitlines() + end_idx = None + for i in range(1, len(lines)): + if lines[i].strip() == "---": + end_idx = i + break + if end_idx is not None: + content = "\n".join(lines[end_idx + 1 :]).strip() + else: + content = "" + else: + content = text + + return { + "frontmatter": frontmatter, + "content": content, + } diff --git a/astrbot/_internal/skills/prompt_builder.py b/astrbot/_internal/skills/prompt_builder.py new file mode 100644 index 000000000..940db1c22 --- /dev/null +++ b/astrbot/_internal/skills/prompt_builder.py @@ -0,0 +1,158 @@ +"""Skills prompt builder - builds the system prompt section for skills.""" + +from __future__ import annotations + +import json +import os +import re +import shlex + +from .manager import SkillInfo + +# Regex for sanitizing paths used in prompt examples — only allow +# safe path characters to prevent prompt injection via crafted skill paths. +_SAFE_PATH_RE = re.compile(r"[^\w./ ,()'\-]", re.UNICODE) +_WINDOWS_DRIVE_PATH_RE = re.compile(r"^[A-Za-z]:(?:/|\\)") +_WINDOWS_UNC_PATH_RE = re.compile(r"^(//|\\\\)[^/\\]+[/\\][^/\\]+") +_CONTROL_CHARS_RE = re.compile(r"[\x00-\x1F\x7F]") + +SANDBOX_WORKSPACE_ROOT = "/workspace" +SANDBOX_SKILLS_ROOT = "skills" + + +def _is_windows_prompt_path(path: str) -> bool: + if os.name != "nt": + return False + return bool(_WINDOWS_DRIVE_PATH_RE.match(path) or _WINDOWS_UNC_PATH_RE.match(path)) + + +def _sanitize_prompt_path_for_prompt(path: str) -> str: + if not path: + return "" + + if _WINDOWS_DRIVE_PATH_RE.match(path) or _WINDOWS_UNC_PATH_RE.match(path): + path = path.replace("\\", "/") + + drive_prefix = "" + if _WINDOWS_DRIVE_PATH_RE.match(path): + drive_prefix = path[:2] + path = path[2:] + + path = path.replace("`", "") + path = _CONTROL_CHARS_RE.sub("", path) + sanitized = _SAFE_PATH_RE.sub("", path) + return f"{drive_prefix}{sanitized}" + + +def _sanitize_prompt_description(description: str) -> str: + description = description.replace("`", "") + description = _CONTROL_CHARS_RE.sub(" ", description) + description = " ".join(description.split()) + return description + + +_SKILL_NAME_RE = re.compile(r"^[A-Za-z0-9._-]+$") + + +def _sanitize_skill_display_name(name: str) -> str: + if _SKILL_NAME_RE.fullmatch(name): + return name + return "" + + +def _default_sandbox_skill_path(name: str) -> str: + return f"{SANDBOX_WORKSPACE_ROOT}/{SANDBOX_SKILLS_ROOT}/{name}/SKILL.md" + + +def _build_skill_read_command_example(path: str) -> str: + if path == "//SKILL.md": + return f"cat {path}" + if _is_windows_prompt_path(path): + command = "type" + path_arg = f'"{os.path.normpath(path)}"' + else: + command = "cat" + path_arg = shlex.quote(path) + return f"{command} {path_arg}" + + +def build_skills_prompt(skills: list[SkillInfo]) -> str: + """Build the skills section of the system prompt. + + Generates a markdown-formatted skill inventory for the LLM. Only + ``name`` and ``description`` are shown upfront; the LLM must read + the full ``SKILL.md`` before execution (progressive disclosure). + """ + skills_lines: list[str] = [] + example_path = "" + for skill in skills: + display_name = _sanitize_skill_display_name(skill.name) + + description = skill.description or "No description" + if skill.source_type == "sandbox_only": + description = _sanitize_prompt_description(description) + if not description: + description = "Read SKILL.md for details." + + if skill.source_type == "sandbox_only": + # Prefer the actual path from sandbox cache if available + rendered_path = _sanitize_prompt_path_for_prompt(skill.path) + if not rendered_path: + rendered_path = _default_sandbox_skill_path(skill.name) + else: + rendered_path = _sanitize_prompt_path_for_prompt(skill.path) + if not rendered_path: + rendered_path = "//SKILL.md" + + entry = f"- **{display_name}**: {description}\n File: `{rendered_path}`" + if skill.input_schema: + entry += f"\n Input Schema: {json.dumps(skill.input_schema, ensure_ascii=False)}" + if skill.output_schema: + entry += f"\n Output Schema: {json.dumps(skill.output_schema, ensure_ascii=False)}" + skills_lines.append(entry) + if not example_path: + example_path = rendered_path + skills_block = "\n".join(skills_lines) + # Sanitize example_path — it may originate from sandbox cache (untrusted) + if example_path == "//SKILL.md": + example_path = "//SKILL.md" + else: + example_path = _sanitize_prompt_path_for_prompt(example_path) + example_path = example_path or "//SKILL.md" + example_command = _build_skill_read_command_example(example_path) + + return ( + "## Skills\n\n" + "You have specialized skills — reusable instruction bundles stored " + "in `SKILL.md` files. Each skill has a **name** and a **description** " + "that tells you what it does and when to use it.\n\n" + "### Available skills\n\n" + f"{skills_block}\n\n" + "### Skill rules\n\n" + "1. **Discovery** — The list above is the complete skill inventory " + "for this session. Full instructions are in the referenced " + "`SKILL.md` file.\n" + "2. **When to trigger** — Use a skill if the user names it " + "explicitly, or if the task clearly matches the skill's description. " + "*Never silently skip a matching skill* — either use it or briefly " + "explain why you chose not to.\n" + "3. **Mandatory grounding** — Before executing any skill you MUST " + "first read its `SKILL.md` by running a shell command compatible " + "with the current runtime shell and using the **absolute path** " + f"shown above (e.g. `{example_command}`). " + "Never rely on memory or assumptions about a skill's content.\n" + "4. **Progressive disclosure** — Load only what is directly " + "referenced from `SKILL.md`:\n" + " - If `scripts/` exist, prefer running or patching them over " + "rewriting code from scratch.\n" + " - If `assets/` or templates exist, reuse them.\n" + " - Do NOT bulk-load every file in the skill directory.\n" + "5. **Coordination** — When multiple skills apply, pick the minimal " + "set needed. Announce which skill(s) you are using and why " + "(one short line). Prefer `astrbot_*` tools when running skill " + "scripts.\n" + "6. **Context hygiene** — Avoid deep reference chasing; open only " + "files that are directly linked from `SKILL.md`.\n" + "7. **Failure handling** — If a skill cannot be applied, state the " + "issue clearly and continue with the best alternative.\n" + ) diff --git a/astrbot/_internal/skills/to_tool.py b/astrbot/_internal/skills/to_tool.py new file mode 100644 index 000000000..0b2c6a717 --- /dev/null +++ b/astrbot/_internal/skills/to_tool.py @@ -0,0 +1,79 @@ +"""Skill to Tool converter - converts skills with input_schema to FunctionTool.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import Any + +from pydantic.dataclasses import dataclass + +from astrbot.core.agent.tool import FunctionTool + +from .manager import SkillInfo + + +@dataclass +class SkillToToolConverter: + """Converter that transforms skills with input_schema into FunctionTool instances. + + This enables skills to be used as callable tools in the agent's function calling + system, providing a structured way to invoke skills with parameters. + """ + + def can_convert(self, skill: SkillInfo) -> bool: + """Check if a skill can be converted to a FunctionTool. + + Args: + skill: The skill to check + + Returns: + True if the skill has an input_schema, False otherwise + """ + return skill.input_schema is not None + + def convert(self, skill: SkillInfo) -> FunctionTool: + """Convert a skill to a FunctionTool. + + Args: + skill: The skill to convert. Must have an input_schema. + + Returns: + A FunctionTool with name=f"skill_{skill.name}", description, parameters, + and a handler that executes the skill. + + Raises: + ValueError: If the skill does not have an input_schema + """ + if not self.can_convert(skill): + raise ValueError( + f"Skill '{skill.name}' cannot be converted to FunctionTool: " + "no input_schema defined" + ) + + tool_name = f"skill_{skill.name}" + tool_description = skill.description or f"Skill: {skill.name}" + + # Create a handler that executes the skill + # The actual execution is delegated to the skill execution system + async def skill_handler( + context: Any = None, + **kwargs: Any, + ) -> str | None: + # This is a placeholder handler that integrates with the skill system + # The actual skill execution would happen here through the skill manager + # For now, we return a message indicating the skill should be executed + # through the standard skill execution flow + return f"Skill '{skill.name}' execution requested with params: {kwargs}" + + handler: Callable[..., Awaitable[str | None]] = skill_handler + + return FunctionTool( + name=tool_name, + description=tool_description, + parameters=skill.input_schema, # type: ignore[arg-type] + handler=handler, + handler_module_path="astrbot._internal.skills.to_tool", + active=True, + is_background_task=False, + source="skill", + ) diff --git a/astrbot/_internal/tools/__init__.py b/astrbot/_internal/tools/__init__.py new file mode 100644 index 000000000..d34d22f55 --- /dev/null +++ b/astrbot/_internal/tools/__init__.py @@ -0,0 +1,9 @@ +"""Internal tools module for AstrBot. + +This module provides internal tool abstractions and registry functionality +for the AstrBot framework. +""" + +from astrbot._internal.tools.base import FunctionTool, ToolExecResult + +__all__ = ["FunctionTool", "ToolExecResult"] diff --git a/astrbot/_internal/tools/base.py b/astrbot/_internal/tools/base.py new file mode 100644 index 000000000..e69bf7cbc --- /dev/null +++ b/astrbot/_internal/tools/base.py @@ -0,0 +1,442 @@ +"""Internal base tool definitions for AstrBot. + +This module provides the core tool abstractions used throughout AstrBot, +including tool schemas, callable function tools, and tool sets for +managing multiple tools. +""" + +from __future__ import annotations + +import copy +from collections.abc import AsyncGenerator, Awaitable, Callable +from typing import Any, Generic + +import jsonschema +import mcp +from pydantic import Field, model_validator +from pydantic.dataclasses import dataclass + +from astrbot.core.agent.run_context import ContextWrapper, TContext +from astrbot.core.message.message_event_result import MessageEventResult + +ParametersType = dict[str, Any] +ToolExecResult = str | mcp.types.CallToolResult + + +@dataclass +class ToolSchema: + """A class representing the schema of a tool for function calling. + + ToolSchema defines the interface contract for a tool, including its + name, description, and parameter specification in JSON Schema format. + The parameters are validated against the JSON Schema Draft 2020-12. + """ + + name: str + """The name of the tool.""" + + description: str + """The description of the tool.""" + + parameters: ParametersType + """The parameters of the tool, in JSON Schema format.""" + + @model_validator(mode="after") + def validate_parameters(self) -> ToolSchema: + """Validate that parameters conform to JSON Schema Draft 2020-12.""" + jsonschema.validate( + self.parameters, jsonschema.Draft202012Validator.META_SCHEMA + ) + return self + + +@dataclass +class FunctionTool(ToolSchema, Generic[TContext]): + """A callable tool, for function calling. + + FunctionTool represents an executable tool that can be called by the LLM. + It extends ToolSchema with a handler callable that performs the actual + tool execution. The handler should be an async function. + + Type Parameters: + TContext: The context type for tool execution. + + Attributes: + handler: The async callable that implements the tool's functionality. + handler_module_path: Module path of the handler function for serialization. + active: Whether the tool is active and should be used. + is_background_task: Whether this tool runs as a background task. + source: Origin of the tool ('plugin', 'internal', or 'mcp'). + """ + + handler: ( + Callable[..., Awaitable[str | None] | AsyncGenerator[MessageEventResult, None]] + | None + ) = None + """A callable that implements the tool's functionality. It should be an async function.""" + + handler_module_path: str | None = None + """ + The module path of the handler function. This is empty when the origin is mcp. + This field must be retained, as the handler will be wrapped in functools.partial during initialization, + causing the handler's __module__ to be functools + """ + active: bool = True + """ + Whether the tool is active. This field is a special field for AstrBot. + You can ignore it when integrating with other frameworks. + """ + is_background_task: bool = False + """ + Declare this tool as a background task. Background tasks return immediately + with a task identifier while the real work continues asynchronously. + """ + source: str = "plugin" + """ + Origin of this tool: 'plugin' (from star plugins), 'internal' (AstrBot built-in), + or 'mcp' (from MCP servers). Used by WebUI for display grouping. + """ + + def __repr__(self) -> str: + return f"FuncTool(name={self.name}, parameters={self.parameters}, description={self.description})" + + async def call( + self, context: ContextWrapper[TContext], **kwargs: Any + ) -> ToolExecResult: + """Run the tool with the given arguments. The handler field has priority. + + Raises: + NotImplementedError: If no handler is set. + """ + raise NotImplementedError( + "FunctionTool.call() must be implemented by subclasses or set a handler." + ) + + +@dataclass +class ToolSet: + """A set of function tools that can be used in function calling. + + This class provides methods to add, remove, and retrieve tools, as well as + convert the tools to different API formats (OpenAI, Anthropic, Google GenAI). + + Tools can be added with add_tool(), which handles duplicate names by + preferring active tools. The normalize() method sorts tools by name + for deterministic serialization. + """ + + tools: list[FunctionTool] = Field(default_factory=list) + + def empty(self) -> bool: + """Check if the tool set is empty.""" + return len(self.tools) == 0 + + def add_tool(self, tool: FunctionTool) -> None: + """Add a tool to the set. + + If a tool with the same name already exists: + - Prefer the one that is active (active=True) + - If both have the same active state, use the new one (overwrite) + """ + for i, existing_tool in enumerate(self.tools): + if existing_tool.name == tool.name: + # Use getattr with a default of True for compatibility with tools + # that may not define an `active` attribute (e.g., mocks). + existing_active = bool(getattr(existing_tool, "active", True)) + new_active = bool(getattr(tool, "active", True)) + # Overwrite if new tool is active, or if existing tool is not active + if new_active or not existing_active: + self.tools[i] = tool + return + self.tools.append(tool) + + def remove_tool(self, name: str) -> None: + """Remove a tool by its name.""" + self.tools = [tool for tool in self.tools if tool.name != name] + + def normalize(self) -> None: + """Sort tools by name for deterministic serialization. + + This ensures the serialized tool schema sent to the LLM is + identical across requests regardless of registration/injection + order, enabling LLM provider prefix cache hits. + """ + self.tools.sort(key=lambda t: t.name) + + def get_tool(self, name: str) -> FunctionTool | None: + """Get a tool by its name.""" + for tool in self.tools: + if tool.name == name: + return tool + return None + + def get_light_tool_set(self) -> ToolSet: + """Return a light tool set with only name/description.""" + light_tools = [] + for tool in self.tools: + if hasattr(tool, "active") and not tool.active: + continue + light_params = { + "type": "object", + "properties": {}, + } + light_tools.append( + FunctionTool( + name=tool.name, + description=tool.description, + parameters=light_params, + handler=None, + ) + ) + return ToolSet(light_tools) + + def get_param_only_tool_set(self) -> ToolSet: + """Return a tool set with name/parameters only (no description).""" + param_tools = [] + for tool in self.tools: + if hasattr(tool, "active") and not tool.active: + continue + params = ( + copy.deepcopy(tool.parameters) + if tool.parameters + else {"type": "object", "properties": {}} + ) + param_tools.append( + FunctionTool( + name=tool.name, + description="", + parameters=params, + handler=None, + ) + ) + return ToolSet(param_tools) + + def add_func( + self, + name: str, + func_args: list, + desc: str, + handler: Callable[..., Awaitable[Any]], + ) -> None: + """Add a function tool to the set. + + .. deprecated:: 4.0.0 + Use add_tool() instead. + """ + params = { + "type": "object", # hard-coded here + "properties": {}, + } + for param in func_args: + params["properties"][param["name"]] = { + "type": param["type"], + "description": param["description"], + } + _func = FunctionTool( + name=name, + parameters=params, + description=desc, + handler=handler, + ) + self.add_tool(_func) + + def remove_func(self, name: str) -> None: + """Remove a function tool by its name. + + .. deprecated:: 4.0.0 + Use remove_tool() instead. + """ + self.remove_tool(name) + + def get_func(self, name: str) -> FunctionTool | None: + """Get all function tools. + + .. deprecated:: 4.0.0 + Use get_tool() instead. + """ + return self.get_tool(name) + + @property + def func_list(self) -> list[FunctionTool]: + """Get the list of function tools.""" + return self.tools + + def openai_schema(self, omit_empty_parameter_field: bool = False) -> list[dict]: + """Convert tools to OpenAI API function calling schema format.""" + result = [] + for tool in self.tools: + func_def: dict[str, Any] = { + "type": "function", + "function": {"name": tool.name}, + } + if tool.description: + func_def["function"]["description"] = tool.description + + if tool.parameters is not None: + if ( + tool.parameters and tool.parameters.get("properties") + ) or not omit_empty_parameter_field: + func_def["function"]["parameters"] = tool.parameters # type: ignore[index] + + result.append(func_def) + return result + + def anthropic_schema(self) -> list[dict]: + """Convert tools to Anthropic API format.""" + result = [] + for tool in self.tools: + input_schema = {"type": "object"} + if tool.parameters: + input_schema["properties"] = tool.parameters.get("properties", {}) + input_schema["required"] = tool.parameters.get("required", []) + tool_def = {"name": tool.name, "input_schema": input_schema} + if tool.description: + tool_def["description"] = tool.description + result.append(tool_def) + return result + + def google_schema(self) -> dict: + """Convert tools to Google GenAI API format.""" + + def convert_schema(schema: dict) -> dict: + """Convert schema to Gemini API format.""" + supported_types = { + "string", + "number", + "integer", + "boolean", + "array", + "object", + "null", + } + supported_formats = { + "string": {"enum", "date-time"}, + "integer": {"int32", "int64"}, + "number": {"float", "double"}, + } + + if "anyOf" in schema: + return {"anyOf": [convert_schema(s) for s in schema["anyOf"]]} + + result = {} + + # Avoid side effects by not modifying the original schema + origin_type = schema.get("type") + target_type = origin_type + + # Compatibility fix: Gemini API expects 'type' to be a string (enum), + # but standard JSON Schema (MCP) allows lists (e.g. ["string", "null"]). + # We fallback to the first non-null type. + if isinstance(origin_type, list): + target_type = next((t for t in origin_type if t != "null"), "string") + + if target_type in supported_types: + result["type"] = target_type + if "format" in schema and schema["format"] in supported_formats.get( + result["type"], + set(), + ): + result["format"] = schema["format"] + else: + result["type"] = "null" + + support_fields = { + "title", + "description", + "enum", + "minimum", + "maximum", + "maxItems", + "minItems", + "nullable", + "required", + } + result.update({k: schema[k] for k in support_fields if k in schema}) + + if "properties" in schema: + properties = {} + for key, value in schema["properties"].items(): + prop_value = convert_schema(value) + if "default" in prop_value: + del prop_value["default"] + # see #5217 + if "additionalProperties" in prop_value: + del prop_value["additionalProperties"] + properties[key] = prop_value + + if properties: + result["properties"] = properties + + if target_type == "array": + items_schema = schema.get("items") + if isinstance(items_schema, dict): + result["items"] = convert_schema(items_schema) + else: + # Gemini requires array schemas to include an `items` schema. + # JSON Schema allows omitting it, so fall back to a permissive + # string item schema instead of emitting an invalid declaration. + result["items"] = {"type": "string"} + + return result + + tools = [] + for tool in self.tools: + d: dict[str, Any] = {"name": tool.name} + if tool.description: + d["description"] = tool.description + if tool.parameters: + d["parameters"] = convert_schema(tool.parameters) + tools.append(d) + + declarations = {} + if tools: + declarations["function_declarations"] = tools + return declarations + + def get_func_desc_openai_style(self, omit_empty_parameter_field: bool = False): + """Get OpenAI style function descriptions. + + .. deprecated:: 4.0.0 + Use openai_schema() instead. + """ + return self.openai_schema(omit_empty_parameter_field) + + def get_func_desc_anthropic_style(self): + """Get Anthropic style function descriptions. + + .. deprecated:: 4.0.0 + Use anthropic_schema() instead. + """ + return self.anthropic_schema() + + def get_func_desc_google_genai_style(self): + """Get Google GenAI style function descriptions. + + .. deprecated:: 4.0.0 + Use google_schema() instead. + """ + return self.google_schema() + + def names(self) -> list[str]: + """Get a list of all tool names.""" + return [tool.name for tool in self.tools] + + def merge(self, other: ToolSet) -> None: + """Merge another ToolSet into this one.""" + for tool in other.tools: + self.add_tool(tool) + + def __len__(self) -> int: + return len(self.tools) + + def __bool__(self) -> bool: + return len(self.tools) > 0 + + def __iter__(self): + return iter(self.tools) + + def __repr__(self) -> str: + return f"ToolSet(tools={self.tools})" + + def __str__(self) -> str: + return f"ToolSet(tools={self.tools})" diff --git a/astrbot/_internal/tools/builtin/__init__.py b/astrbot/_internal/tools/builtin/__init__.py new file mode 100644 index 000000000..b05c44cfc --- /dev/null +++ b/astrbot/_internal/tools/builtin/__init__.py @@ -0,0 +1,52 @@ +"""Built-in tools for AstrBot. + +This module provides access to all AstrBot built-in tools that are part of +the internal tool providers system. +""" + +from astrbot._internal.tools.base import FunctionTool +from astrbot._internal.tools.builtin.cron import ( + CREATE_CRON_JOB_TOOL, + DELETE_CRON_JOB_TOOL, + LIST_CRON_JOBS_TOOL, +) +from astrbot._internal.tools.builtin.cron import ( + get_all_tools as cron_get_all_tools, +) +from astrbot._internal.tools.builtin.kb_query import ( + KNOWLEDGE_BASE_QUERY_TOOL, +) +from astrbot._internal.tools.builtin.kb_query import ( + get_all_tools as kb_query_get_all_tools, +) +from astrbot._internal.tools.builtin.send_message import ( + SEND_MESSAGE_TO_USER_TOOL, +) +from astrbot._internal.tools.builtin.send_message import ( + get_all_tools as send_message_get_all_tools, +) + + +def get_all_tools() -> list[FunctionTool]: + """Return all built-in tools for registration. + + This aggregates tools from all built-in tool modules: + - cron tools (create/delete/list future tasks) + - knowledge base query tool + - send message tool + """ + tools: list[FunctionTool] = [] + tools.extend(cron_get_all_tools()) + tools.extend(kb_query_get_all_tools()) + tools.extend(send_message_get_all_tools()) + return tools + + +__all__ = [ + "CREATE_CRON_JOB_TOOL", + "DELETE_CRON_JOB_TOOL", + "KNOWLEDGE_BASE_QUERY_TOOL", + "LIST_CRON_JOBS_TOOL", + "SEND_MESSAGE_TO_USER_TOOL", + "get_all_tools", +] diff --git a/astrbot/_internal/tools/builtin/cron.py b/astrbot/_internal/tools/builtin/cron.py new file mode 100644 index 000000000..ec53958f2 --- /dev/null +++ b/astrbot/_internal/tools/builtin/cron.py @@ -0,0 +1,201 @@ +from datetime import datetime +from typing import Any + +from pydantic import Field +from pydantic.dataclasses import dataclass + +from astrbot._internal.tools.base import FunctionTool, ToolExecResult +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.astr_agent_context import AstrAgentContext + + +def _extract_job_session(job: Any) -> str | None: + payload = getattr(job, "payload", None) + if not isinstance(payload, dict): + return None + session = payload.get("session") + return str(session) if session is not None else None + + +@dataclass +class CreateActiveCronTool(FunctionTool[AstrAgentContext]): + name: str = "create_future_task" + description: str = ( + "Create a future task for your future. Supports recurring cron expressions or one-time run_at datetime. " + "Use this when you or the user want scheduled follow-up or proactive actions." + ) + parameters: dict = Field( + default_factory=lambda: { + "type": "object", + "properties": { + "cron_expression": { + "type": "string", + "description": "Cron expression defining recurring schedule (e.g., '0 8 * * *' or '0 23 * * mon-fri'). Prefer named weekdays like 'mon-fri' or 'sat,sun' instead of numeric day-of-week ranges such as '1-5' to avoid ambiguity across cron implementations.", + }, + "run_at": { + "type": "string", + "description": "ISO datetime for one-time execution, e.g., 2026-02-02T08:00:00+08:00. Use with run_once=true.", + }, + "note": { + "type": "string", + "description": "Detailed instructions for your future agent to execute when it wakes.", + }, + "name": { + "type": "string", + "description": "Optional label to recognize this future task.", + }, + "run_once": { + "type": "boolean", + "description": "If true, the task will run only once and then be deleted. Use run_at to specify the time.", + }, + }, + "required": ["note"], + } + ) + + async def call( + self, context: ContextWrapper[AstrAgentContext], **kwargs + ) -> ToolExecResult: + cron_mgr = context.context.context.cron_manager + if cron_mgr is None: + return "error: cron manager is not available." + + cron_expression = kwargs.get("cron_expression") + run_at = kwargs.get("run_at") + run_once = bool(kwargs.get("run_once", False)) + note = str(kwargs.get("note", "")).strip() + name = str(kwargs.get("name") or "").strip() or "active_agent_task" + + if not note: + return "error: note is required." + if run_once and not run_at: + return "error: run_at is required when run_once=true." + if (not run_once) and not cron_expression: + return "error: cron_expression is required when run_once=false." + if run_once and cron_expression: + cron_expression = None + run_at_dt = None + if run_at: + try: + run_at_dt = datetime.fromisoformat(str(run_at)) + except Exception: + return "error: run_at must be ISO datetime, e.g., 2026-02-02T08:00:00+08:00" + + payload = { + "session": context.context.event.unified_msg_origin, + "sender_id": context.context.event.get_sender_id(), + "note": note, + "origin": "tool", + } + + job = await cron_mgr.add_active_job( + name=name, + cron_expression=str(cron_expression) if cron_expression else None, + payload=payload, + description=note, + run_once=run_once, + run_at=run_at_dt, + ) + next_run = job.next_run_time or run_at_dt + suffix = ( + f"one-time at {next_run}" + if run_once + else f"expression '{cron_expression}' (next {next_run})" + ) + return f"Scheduled future task {job.job_id} ({job.name}) {suffix}." + + +@dataclass +class DeleteCronJobTool(FunctionTool[AstrAgentContext]): + name: str = "delete_future_task" + description: str = "Delete a future task (cron job) by its job_id." + parameters: dict = Field( + default_factory=lambda: { + "type": "object", + "properties": { + "job_id": { + "type": "string", + "description": "The job_id returned when the job was created.", + } + }, + "required": ["job_id"], + } + ) + + async def call( + self, context: ContextWrapper[AstrAgentContext], **kwargs + ) -> ToolExecResult: + cron_mgr = context.context.context.cron_manager + if cron_mgr is None: + return "error: cron manager is not available." + current_umo = context.context.event.unified_msg_origin + job_id = kwargs.get("job_id") + if not job_id: + return "error: job_id is required." + job = await cron_mgr.db.get_cron_job(str(job_id)) + if not job: + return f"error: cron job {job_id} not found." + if _extract_job_session(job) != current_umo: + return "error: you can only delete future tasks in the current umo." + await cron_mgr.delete_job(str(job_id)) + return f"Deleted cron job {job_id}." + + +@dataclass +class ListCronJobsTool(FunctionTool[AstrAgentContext]): + name: str = "list_future_tasks" + description: str = "List existing future tasks (cron jobs) for inspection." + parameters: dict = Field( + default_factory=lambda: { + "type": "object", + "properties": { + "job_type": { + "type": "string", + "description": "Optional filter: basic or active_agent.", + } + }, + } + ) + + async def call( + self, context: ContextWrapper[AstrAgentContext], **kwargs + ) -> ToolExecResult: + cron_mgr = context.context.context.cron_manager + if cron_mgr is None: + return "error: cron manager is not available." + current_umo = context.context.event.unified_msg_origin + job_type = kwargs.get("job_type") + jobs = [ + job + for job in await cron_mgr.list_jobs(job_type) + if _extract_job_session(job) == current_umo + ] + if not jobs: + return "No cron jobs found." + lines = [] + for j in jobs: + lines.append( + f"{j.job_id} | {j.name} | {j.job_type} | run_once={getattr(j, 'run_once', False)} | enabled={j.enabled} | next={j.next_run_time}" + ) + return "\n".join(lines) + + +CREATE_CRON_JOB_TOOL = CreateActiveCronTool() +DELETE_CRON_JOB_TOOL = DeleteCronJobTool() +LIST_CRON_JOBS_TOOL = ListCronJobsTool() + + +def get_all_tools() -> list[FunctionTool]: + """Return all cron-related tools for registration.""" + return [CREATE_CRON_JOB_TOOL, DELETE_CRON_JOB_TOOL, LIST_CRON_JOBS_TOOL] + + +__all__ = [ + "CREATE_CRON_JOB_TOOL", + "DELETE_CRON_JOB_TOOL", + "LIST_CRON_JOBS_TOOL", + "CreateActiveCronTool", + "DeleteCronJobTool", + "ListCronJobsTool", + "get_all_tools", +] diff --git a/astrbot/_internal/tools/builtin/kb_query.py b/astrbot/_internal/tools/builtin/kb_query.py new file mode 100644 index 000000000..f6392bafd --- /dev/null +++ b/astrbot/_internal/tools/builtin/kb_query.py @@ -0,0 +1,139 @@ +"""Knowledge base query tool and retrieval logic. + +Extracted from ``astr_main_agent_resources.py`` to its own module. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pydantic import Field +from pydantic.dataclasses import dataclass + +from astrbot._internal.tools.base import FunctionTool, ToolExecResult +from astrbot.api import logger, sp +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.astr_agent_context import AstrAgentContext + +if TYPE_CHECKING: + from astrbot.core.star.context import Context + + +@dataclass +class KnowledgeBaseQueryTool(FunctionTool[AstrAgentContext]): + name: str = "astr_kb_search" + description: str = ( + "Query the knowledge base for facts or relevant context. " + "Use this tool when the user's question requires factual information, " + "definitions, background knowledge, or previously indexed content. " + "Only send short keywords or a concise question as the query." + ) + parameters: dict = Field( + default_factory=lambda: { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "A concise keyword query for the knowledge base.", + }, + }, + "required": ["query"], + } + ) + + async def call( + self, context: ContextWrapper[AstrAgentContext], **kwargs + ) -> ToolExecResult: + query = kwargs.get("query", "") + if not query: + return "error: Query parameter is empty." + result = await retrieve_knowledge_base( + query=kwargs.get("query", ""), + umo=context.context.event.unified_msg_origin, + context=context.context.context, + ) + if not result: + return "No relevant knowledge found." + return result + + +async def retrieve_knowledge_base( + query: str, + umo: str, + context: Context, +) -> str | None: + """Inject knowledge base context into the provider request + + Args: + query: The search query string + umo: Unique message object (session ID) + context: Star context + """ + kb_mgr = context.kb_manager + config = context.get_config(umo=umo) + + # 1. Prefer session-level config + session_config = await sp.session_get(umo, "kb_config", default={}) + + if session_config and "kb_ids" in session_config: + kb_ids = session_config.get("kb_ids", []) + + if not kb_ids: + logger.info(f"[知识库] 会话 {umo} 已被配置为不使用知识库") + return + + top_k = session_config.get("top_k", 5) + + kb_names = [] + invalid_kb_ids = [] + for kb_id in kb_ids: + kb_helper = await kb_mgr.get_kb(kb_id) + if kb_helper: + kb_names.append(kb_helper.kb.kb_name) + else: + logger.warning(f"[知识库] 知识库不存在或未加载: {kb_id}") + invalid_kb_ids.append(kb_id) + + if invalid_kb_ids: + logger.warning( + f"[知识库] 会话 {umo} 配置的以下知识库无效: {invalid_kb_ids}", + ) + + if not kb_names: + return + + logger.debug(f"[知识库] 使用会话级配置,知识库数量: {len(kb_names)}") + else: + kb_names = config.get("kb_names", []) + top_k = config.get("kb_final_top_k", 5) + logger.debug(f"[知识库] 使用全局配置,知识库数量: {len(kb_names)}") + + top_k_fusion = config.get("kb_fusion_top_k", 20) + + if not kb_names: + return + + logger.debug(f"[知识库] 开始检索知识库,数量: {len(kb_names)}, top_k={top_k}") + kb_context = await kb_mgr.retrieve( + query=query, + kb_names=kb_names, + top_k_fusion=top_k_fusion, + top_m_final=top_k, + ) + + if not kb_context: + return + + formatted = kb_context.get("context_text", "") + if formatted: + results = kb_context.get("results", []) + logger.debug(f"[知识库] 为会话 {umo} 注入了 {len(results)} 条相关知识块") + return formatted + + +KNOWLEDGE_BASE_QUERY_TOOL = KnowledgeBaseQueryTool() + + +def get_all_tools() -> list[FunctionTool]: + """Return all knowledge-base tools for registration.""" + return [KNOWLEDGE_BASE_QUERY_TOOL] diff --git a/astrbot/_internal/tools/builtin/send_message.py b/astrbot/_internal/tools/builtin/send_message.py new file mode 100644 index 000000000..9125c6daa --- /dev/null +++ b/astrbot/_internal/tools/builtin/send_message.py @@ -0,0 +1,226 @@ +"""SendMessageToUserTool — proactive message delivery to users. + +Extracted from ``astr_main_agent_resources.py`` to its own module. +""" + +from __future__ import annotations + +import json +import os +import uuid +from typing import Any, TypedDict, cast + +import anyio +from pydantic import Field +from pydantic.dataclasses import dataclass + +import astrbot.core.message.components as Comp +from astrbot._internal.tools.base import FunctionTool, ToolExecResult +from astrbot.api import logger +from astrbot.core.agent.run_context import ContextWrapper +from astrbot.core.astr_agent_context import AstrAgentContext +from astrbot.core.computer.computer_client import get_booter +from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.platform.message_session import MessageSession +from astrbot.core.utils.astrbot_path import get_astrbot_temp_path + + +class MessageComponent(TypedDict, total=False): + """Type-safe message component structure.""" + + type: str + text: str + path: str + url: str + mention_user_id: str + + +@dataclass +class SendMessageToUserTool(FunctionTool[AstrAgentContext]): + name: str = "send_message_to_user" + description: str = "Directly send message to the user. Only use this tool when you need to proactively message the user. Otherwise you can directly output the reply in the conversation." + + parameters: dict = Field( + default_factory=lambda: { + "type": "object", + "properties": { + "messages": { + "type": "array", + "description": "An ordered list of message components to send. `mention_user` type can be used to mention the user.", + "items": { + "type": "object", + "additionalProperties": {"type": "string"}, + }, + }, + }, + "required": ["messages"], + } + ) + + async def _resolve_path_from_sandbox( + self, context: ContextWrapper[AstrAgentContext], path: str + ) -> tuple[str, bool]: + """ + If the path exists locally, return it directly. + Otherwise, check if it exists in the sandbox and download it. + + bool: indicates whether the file was downloaded from sandbox. + """ + if await anyio.Path(path).exists(): + return path, False + + # Try to check if the file exists in the sandbox + try: + sb = await get_booter( + context.context.context, + context.context.event.unified_msg_origin, + ) + # Use shell to check if the file exists in sandbox + import shlex + + result = await sb.shell.exec( + f"test -f {shlex.quote(path)} && echo '_&exists_'" + ) + if "_&exists_" in json.dumps(result): + # Download the file from sandbox + name = anyio.Path(path).name + local_path = os.path.join( + get_astrbot_temp_path(), f"sandbox_{uuid.uuid4().hex[:4]}_{name}" + ) + await sb.download_file(path, local_path) + logger.info(f"Downloaded file from sandbox: {path} -> {local_path}") + return local_path, True + except Exception as e: + logger.warning(f"Failed to check/download file from sandbox: {e}") + + # Return the original path (will likely fail later, but that's expected) + return path, False + + async def call( + self, context: ContextWrapper[AstrAgentContext], **kwargs: Any + ) -> ToolExecResult: + session: str | MessageSession = ( + kwargs.get("session") or context.context.event.unified_msg_origin + ) + messages: list[dict[str, Any]] | None = kwargs.get("messages") + + if not isinstance(messages, list) or not messages: + return "error: messages parameter is empty or invalid." + + components: list[Comp.BaseMessageComponent] = [] + + for idx, msg in enumerate(messages): + if not isinstance(msg, dict): + return f"error: messages[{idx}] should be an object." + + msg_dict: dict[str, Any] = cast(dict[str, Any], msg) + + if "type" not in msg_dict: + return f"error: messages[{idx}].type is required." + msg_type = str(msg_dict["type"]).lower() + + _file_from_sandbox = False + + try: + if msg_type == "plain": + text = str(msg_dict.get("text", "")).strip() + if not text: + return f"error: messages[{idx}].text is required for plain component." + components.append(Comp.Plain(text=text)) + elif msg_type == "image": + path = msg_dict.get("path") + url = msg_dict.get("url") + if path: + ( + local_path, + _file_from_sandbox, + ) = await self._resolve_path_from_sandbox(context, path) + components.append(Comp.Image.fromFileSystem(path=local_path)) + elif url: + components.append(Comp.Image.fromURL(url=url)) + else: + return f"error: messages[{idx}] must include path or url for image component." + elif msg_type == "record": + path = msg_dict.get("path") + url = msg_dict.get("url") + if path: + ( + local_path, + _file_from_sandbox, + ) = await self._resolve_path_from_sandbox(context, path) + components.append(Comp.Record.fromFileSystem(path=local_path)) + elif url: + components.append(Comp.Record.fromURL(url=url)) + else: + return f"error: messages[{idx}] must include path or url for record component." + elif msg_type == "video": + path = msg_dict.get("path") + url = msg_dict.get("url") + if path: + ( + local_path, + _file_from_sandbox, + ) = await self._resolve_path_from_sandbox(context, path) + components.append(Comp.Video.fromFileSystem(path=local_path)) + elif url: + components.append(Comp.Video.fromURL(url=url)) + else: + return f"error: messages[{idx}] must include path or url for video component." + elif msg_type == "file": + path = msg_dict.get("path") + url = msg_dict.get("url") + name = ( + msg_dict.get("text") + or (os.path.basename(path) if path else "") + or (os.path.basename(url) if url else "") + or "file" + ) + if path: + ( + local_path, + _file_from_sandbox, + ) = await self._resolve_path_from_sandbox(context, path) + components.append(Comp.File(name=name, file=local_path)) + elif url: + components.append(Comp.File(name=name, url=url)) + else: + return f"error: messages[{idx}] must include path or url for file component." + elif msg_type == "mention_user": + mention_user_id = msg_dict.get("mention_user_id") + if not mention_user_id: + return f"error: messages[{idx}].mention_user_id is required for mention_user component." + components.append( + Comp.At( + qq=mention_user_id, + ), + ) + else: + return ( + f"error: unsupported message type '{msg_type}' at index {idx}." + ) + except Exception as exc: + return f"error: failed to build messages[{idx}] component: {exc}" + + try: + target_session = ( + MessageSession.from_str(session) + if isinstance(session, str) + else session + ) + except Exception as e: + return f"error: invalid session: {e}" + + await context.context.context.send_message( + target_session, + MessageChain(chain=components), + ) + + return f"Message sent to session {target_session}" + + +SEND_MESSAGE_TO_USER_TOOL = SendMessageToUserTool() + + +def get_all_tools() -> list[FunctionTool]: + """Return all send-message tools for registration.""" + return [SEND_MESSAGE_TO_USER_TOOL] diff --git a/astrbot/_internal/tools/providers/__init__.py b/astrbot/_internal/tools/providers/__init__.py new file mode 100644 index 000000000..ca91fa43d --- /dev/null +++ b/astrbot/_internal/tools/providers/__init__.py @@ -0,0 +1,25 @@ +"""Tool providers for AstrBot. + +This module provides different tool providers that supply tools +through a unified interface: + +- InternalToolProvider: Provides built-in AstrBot tools (cron, kb_query, send_message) +- PluginToolProvider: Provides tools registered by star plugins +- ComputerToolProvider: Provides computer-use tools (shell, Python, file ops, etc.) +""" + +from astrbot._internal.tools.providers.computer import ( + ComputerToolProvider, +) +from astrbot._internal.tools.providers.internal import ( + InternalToolProvider, +) +from astrbot._internal.tools.providers.plugin import ( + PluginToolProvider, +) + +__all__ = [ + "ComputerToolProvider", + "InternalToolProvider", + "PluginToolProvider", +] diff --git a/astrbot/_internal/tools/providers/computer.py b/astrbot/_internal/tools/providers/computer.py new file mode 100644 index 000000000..2ddcb9d8f --- /dev/null +++ b/astrbot/_internal/tools/providers/computer.py @@ -0,0 +1,45 @@ +"""Computer tools provider for AstrBot. + +This provider wraps the ComputerToolProvider from the computer module +to ensure computer tools are available through the unified provider interface. +""" + +from typing import TYPE_CHECKING + +from astrbot._internal.tools.base import FunctionTool + +if TYPE_CHECKING: + pass + + +class ComputerToolProvider: + """Provider for computer-use tools (local/sandbox). + + This class wraps the existing ``ComputerToolProvider`` from + ``astrbot.core.computer.computer_tool_provider`` to integrate + computer tools into the unified provider interface. + + The computer tools include shell execution, Python code execution, + file operations, browser automation, and skill management tools. + """ + + @staticmethod + def get_all_tools() -> list[FunctionTool]: + """Return all computer-use tools across all runtimes. + + Delegates to ``ComputerToolProvider.get_all_tools()`` which + collects tools from all runtimes (local, sandbox, browser, neo). + + Creates **fresh instances** separate from the runtime caches so + that setting ``active=False`` on them does not affect runtime + behaviour. These registration-only instances let the WebUI display + and assign tools without injecting them into actual LLM requests. + + Returns: + list[FunctionTool]: A list of all computer FunctionTool instances. + """ + from astrbot.core.computer.computer_tool_provider import ( + ComputerToolProvider as CoreComputerToolProvider, + ) + + return CoreComputerToolProvider.get_all_tools() diff --git a/astrbot/_internal/tools/providers/internal.py b/astrbot/_internal/tools/providers/internal.py new file mode 100644 index 000000000..b086c0754 --- /dev/null +++ b/astrbot/_internal/tools/providers/internal.py @@ -0,0 +1,73 @@ +"""Internal tools provider for AstrBot. + +This provider wraps the logic for loading built-in internal tools from +the provider modules: cron_tools, kb_query, and send_message. +""" + +from typing import TYPE_CHECKING + +from astrbot import logger +from astrbot._internal.tools.base import FunctionTool + +if TYPE_CHECKING: + pass + + +# Provider modules that supply internal tools +_INTERNAL_PROVIDER_MODULES: list[str] = [ + "astrbot.core.tools.cron_tools", + "astrbot.core.tools.kb_query", + "astrbot.core.tools.send_message", +] + + +class InternalToolProvider: + """Provider for AstrBot built-in internal tools. + + This class wraps the logic previously found in + ``FunctionToolManager._INTERNAL_TOOL_PROVIDERS`` and provides + a unified interface for loading tools from the internal provider + modules. + + Each provider module is expected to expose a ``get_all_tools()`` + function that returns a list of ``FunctionTool`` instances. + + Tools are marked with ``source='internal'`` so the WebUI can + distinguish them from plugin and MCP tools. + """ + + @staticmethod + def get_all_tools() -> list[FunctionTool]: + """Return all internal tools from all provider modules. + + Iterates through the provider modules and collects tools + from each module's ``get_all_tools()`` function. + + Returns: + list[FunctionTool]: A list of all internal FunctionTool instances. + """ + all_tools: list[FunctionTool] = [] + existing_names: set[str] = set() + + for module_path in _INTERNAL_PROVIDER_MODULES: + try: + import importlib + + mod = importlib.import_module(module_path) + provider_tools = mod.get_all_tools() + except Exception as e: + logger.warning( + "Failed to load internal tool provider %s: %s", + module_path, + e, + ) + continue + + for tool in provider_tools: + tool.source = "internal" + if tool.name not in existing_names: + all_tools.append(tool) + existing_names.add(tool.name) + logger.debug("Loaded internal tool: %s", tool.name) + + return all_tools diff --git a/astrbot/_internal/tools/providers/plugin.py b/astrbot/_internal/tools/providers/plugin.py new file mode 100644 index 000000000..eb1b35515 --- /dev/null +++ b/astrbot/_internal/tools/providers/plugin.py @@ -0,0 +1,73 @@ +"""Plugin tools provider for AstrBot. + +This provider handles loading tools from star plugins. Plugin tools +are discovered through the star plugin system and made available +through the unified provider interface. +""" + +from astrbot._internal.tools.base import FunctionTool + + +class PluginToolProvider: + """Provider for tools from star plugins. + + This class handles loading tools that are registered by star plugins. + Plugin tools are discovered through the plugin system and integrated + into the tool registry. + + Note: Plugin tools are typically registered dynamically through the + plugin context (``Context.register_llm_tool()``) and are managed + by the ``FunctionToolManager`` in the provider module. + + This provider class serves as an integration point for the plugin + tool system with the unified internal tools architecture. + """ + + @staticmethod + def get_all_tools() -> list[FunctionTool]: + """Return all plugin-registered tools. + + This method retrieves tools that have been registered by plugins + through the ``FunctionToolManager``. It accesses the global + ``llm_tools`` instance from ``astrbot.core.provider.register``. + + Returns: + list[FunctionTool]: A list of all plugin FunctionTool instances. + """ + from astrbot.core.provider.register import llm_tools + from astrbot.core.star.star import star_map + + # Get all tools from the FunctionToolManager that are from plugins + plugin_tools: list[FunctionTool] = [] + existing_names: set[str] = set() + + for tool in llm_tools.func_list: + # Only include tools that are marked as 'plugin' source + # and belong to an activated plugin + if tool.source == "plugin": + if tool.name not in existing_names: + if tool.handler_module_path: + star_meta = star_map.get(tool.handler_module_path) + if star_meta and star_meta.activated: + plugin_tools.append(tool) + existing_names.add(tool.name) + else: + # Tools without handler_module_path are treated as plugin tools + plugin_tools.append(tool) + existing_names.add(tool.name) + + return plugin_tools + + @staticmethod + def get_tool(name: str) -> FunctionTool | None: + """Get a specific plugin tool by name. + + Args: + name: The name of the tool to retrieve. + + Returns: + FunctionTool | None: The tool if found, None otherwise. + """ + from astrbot.core.provider.register import llm_tools + + return llm_tools.get_func(name) diff --git a/astrbot/_internal/tools/registry.py b/astrbot/_internal/tools/registry.py new file mode 100644 index 000000000..bc923f36d --- /dev/null +++ b/astrbot/_internal/tools/registry.py @@ -0,0 +1,821 @@ +"""FunctionTool registry and manager for AstrBot. + +This module provides the FunctionToolManager class that serves as the central +registry for all function tools (built-in, plugin, and MCP). +""" + +from __future__ import annotations + +import asyncio +import json +import os +import threading +import urllib.parse +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, ClassVar + +import aiofiles +import anyio + +from astrbot import logger +from astrbot.core import sp +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + +from .base import FunctionTool, ToolSet + +# Deferred imports to avoid circular dependency +_mcp_client_module = None +_mcp_tool_module = None + + +def _get_mcp_client(): + global _mcp_client_module + if _mcp_client_module is None: + from astrbot._internal.mcp import client as module + _mcp_client_module = module + return _mcp_client_module + + +def _get_mcp_tool(): + global _mcp_tool_module + if _mcp_tool_module is None: + from astrbot._internal.mcp import tool as module + _mcp_tool_module = module + return _mcp_tool_module + + +DEFAULT_MCP_CONFIG = {"mcpServers": {}} + +DEFAULT_MCP_INIT_TIMEOUT_SECONDS = 180.0 +DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS = 180.0 +MCP_INIT_TIMEOUT_ENV = "ASTRBOT_MCP_INIT_TIMEOUT" +ENABLE_MCP_TIMEOUT_ENV = "ASTRBOT_MCP_ENABLE_TIMEOUT" +MAX_MCP_TIMEOUT_SECONDS = 300.0 + + +class MCPInitError(Exception): + """Base exception for MCP initialization failures.""" + + +class MCPInitTimeoutError(asyncio.TimeoutError, MCPInitError): + """Raised when MCP client initialization exceeds the configured timeout.""" + + +class MCPAllServicesFailedError(MCPInitError): + """Raised when all configured MCP services fail to initialize.""" + + +class MCPShutdownTimeoutError(asyncio.TimeoutError): + """Raised when MCP shutdown exceeds the configured timeout.""" + + def __init__(self, names: list[str], timeout: float) -> None: + self.names = names + self.timeout = timeout + message = f"MCP 服务关闭超时({timeout:g} 秒):{', '.join(names)}" + super().__init__(message) + + +@dataclass +class MCPInitSummary: + total: int + success: int + failed: list[str] + + +@dataclass +class _MCPServerRuntime: + name: str + client: Any # MCPClient + shutdown_event: asyncio.Event + lifecycle_task: asyncio.Task[None] + + +class _MCPClientDictView(Mapping[str, Any]): + """Read-only view of MCP clients derived from runtime state.""" + + def __init__(self, runtime: dict[str, _MCPServerRuntime]) -> None: + self._runtime = runtime + + def __getitem__(self, key: str) -> Any: + return self._runtime[key].client + + def __iter__(self): + return iter(self._runtime) + + def __len__(self) -> int: + return len(self._runtime) + + +def _resolve_timeout( + timeout: float | str | None = None, + *, + env_name: str = MCP_INIT_TIMEOUT_ENV, + default: float = DEFAULT_MCP_INIT_TIMEOUT_SECONDS, +) -> float: + """Resolve timeout with precedence: explicit argument > env value > default.""" + source = f"环境变量 {env_name}" + if timeout is None: + timeout = os.getenv(env_name, str(default)) + else: + source = "显式参数 timeout" + + try: + timeout_value = float(timeout) + except (TypeError, ValueError): + logger.warning( + f"超时配置({source})={timeout!r} 无效,使用默认值 {default:g} 秒｡" + ) + return default + + if timeout_value <= 0: + logger.warning( + f"超时配置({source})={timeout_value:g} 必须大于 0,使用默认值 {default:g} 秒｡" + ) + return default + + if timeout_value > MAX_MCP_TIMEOUT_SECONDS: + logger.warning( + f"超时配置({source})={timeout_value:g} 过大,已限制为最大值 " + f"{MAX_MCP_TIMEOUT_SECONDS:g} 秒,以避免长时间等待｡" + ) + return MAX_MCP_TIMEOUT_SECONDS + + return timeout_value + + +SUPPORTED_TYPES = [ + "string", + "number", + "object", + "array", + "boolean", +] + +PY_TO_JSON_TYPE = { + "int": "number", + "float": "number", + "bool": "boolean", + "str": "string", + "dict": "object", + "list": "array", + "tuple": "array", + "set": "array", +} + + +class FunctionToolManager: + """Central registry for all function tools in AstrBot. + + This class manages: + - Built-in tools (cron, KB query, send message, computer) + - Plugin tools + - MCP tools (from external MCP servers) + + Tools are stored in func_list and can be queried by name. + MCP servers are tracked separately in _mcp_server_runtime. + """ + + def __init__(self) -> None: + self.func_list: list[FunctionTool] = [] + self._mcp_server_runtime: dict[str, _MCPServerRuntime] = {} + self._mcp_server_runtime_view = MappingProxyType(self._mcp_server_runtime) + self._mcp_client_dict_view = _MCPClientDictView(self._mcp_server_runtime) + self._timeout_mismatch_warned = False + self._timeout_warn_lock = threading.Lock() + self._runtime_lock = asyncio.Lock() + self._mcp_starting: set[str] = set() + self._init_timeout_default = _resolve_timeout( + timeout=None, + env_name=MCP_INIT_TIMEOUT_ENV, + default=DEFAULT_MCP_INIT_TIMEOUT_SECONDS, + ) + self._enable_timeout_default = _resolve_timeout( + timeout=None, + env_name=ENABLE_MCP_TIMEOUT_ENV, + default=DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS, + ) + self._warn_on_timeout_mismatch( + self._init_timeout_default, + self._enable_timeout_default, + ) + + @property + def mcp_client_dict(self) -> Mapping[str, Any]: + """Read-only view of MCP clients.""" + return self._mcp_client_dict_view + + @property + def mcp_server_runtime_view(self) -> Mapping[str, _MCPServerRuntime]: + """Read-only view of MCP runtime metadata.""" + return self._mcp_server_runtime_view + + @property + def mcp_server_runtime(self) -> Mapping[str, _MCPServerRuntime]: + """Backward-compatible read-only view (deprecated).""" + return self._mcp_server_runtime_view + + def empty(self) -> bool: + return len(self.func_list) == 0 + + def spec_to_func( + self, + name: str, + func_args: list[dict], + desc: str, + handler: Callable[..., Awaitable[Any] | AsyncGenerator[Any]], + ) -> FunctionTool: + params = { + "type": "object", + "properties": {}, + } + for param in func_args: + p = param.copy() + p.pop("name", None) + params["properties"][param["name"]] = p + return FunctionTool( + name=name, + parameters=params, + description=desc, + handler=handler, + ) + + def add_func( + self, + name: str, + func_args: list, + desc: str, + handler: Callable[..., Awaitable[Any] | AsyncGenerator[Any]], + ) -> None: + """Add a function tool.""" + self.remove_func(name) + self.func_list.append( + self.spec_to_func( + name=name, + func_args=func_args, + desc=desc, + handler=handler, + ), + ) + logger.info(f"添加函数调用工具: {name}") + + def remove_func(self, name: str) -> None: + """Remove a function tool by name.""" + for i, f in enumerate(self.func_list): + if f.name == name: + self.func_list.pop(i) + break + + def get_func(self, name: str) -> FunctionTool | None: + """Get a function tool by name. + + Prefers active tools, falls back to last registered. + """ + for f in reversed(self.func_list): + if f.name == name and getattr(f, "active", True): + return f + for f in reversed(self.func_list): + if f.name == name: + return f + return None + + def get_full_tool_set(self) -> ToolSet: + """Get all tools as a ToolSet.""" + tool_set = ToolSet() + for tool in self.func_list: + tool_set.add_tool(tool) + return tool_set + + @staticmethod + def _log_safe_mcp_debug_config(cfg: dict) -> None: + """Log sanitized MCP config for debugging.""" + if "command" in cfg: + cmd = cfg["command"] + executable = str(cmd[0] if isinstance(cmd, (list, tuple)) and cmd else cmd) + args_val = cfg.get("args", []) + args_count = ( + len(args_val) + if isinstance(args_val, (list, tuple)) + else (0 if args_val is None else 1) + ) + logger.debug(f" 命令可执行文件: {executable}, 参数数量: {args_count}") + return + + if "url" in cfg: + parsed = urllib.parse.urlparse(str(cfg["url"])) + host = parsed.hostname or "" + scheme = parsed.scheme or "unknown" + try: + port = f":{parsed.port}" if parsed.port else "" + except ValueError: + port = "" + logger.debug(f" 主机: {scheme}://{host}{port}") + + async def init_mcp_clients( + self, raise_on_all_failed: bool = False + ) -> MCPInitSummary: + """Initialize MCP clients from mcp_server.json config.""" + data_dir = get_astrbot_data_path() + + mcp_json_file = os.path.join(data_dir, "mcp_server.json") + mcp_json_path = anyio.Path(mcp_json_file) + if not await mcp_json_path.exists(): + async with aiofiles.open(mcp_json_file, "w", encoding="utf-8") as f: + await f.write( + json.dumps(DEFAULT_MCP_CONFIG, ensure_ascii=False, indent=4) + ) + logger.info(f"未找到 MCP 服务配置文件,已创建默认配置文件 {mcp_json_file}") + return MCPInitSummary(total=0, success=0, failed=[]) + + async with aiofiles.open(mcp_json_file, encoding="utf-8") as f: + mcp_server_json_obj: dict[str, dict] = json.loads(await f.read())[ + "mcpServers" + ] + + init_timeout = self._init_timeout_default + timeout_display = f"{init_timeout:g}" + + active_configs: list[tuple[str, dict, asyncio.Event]] = [] + for name, cfg in mcp_server_json_obj.items(): + if cfg.get("active", True): + shutdown_event = asyncio.Event() + active_configs.append((name, cfg, shutdown_event)) + + if not active_configs: + return MCPInitSummary(total=0, success=0, failed=[]) + + logger.info(f"等待 {len(active_configs)} 个 MCP 服务初始化...") + + init_tasks = [ + asyncio.create_task( + self._start_mcp_server( + name=name, + cfg=cfg, + shutdown_event=shutdown_event, + init_timeout=init_timeout, + ), + name=f"mcp-init:{name}", + ) + for (name, cfg, shutdown_event) in active_configs + ] + results = await asyncio.gather(*init_tasks, return_exceptions=True) + + success_count = 0 + failed_services: list[str] = [] + + for (name, cfg, _), result in zip(active_configs, results, strict=False): + if isinstance(result, Exception): + if isinstance(result, MCPInitTimeoutError): + logger.error( + f"Connected to MCP server {name} timeout ({timeout_display} seconds)" + ) + else: + logger.error(f"Failed to initialize MCP server {name}: {result}") + self._log_safe_mcp_debug_config(cfg) + failed_services.append(name) + async with self._runtime_lock: + self._mcp_server_runtime.pop(name, None) + continue + + success_count += 1 + + if failed_services: + logger.warning( + f"The following MCP services failed to initialize: {', '.join(failed_services)}. " + f"Please check the mcp_server.json file and server availability." + ) + + summary = MCPInitSummary( + total=len(active_configs), success=success_count, failed=failed_services + ) + logger.info( + f"MCP services initialization completed: {summary.success}/{summary.total} successful, {len(summary.failed)} failed." + ) + if summary.total > 0 and summary.success == 0: + msg = "All MCP services failed to initialize, please check the mcp_server.json and server availability." + if raise_on_all_failed: + raise MCPAllServicesFailedError(msg) + logger.error(msg) + return summary + + async def _start_mcp_server( + self, + name: str, + cfg: dict, + *, + shutdown_event: asyncio.Event | None = None, + init_timeout: float, + ) -> None: + """Start an MCP server with timeout.""" + async with self._runtime_lock: + if name in self._mcp_server_runtime or name in self._mcp_starting: + logger.warning( + f"Connected to MCP server {name}, ignoring this startup request (timeout={init_timeout:g})." + ) + self._log_safe_mcp_debug_config(cfg) + return + self._mcp_starting.add(name) + + if shutdown_event is None: + shutdown_event = asyncio.Event() + + mcp_client: Any = None + try: + mcp_client = await asyncio.wait_for( + self._init_mcp_client(name, cfg), + timeout=init_timeout, + ) + except asyncio.TimeoutError as exc: + raise MCPInitTimeoutError( + f"Connected to MCP server {name} timeout ({init_timeout:g} seconds)" + ) from exc + except Exception: + logger.error(f"Failed to initialize MCP client {name}", exc_info=True) + raise + finally: + if mcp_client is None: + async with self._runtime_lock: + self._mcp_starting.discard(name) + + async def lifecycle() -> None: + try: + await shutdown_event.wait() + logger.info(f"Received shutdown signal for MCP client {name}") + except asyncio.CancelledError: + logger.debug(f"MCP client {name} task was cancelled") + raise + finally: + await self._terminate_mcp_client(name) + + lifecycle_task = asyncio.create_task(lifecycle(), name=f"mcp-client:{name}") + async with self._runtime_lock: + self._mcp_server_runtime[name] = _MCPServerRuntime( + name=name, + client=mcp_client, + shutdown_event=shutdown_event, + lifecycle_task=lifecycle_task, + ) + self._mcp_starting.discard(name) + + async def _shutdown_runtimes( + self, + runtimes: list[_MCPServerRuntime], + shutdown_timeout: float, + *, + strict: bool = True, + ) -> list[str]: + """Shutdown runtimes and wait for lifecycle tasks.""" + lifecycle_tasks = [ + runtime.lifecycle_task + for runtime in runtimes + if not runtime.lifecycle_task.done() + ] + if not lifecycle_tasks: + return [] + + for runtime in runtimes: + runtime.shutdown_event.set() + + try: + results = await asyncio.wait_for( + asyncio.gather(*lifecycle_tasks, return_exceptions=True), + timeout=shutdown_timeout, + ) + except asyncio.TimeoutError: + pending_names = [ + runtime.name + for runtime in runtimes + if not runtime.lifecycle_task.done() + ] + for task in lifecycle_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*lifecycle_tasks, return_exceptions=True) + if strict: + raise MCPShutdownTimeoutError(pending_names, shutdown_timeout) + logger.warning( + "MCP server shutdown timeout (%s seconds), the following servers were not fully closed: %s", + f"{shutdown_timeout:g}", + ", ".join(pending_names), + ) + return pending_names + else: + for result in results: + if isinstance(result, asyncio.CancelledError): + logger.debug("MCP lifecycle task was cancelled during shutdown.") + elif isinstance(result, Exception): + logger.error( + "MCP lifecycle task failed during shutdown.", + exc_info=(type(result), result, result.__traceback__), + ) + return [] + + async def _cleanup_mcp_client_safely( + self, mcp_client: Any, name: str + ) -> None: + """Safely cleanup an MCP client.""" + try: + await mcp_client.cleanup() + except Exception as cleanup_exc: + logger.error( + f"Failed to cleanup MCP client resources {name}: {cleanup_exc}" + ) + + async def _init_mcp_client(self, name: str, config: dict) -> Any: + """Initialize a single MCP client.""" + mcp_mod = _get_mcp_client() + MCPClient = mcp_mod.MCPClient + mcp_tool_mod = _get_mcp_tool() + MCPTool = mcp_tool_mod.MCPTool + + mcp_client = MCPClient() + mcp_client.name = name + try: + await mcp_client.connect_to_server(config, name) + tools_res = await mcp_client.list_tools_and_save() + except asyncio.CancelledError: + await self._cleanup_mcp_client_safely(mcp_client, name) + raise + except Exception: + await self._cleanup_mcp_client_safely(mcp_client, name) + raise + logger.debug(f"MCP server {name} list tools response: {tools_res}") + tool_names = [tool.name for tool in tools_res.tools] + + # Remove old MCP tools for this server + self.func_list = [ + f + for f in self.func_list + if not (isinstance(f, MCPTool) and f.mcp_server_name == name) + ] + + # Add new MCP tools + for tool in mcp_client.tools: + func_tool = MCPTool( + mcp_tool=tool, + mcp_client=mcp_client, + mcp_server_name=name, + ) + self.func_list.append(func_tool) + + logger.info(f"Connected to MCP server {name}, Tools: {tool_names}") + return mcp_client + + async def _terminate_mcp_client(self, name: str) -> None: + """Terminate and cleanup an MCP client.""" + async with self._runtime_lock: + runtime = self._mcp_server_runtime.get(name) + if runtime: + client = runtime.client + await self._cleanup_mcp_client_safely(client, name) + self.func_list = [ + f + for f in self.func_list + if not (isinstance(f, _get_mcp_tool().MCPTool) and f.mcp_server_name == name) + ] + async with self._runtime_lock: + self._mcp_server_runtime.pop(name, None) + self._mcp_starting.discard(name) + logger.info(f"Disconnected from MCP server {name}") + return + + self.func_list = [ + f + for f in self.func_list + if not (isinstance(f, _get_mcp_tool().MCPTool) and f.mcp_server_name == name) + ] + async with self._runtime_lock: + self._mcp_starting.discard(name) + + async def test_mcp_server_connection(self, config: dict) -> list[str]: + """Test connection to an MCP server.""" + mcp_mod = _get_mcp_client() + MCPClient = mcp_mod.MCPClient + _prepare_config = mcp_mod._prepare_config + _quick_test_mcp_connection = mcp_mod._quick_test_mcp_connection + + if "url" in config: + cfg = _prepare_config(config.copy()) + success, error_msg = await _quick_test_mcp_connection(cfg) + if not success: + raise Exception(error_msg) + + mcp_client = MCPClient() + try: + logger.debug(f"testing MCP server connection with config: {config}") + await mcp_client.connect_to_server(config, "test") + tools_res = await mcp_client.list_tools_and_save() + tool_names = [tool.name for tool in tools_res.tools] + finally: + logger.debug("Cleaning up MCP client after testing connection.") + await mcp_client.cleanup() + return tool_names + + async def enable_mcp_server( + self, + name: str, + config: dict, + shutdown_event: asyncio.Event | None = None, + init_timeout: float | str | None = None, + ) -> None: + """Enable and initialize an MCP server.""" + if init_timeout is None: + timeout_value = self._enable_timeout_default + else: + timeout_value = _resolve_timeout( + timeout=init_timeout, + env_name=ENABLE_MCP_TIMEOUT_ENV, + default=self._enable_timeout_default, + ) + await self._start_mcp_server( + name=name, + cfg=config, + shutdown_event=shutdown_event, + init_timeout=timeout_value, + ) + + async def disable_mcp_server( + self, + name: str | None = None, + shutdown_timeout: float = 10, + ) -> None: + """Disable an MCP server by name, or all if name is None.""" + if name: + async with self._runtime_lock: + runtime = self._mcp_server_runtime.get(name) + if runtime is None: + return + await self._shutdown_runtimes([runtime], shutdown_timeout, strict=True) + else: + async with self._runtime_lock: + runtimes = list(self._mcp_server_runtime.values()) + await self._shutdown_runtimes(runtimes, shutdown_timeout, strict=False) + + def _warn_on_timeout_mismatch( + self, + init_timeout: float, + enable_timeout: float, + ) -> None: + if init_timeout == enable_timeout: + return + with self._timeout_warn_lock: + if self._timeout_mismatch_warned: + return + logger.info( + "检测到 MCP 初始化超时与动态启用超时配置不同:" + "初始化使用 %s 秒,动态启用使用 %s 秒｡如需一致,请设置相同值｡", + f"{init_timeout:g}", + f"{enable_timeout:g}", + ) + self._timeout_mismatch_warned = True + + def get_func_desc_openai_style(self, omit_empty_parameter_field=False) -> list: + """Get OpenAI-style function descriptions for active tools.""" + tools = [f for f in self.func_list if f.active] + toolset = ToolSet(tools) + return toolset.openai_schema( + omit_empty_parameter_field=omit_empty_parameter_field, + ) + + def get_func_desc_anthropic_style(self) -> list: + """Get Anthropic-style function descriptions for active tools.""" + tools = [f for f in self.func_list if f.active] + toolset = ToolSet(tools) + return toolset.anthropic_schema() + + def get_func_desc_google_genai_style(self) -> dict: + """Get Google GenAI-style function descriptions for active tools.""" + tools = [f for f in self.func_list if f.active] + toolset = ToolSet(tools) + return toolset.google_schema() + + def deactivate_llm_tool(self, name: str) -> bool: + """Deactivate a registered function tool.""" + func_tool = self.get_func(name) + if func_tool is not None: + func_tool.active = False + + inactivated_llm_tools: list = sp.get( + "inactivated_llm_tools", + [], + scope="global", + scope_id="global", + ) + if name not in inactivated_llm_tools: + inactivated_llm_tools.append(name) + sp.put( + "inactivated_llm_tools", + inactivated_llm_tools, + scope="global", + scope_id="global", + ) + + return True + return False + + def activate_llm_tool(self, name: str, star_map: dict) -> bool: + """Activate a registered function tool.""" + func_tool = self.get_func(name) + if func_tool is not None: + if func_tool.handler_module_path in star_map: + if not star_map[func_tool.handler_module_path].activated: + raise ValueError( + f"此函数调用工具所属的插件 {star_map[func_tool.handler_module_path].name} 已被禁用,请先在管理面板启用再激活此工具｡" + ) + + func_tool.active = True + + inactivated_llm_tools: list = sp.get( + "inactivated_llm_tools", + [], + scope="global", + scope_id="global", + ) + if name in inactivated_llm_tools: + inactivated_llm_tools.remove(name) + sp.put( + "inactivated_llm_tools", + inactivated_llm_tools, + scope="global", + scope_id="global", + ) + + return True + return False + + @property + def mcp_config_path(self) -> str: + data_dir = get_astrbot_data_path() + return os.path.join(data_dir, "mcp_server.json") + + def load_mcp_config(self) -> dict: + """Load MCP configuration from file.""" + if not os.path.exists(self.mcp_config_path): + os.makedirs(os.path.dirname(self.mcp_config_path), exist_ok=True) + with open(self.mcp_config_path, "w", encoding="utf-8") as f: + json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4) + return DEFAULT_MCP_CONFIG.copy() + + try: + with open(self.mcp_config_path, encoding="utf-8") as f: + return json.load(f) + except Exception as e: + logger.error(f"加载 MCP 配置失败: {e}") + return DEFAULT_MCP_CONFIG.copy() + + def save_mcp_config(self, config: dict) -> bool: + """Save MCP configuration to file.""" + try: + with open(self.mcp_config_path, "w", encoding="utf-8") as f: + json.dump(config, f, ensure_ascii=False, indent=4) + return True + except Exception as e: + logger.error(f"保存 MCP 配置失败: {e}") + return False + + # Module paths for built-in tool providers + _INTERNAL_TOOL_PROVIDERS: ClassVar[list[str]] = [ + "astrbot.core.tools.cron_tools", + "astrbot.core.tools.kb_query", + "astrbot.core.tools.send_message", + "astrbot.core.computer.computer_tool_provider", + ] + + def register_internal_tools(self) -> None: + """Register AstrBot built-in tools from all internal providers. + + Each provider module should expose a get_all_tools() function. + """ + import importlib + + existing_names = {t.name for t in self.func_list} + + for module_path in self._INTERNAL_TOOL_PROVIDERS: + try: + mod = importlib.import_module(module_path) + provider_tools = mod.get_all_tools() + except Exception as e: + logger.warning( + "Failed to load internal tool provider %s: %s", + module_path, + e, + ) + continue + + for tool in provider_tools: + tool.source = "internal" + if tool.name not in existing_names: + self.func_list.append(tool) + existing_names.add(tool.name) + logger.info("Registered internal tool: %s", tool.name) + + def __str__(self) -> str: + return str(self.func_list) + + def __repr__(self) -> str: + return str(self.func_list) + + +# Alias for backward compatibility +FuncCall = FunctionToolManager diff --git a/astrbot/api/__init__.py b/astrbot/api/__init__.py index 5d15dedc2..b1f696401 100644 --- a/astrbot/api/__init__.py +++ b/astrbot/api/__init__.py @@ -1,4 +1,37 @@ +""" +AstrBot Public API. + +This package exposes the public interface for extending and integrating with +AstrBot. All exports from this module are guaranteed to be stable across +minor version updates. + +Modules: + tools: Tool registration and management API + mcp: Model Context Protocol server and tool API + skills: Skill management and conversion API +""" + from astrbot import logger + +# MCP API +from astrbot.api.mcp import ( + MCPClient, + MCPTool, + get_mcp_servers, + register_mcp_server, + unregister_mcp_server, +) + +# Skills API +from astrbot.api.skills import ( + SkillInfo, + SkillManager, + get_skill_manager, + skill_to_tool, +) + +# Tools API +from astrbot.api.tools import ToolRegistry, get_registry, tool from astrbot.core import html_renderer, sp from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor @@ -10,10 +43,22 @@ __all__ = [ "AstrBotConfig", "BaseFunctionToolExecutor", "FunctionTool", + "MCPClient", + "MCPTool", + "SkillInfo", + "SkillManager", + "ToolRegistry", "ToolSet", "agent", + "get_mcp_servers", + "get_registry", + "get_skill_manager", "html_renderer", "llm_tool", "logger", + "register_mcp_server", + "skill_to_tool", "sp", + "tool", + "unregister_mcp_server", ] diff --git a/astrbot/api/mcp.py b/astrbot/api/mcp.py new file mode 100644 index 000000000..c3152b82a --- /dev/null +++ b/astrbot/api/mcp.py @@ -0,0 +1,93 @@ +""" +MCP (Model Context Protocol) Public API for AstrBot. + +Example: + from astrbot.api.mcp import get_mcp_servers, register_mcp_server + + # List connected servers + servers = get_mcp_servers() + + # Register stdio MCP server + await register_mcp_server( + name="weather", + command="uv", + args=["tool", "run", "weather-mcp"], + ) + + # Register SSE server + await register_mcp_server( + name="fileserver", + url="http://localhost:8080/sse", + transport="sse", + ) +""" + +from __future__ import annotations + +from typing import Any + +from astrbot.core.agent.mcp_client import MCPClient, MCPTool +from astrbot.core.provider.func_tool_manager import FunctionToolManager + +__all__ = [ + "MCPClient", + "MCPTool", + "get_mcp_servers", + "register_mcp_server", + "unregister_mcp_server", +] + + +def get_mcp_servers() -> dict[str, MCPClient]: + """Get all connected MCP servers.""" + from astrbot.core.provider.register import llm_tools as func_tool_manager + + manager: FunctionToolManager = func_tool_manager + return dict(manager.mcp_client_dict) + + +async def register_mcp_server( + name: str, + command: str | None = None, + args: list[str] | None = None, + url: str | None = None, + transport: str | None = None, + **kwargs: Any, +) -> None: + """Register and connect to an MCP server. + + Args: + name: Unique name for this server + command: Command to run (for stdio transport) + args: Command arguments + url: URL (for SSE/Streamable HTTP transports) + transport: "sse", "streamable_http", or None for stdio + + Example - Stdio: + await register_mcp_server(name="weather", command="uv", + args=["tool", "run", "weather-mcp"]) + """ + from astrbot.core.provider.register import llm_tools as func_tool_manager + + manager: FunctionToolManager = func_tool_manager + + config: dict[str, Any] = {} + if command is not None: + config["command"] = command + if args is not None: + config["args"] = args + if url is not None: + config["url"] = url + if transport is not None: + config["transport"] = transport + config.update(kwargs) + + await manager.enable_mcp_server(name=name, config=config) + + +async def unregister_mcp_server(name: str) -> None: + """Disconnect and remove an MCP server.""" + from astrbot.core.provider.register import llm_tools as func_tool_manager + + manager: FunctionToolManager = func_tool_manager + await manager.disable_mcp_server(name=name) diff --git a/astrbot/api/skills.py b/astrbot/api/skills.py new file mode 100644 index 000000000..25828074f --- /dev/null +++ b/astrbot/api/skills.py @@ -0,0 +1,52 @@ +""" +Skills Public API for AstrBot. + +Two skill types: +1. Prompt-based: SKILL.md files injected into system prompt +2. Tool-based: Skills with input_schema converted to FunctionTool + +Example: + from astrbot.api.skills import get_skill_manager, skill_to_tool + + # List skills + mgr = get_skill_manager() + skills = mgr.list_skills() + + # Convert tool-based skill to FunctionTool + tool_skills = [s for s in skills if s.input_schema] + if tool_skills: + func_tool = skill_to_tool(tool_skills[0]) +""" + +from __future__ import annotations + +from astrbot.core.agent.tool import FunctionTool +from astrbot.core.skills.skill_manager import SkillInfo, SkillManager + +__all__ = ["SkillInfo", "SkillManager", "get_skill_manager", "skill_to_tool"] + + +def get_skill_manager() -> SkillManager: + """Get the global SkillManager instance.""" + return SkillManager() + + +def skill_to_tool(skill: SkillInfo) -> FunctionTool | None: + """Convert a tool-based skill (with input_schema) to a FunctionTool. + + Args: + skill: A SkillInfo instance with an input_schema + + Returns: + A FunctionTool, or None if the skill has no input_schema + """ + if not skill.input_schema: + return None + + return FunctionTool( + name=f"skill_{skill.name}", + description=skill.description or f"Skill: {skill.name}", + parameters=skill.input_schema, + handler=None, + source="skill", + ) diff --git a/astrbot/api/tools.py b/astrbot/api/tools.py new file mode 100644 index 000000000..1c2492be7 --- /dev/null +++ b/astrbot/api/tools.py @@ -0,0 +1,110 @@ +""" +Tools Public API for AstrBot. + +Example: + from astrbot.api.tools import tool, get_registry + + @tool(name="weather", description="Get weather", parameters={...}) + async def get_weather(city: str) -> str: + return f"Weather in {city} is sunny" + + registry = get_registry() + tools = registry.list_tools() +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from functools import wraps +from typing import Any + +from astrbot.core.agent.tool import FunctionTool + +__all__ = ["FunctionTool", "ToolRegistry", "get_registry", "tool"] + + +class ToolRegistry: + """Simple wrapper around FunctionToolManager for tool registration.""" + + _instance: ToolRegistry | None = None + + def __init__(self) -> None: + from astrbot.core.provider.register import llm_tools as func_tool_manager + + self._manager = func_tool_manager + + @classmethod + def get_instance(cls) -> ToolRegistry: + """Get the singleton ToolRegistry instance.""" + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def register(self, tool: FunctionTool) -> None: + """Register a FunctionTool.""" + self._manager.func_list.append(tool) + + def unregister(self, name: str) -> bool: + """Unregister a tool by name. Returns True if found and removed.""" + for i, f in enumerate(self._manager.func_list): + if f.name == name: + self._manager.func_list.pop(i) + return True + return False + + def list_tools(self) -> list[FunctionTool]: + """List all registered tools.""" + return self._manager.func_list.copy() + + def get_tool(self, name: str) -> FunctionTool | None: + """Get a tool by name.""" + return self._manager.get_func(name) + + +def get_registry() -> ToolRegistry: + """Get the global ToolRegistry instance.""" + return ToolRegistry.get_instance() + + +def tool( + name: str, + description: str, + parameters: dict[str, Any] | None = None, +) -> Callable[ + [Callable[..., Awaitable[str | None]]], Callable[..., Awaitable[str | None]] +]: + """Decorator to register an async function as a tool. + + Args: + name: Tool name (used by LLM to invoke it) + description: What the tool does + parameters: JSON Schema for parameters (optional) + + Example: + @tool(name="weather", description="Get weather for a city", parameters={...}) + async def get_weather(city: str) -> str: + return f"The weather in {city} is sunny" + """ + if parameters is None: + parameters = {"type": "object", "properties": {}} + + def decorator( + func: Callable[..., Awaitable[str | None]], + ) -> Callable[..., Awaitable[str | None]]: + func_tool = FunctionTool( + name=name, + description=description, + parameters=parameters, + handler=func, + handler_module_path=getattr(func, "__module__", ""), + source="api", + ) + get_registry().register(func_tool) + + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> str | None: + return await func(*args, **kwargs) + + return wrapper + + return decorator