AstrBot 3.5.26__tar.gz → 4.0.0b2__tar.gz

{astrbot-3.5.26 → astrbot-4.0.0b2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AstrBot
-Version: 3.5.26
+Version: 4.0.0b2
 Summary: 易上手的多平台 LLM 聊天机器人及开发框架
 License-File: LICENSE
 Requires-Python: >=3.10
@@ -17,6 +17,7 @@ Requires-Dist: colorlog>=6.9.0
 Requires-Dist: cryptography>=44.0.3
 Requires-Dist: dashscope>=1.23.2
 Requires-Dist: defusedxml>=0.7.1
+Requires-Dist: deprecated>=1.2.18
 Requires-Dist: dingtalk-stream>=0.22.1
 Requires-Dist: docstring-parser>=0.16
 Requires-Dist: faiss-cpu>=1.10.0
@@ -41,6 +42,8 @@ Requires-Dist: quart>=0.20.0
 Requires-Dist: readability-lxml>=0.8.4.1
 Requires-Dist: silk-python>=0.2.6
 Requires-Dist: slack-sdk>=3.35.0
+Requires-Dist: sqlalchemy[asyncio]>=2.0.41
+Requires-Dist: sqlmodel>=0.0.24
 Requires-Dist: telegramify-markdown>=0.5.1
 Requires-Dist: watchfiles>=1.0.5
 Requires-Dist: websockets>=15.0.1

astrbot-4.0.0b2/astrbot/api/__init__.py

@@ -0,0 +1,20 @@
+from astrbot.core.config.astrbot_config import AstrBotConfig
+from astrbot import logger
+from astrbot.core import html_renderer
+from astrbot.core import sp
+from astrbot.core.star.register import register_llm_tool as llm_tool
+from astrbot.core.star.register import register_agent as agent
+from astrbot.core.agent.tool import ToolSet, FunctionTool
+from astrbot.core.agent.tool_executor import BaseFunctionToolExecutor
+
+__all__ = [
+    "AstrBotConfig",
+    "logger",
+    "html_renderer",
+    "llm_tool",
+    "agent",
+    "sp",
+    "ToolSet",
+    "FunctionTool",
+    "BaseFunctionToolExecutor",
+]

{astrbot-3.5.26 → astrbot-4.0.0b2}/astrbot/core/__init__.py

@@ -1,5 +1,4 @@
 import os
-import asyncio
 from .log import LogManager, LogBroker  # noqa
 from astrbot.core.utils.t2i.renderer import HtmlRenderer
 from astrbot.core.utils.shared_preferences import SharedPreferences
@@ -21,7 +20,7 @@ html_renderer = HtmlRenderer(t2i_base_url)
 logger = LogManager.GetLogger(log_name="astrbot")
 db_helper = SQLiteDatabase(DB_PATH)
 # 简单的偏好设置存储, 这里后续应该存储到数据库中, 一些部分可以存储到配置中
-sp = SharedPreferences()
+sp = SharedPreferences(db_helper=db_helper)
 # 文件令牌服务
 file_token_service = FileTokenService()
 pip_installer = PipInstaller(

astrbot-4.0.0b2/astrbot/core/agent/agent.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from .tool import FunctionTool
+from typing import Generic
+from .run_context import TContext
+from .hooks import BaseAgentRunHooks
+
+
+@dataclass
+class Agent(Generic[TContext]):
+    name: str
+    instructions: str | None = None
+    tools: list[str, FunctionTool] | None = None
+    run_hooks: BaseAgentRunHooks[TContext] | None = None

astrbot-4.0.0b2/astrbot/core/agent/handoff.py

@@ -0,0 +1,34 @@
+from typing import Generic
+from .tool import FunctionTool
+from .agent import Agent
+from .run_context import TContext
+
+
+class HandoffTool(FunctionTool, Generic[TContext]):
+    """Handoff tool for delegating tasks to another agent."""
+
+    def __init__(
+        self, agent: Agent[TContext], parameters: dict | None = None, **kwargs
+    ):
+        self.agent = agent
+        super().__init__(
+            name=f"transfer_to_{agent.name}",
+            parameters=parameters or self.default_parameters(),
+            description=agent.instructions or self.default_description(agent.name),
+            **kwargs,
+        )
+
+    def default_parameters(self) -> dict:
+        return {
+            "type": "object",
+            "properties": {
+                "input": {
+                    "type": "string",
+                    "description": "The input to be handed off to another agent. This should be a clear and concise request or task.",
+                },
+            },
+        }
+
+    def default_description(self, agent_name: str | None) -> str:
+        agent_name = agent_name or "another"
+        return f"Delegate tasks to {self.name} agent to handle the request."

astrbot-4.0.0b2/astrbot/core/agent/hooks.py

@@ -0,0 +1,27 @@
+import mcp
+from dataclasses import dataclass
+from .run_context import ContextWrapper, TContext
+from typing import Generic
+from astrbot.core.provider.entities import LLMResponse
+from astrbot.core.agent.tool import FunctionTool
+
+
+@dataclass
+class BaseAgentRunHooks(Generic[TContext]):
+    async def on_agent_begin(self, run_context: ContextWrapper[TContext]): ...
+    async def on_tool_start(
+        self,
+        run_context: ContextWrapper[TContext],
+        tool: FunctionTool,
+        tool_args: dict | None,
+    ): ...
+    async def on_tool_end(
+        self,
+        run_context: ContextWrapper[TContext],
+        tool: FunctionTool,
+        tool_args: dict | None,
+        tool_result: mcp.types.CallToolResult | None,
+    ): ...
+    async def on_agent_done(
+        self, run_context: ContextWrapper[TContext], llm_response: LLMResponse
+    ): ...

astrbot-4.0.0b2/astrbot/core/agent/mcp_client.py

@@ -0,0 +1,208 @@
+import asyncio
+import logging
+from datetime import timedelta
+from typing import Optional
+from contextlib import AsyncExitStack
+from astrbot import logger
+from astrbot.core.utils.log_pipe import LogPipe
+
+try:
+    import mcp
+    from mcp.client.sse import sse_client
+except (ModuleNotFoundError, ImportError):
+    logger.warning("警告: 缺少依赖库 'mcp'，将无法使用 MCP 服务。")
+
+try:
+    from mcp.client.streamable_http import streamablehttp_client
+except (ModuleNotFoundError, ImportError):
+    logger.warning(
+        "警告: 缺少依赖库 'mcp' 或者 mcp 库版本过低，无法使用 Streamable HTTP 连接方式。"
+    )
+
+
+def _prepare_config(config: dict) -> dict:
+    """准备配置，处理嵌套格式"""
+    if "mcpServers" in config and config["mcpServers"]:
+        first_key = next(iter(config["mcpServers"]))
+        config = config["mcpServers"][first_key]
+    config.pop("active", None)
+    return config
+
+
+async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]:
+    """快速测试 MCP 服务器可达性"""
+    import aiohttp
+
+    cfg = _prepare_config(config.copy())
+
+    url = cfg["url"]
+    headers = cfg.get("headers", {})
+    timeout = cfg.get("timeout", 10)
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            if cfg.get("transport") == "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, ""
+                    else:
+                        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, ""
+                    else:
+                        return False, f"HTTP {response.status}: {response.reason}"
+
+    except asyncio.TimeoutError:
+        return False, f"连接超时: {timeout}秒"
+    except Exception as e:
+        return False, f"{e!s}"
+
+
+class MCPClient:
+    def __init__(self):
+        # Initialize session and client objects
+        self.session: Optional[mcp.ClientSession] = None
+        self.exit_stack = AsyncExitStack()
+
+        self.name = None
+        self.active: bool = True
+        self.tools: list[mcp.Tool] = []
+        self.server_errlogs: list[str] = []
+        self.running_event = asyncio.Event()
+
+    async def connect_to_server(self, mcp_server_config: dict, name: str):
+        """连接到 MCP 服务器
+
+        如果 `url` 参数存在：
+            1. 当 transport 指定为 `streamable_http` 时，使用 Streamable HTTP 连接方式。
+            1. 当 transport 指定为 `sse` 时，使用 SSE 连接方式。
+            2. 如果没有指定，默认使用 SSE 的方式连接到 MCP 服务。
+
+        Args:
+            mcp_server_config (dict): Configuration for the MCP server. See https://modelcontextprotocol.io/quickstart/server
+        """
+        cfg = _prepare_config(mcp_server_config.copy())
+
+        def logging_callback(msg: str):
+            # 处理 MCP 服务的错误日志
+            print(f"MCP Server {name} Error: {msg}")
+            self.server_errlogs.append(msg)
+
+        if "url" in cfg:
+            success, error_msg = await _quick_test_mcp_connection(cfg)
+            if not success:
+                raise Exception(error_msg)
+
+            if cfg.get("transport") != "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", 20))
+                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", 20))
+                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:
+            server_params = mcp.StdioServerParameters(
+                **cfg,
+            )
+
+            def callback(msg: str):
+                # 处理 MCP 服务的错误日志
+                self.server_errlogs.append(msg)
+
+            stdio_transport = await self.exit_stack.enter_async_context(
+                mcp.stdio_client(
+                    server_params,
+                    errlog=LogPipe(
+                        level=logging.ERROR,
+                        logger=logger,
+                        identifier=f"MCPServer-{name}",
+                        callback=callback,
+                    ),  # type: ignore
+                ),
+            )
+
+            # 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"""
+        response = await self.session.list_tools()
+        self.tools = response.tools
+        return response
+
+    async def cleanup(self):
+        """Clean up resources"""
+        await self.exit_stack.aclose()
+        self.running_event.set()  # Set the running event to indicate cleanup is done

