mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
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)
This commit is contained in:
63
astrbot/_internal/mcp/__init__.py
Normal file
63
astrbot/_internal/mcp/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
411
astrbot/_internal/mcp/client.py
Normal file
411
astrbot/_internal/mcp/client.py
Normal file
@@ -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
|
||||
55
astrbot/_internal/mcp/config.py
Normal file
55
astrbot/_internal/mcp/config.py
Normal file
@@ -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
|
||||
40
astrbot/_internal/mcp/tool.py
Normal file
40
astrbot/_internal/mcp/tool.py
Normal file
@@ -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),
|
||||
)
|
||||
37
astrbot/_internal/skills/__init__.py
Normal file
37
astrbot/_internal/skills/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
79
astrbot/_internal/skills/loader.py
Normal file
79
astrbot/_internal/skills/loader.py
Normal file
@@ -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
|
||||
617
astrbot/_internal/skills/manager.py
Normal file
617
astrbot/_internal/skills/manager.py
Normal file
@@ -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 "<invalid_skill_name>"
|
||||
|
||||
|
||||
def _build_skill_read_command_example(path: str) -> str:
|
||||
if path == "<skills_root>/<skill_name>/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 = "<skills_root>/<skill_name>/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 == "<skills_root>/<skill_name>/SKILL.md":
|
||||
example_path = "<skills_root>/<skill_name>/SKILL.md"
|
||||
else:
|
||||
example_path = _sanitize_prompt_path_for_prompt(example_path)
|
||||
example_path = example_path or "<skills_root>/<skill_name>/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
|
||||
81
astrbot/_internal/skills/parser.py
Normal file
81
astrbot/_internal/skills/parser.py
Normal file
@@ -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,
|
||||
}
|
||||
158
astrbot/_internal/skills/prompt_builder.py
Normal file
158
astrbot/_internal/skills/prompt_builder.py
Normal file
@@ -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 "<invalid_skill_name>"
|
||||
|
||||
|
||||
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 == "<skills_root>/<skill_name>/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 = "<skills_root>/<skill_name>/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 == "<skills_root>/<skill_name>/SKILL.md":
|
||||
example_path = "<skills_root>/<skill_name>/SKILL.md"
|
||||
else:
|
||||
example_path = _sanitize_prompt_path_for_prompt(example_path)
|
||||
example_path = example_path or "<skills_root>/<skill_name>/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"
|
||||
)
|
||||
79
astrbot/_internal/skills/to_tool.py
Normal file
79
astrbot/_internal/skills/to_tool.py
Normal file
@@ -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",
|
||||
)
|
||||
9
astrbot/_internal/tools/__init__.py
Normal file
9
astrbot/_internal/tools/__init__.py
Normal file
@@ -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"]
|
||||
442
astrbot/_internal/tools/base.py
Normal file
442
astrbot/_internal/tools/base.py
Normal file
@@ -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})"
|
||||
52
astrbot/_internal/tools/builtin/__init__.py
Normal file
52
astrbot/_internal/tools/builtin/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
201
astrbot/_internal/tools/builtin/cron.py
Normal file
201
astrbot/_internal/tools/builtin/cron.py
Normal file
@@ -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",
|
||||
]
|
||||
139
astrbot/_internal/tools/builtin/kb_query.py
Normal file
139
astrbot/_internal/tools/builtin/kb_query.py
Normal file
@@ -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]
|
||||
226
astrbot/_internal/tools/builtin/send_message.py
Normal file
226
astrbot/_internal/tools/builtin/send_message.py
Normal file
@@ -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]
|
||||
25
astrbot/_internal/tools/providers/__init__.py
Normal file
25
astrbot/_internal/tools/providers/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
45
astrbot/_internal/tools/providers/computer.py
Normal file
45
astrbot/_internal/tools/providers/computer.py
Normal file
@@ -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()
|
||||
73
astrbot/_internal/tools/providers/internal.py
Normal file
73
astrbot/_internal/tools/providers/internal.py
Normal file
@@ -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
|
||||
73
astrbot/_internal/tools/providers/plugin.py
Normal file
73
astrbot/_internal/tools/providers/plugin.py
Normal file
@@ -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)
|
||||
821
astrbot/_internal/tools/registry.py
Normal file
821
astrbot/_internal/tools/registry.py
Normal file
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
93
astrbot/api/mcp.py
Normal file
93
astrbot/api/mcp.py
Normal file
@@ -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)
|
||||
52
astrbot/api/skills.py
Normal file
52
astrbot/api/skills.py
Normal file
@@ -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",
|
||||
)
|
||||
110
astrbot/api/tools.py
Normal file
110
astrbot/api/tools.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user