mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
99 lines
2.6 KiB
Python
99 lines
2.6 KiB
Python
"""
|
|
MCP (Model Context Protocol) Public API for AstrBot.
|
|
|
|
This module provides a simple, stable interface for MCP server management,
|
|
delegating to the _internal package.
|
|
|
|
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
|
|
|
|
# Import from _internal package (the canonical source)
|
|
# TODO: fix path - should be protocols.mcp.client
|
|
from astrbot._internal.protocols._mcp.client import McpClient as MCPClient
|
|
from astrbot._internal.protocols._mcp.tool import MCPTool
|
|
|
|
__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 = 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 = 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 = func_tool_manager
|
|
await manager.disable_mcp_server(name=name)
|