astrbot-4.0.0b2/astrbot/core/agent/response.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+import typing as T
+from astrbot.core.message.message_event_result import MessageChain
+
+class AgentResponseData(T.TypedDict):
+    chain: MessageChain
+
+
+@dataclass
+class AgentResponse:
+    type: str
+    data: AgentResponseData

astrbot-4.0.0b2/astrbot/core/agent/run_context.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from typing import Any, Generic
+from typing_extensions import TypeVar
+
+from astrbot.core.platform.astr_message_event import AstrMessageEvent
+
+TContext = TypeVar("TContext", default=Any)
+
+
+@dataclass
+class ContextWrapper(Generic[TContext]):
+    """A context for running an agent, which can be used to pass additional data or state."""
+
+    context: TContext
+    event: AstrMessageEvent
+
+NoContext = ContextWrapper[None]

astrbot-4.0.0b2/astrbot/core/agent/runners/__init__.py

@@ -0,0 +1,3 @@
+from .base import BaseAgentRunner
+
+__all__ = ["BaseAgentRunner"]

{astrbot-3.5.26/astrbot/core/pipeline/process_stage/agent_runner → astrbot-4.0.0b2/astrbot/core/agent/runners}/base.py

@@ -1,32 +1,33 @@
 import abc
 import typing as T
-from dataclasses import dataclass
-from astrbot.core.provider.entities import LLMResponse
-from ....message.message_event_result import MessageChain
 from enum import Enum, auto
+from ..run_context import ContextWrapper, TContext
+from ..response import AgentResponse
+from ..hooks import BaseAgentRunHooks
+from ..tool_executor import BaseFunctionToolExecutor
+from astrbot.core.provider import Provider
+from astrbot.core.provider.entities import LLMResponse
 
 
 class AgentState(Enum):
-    """Agent 状态枚举"""
-    IDLE = auto()          # 初始状态
-    RUNNING = auto()       # 运行中
-    DONE = auto()          # 完成
-    ERROR = auto()         # 错误状态
-
-
-class AgentResponseData(T.TypedDict):
-    chain: MessageChain
-
+    """Defines the state of the agent."""
 
-@dataclass
-class AgentResponse:
-    type: str
-    data: AgentResponseData
+    IDLE = auto()  # Initial state
+    RUNNING = auto()  # Currently processing
+    DONE = auto()  # Completed
+    ERROR = auto()  # Error state
 
 
-class BaseAgentRunner:
+class BaseAgentRunner(T.Generic[TContext]):
     @abc.abstractmethod
-    async def reset(self) -> None:
+    async def reset(
+        self,
+        provider: Provider,
+        run_context: ContextWrapper[TContext],
+        tool_executor: BaseFunctionToolExecutor[TContext],
+        agent_hooks: BaseAgentRunHooks[TContext],
+        **kwargs: T.Any,
+    ) -> None:
         """
         Reset the agent to its initial state.
         This method should be called before starting a new run.