diff --git a/.github/workflows/dashboard_ci.yml b/.github/workflows/dashboard_ci.yml index beb7b8871..63a7f47b3 100644 --- a/.github/workflows/dashboard_ci.yml +++ b/.github/workflows/dashboard_ci.yml @@ -27,9 +27,10 @@ jobs: cache-dependency-path: dashboard/pnpm-lock.yaml - name: Install and Build + working-directory: dashboard run: | - pnpm --dir dashboard install --frozen-lockfile - pnpm --dir dashboard run build + pnpm install --frozen-lockfile + pnpm run build - name: Inject Commit SHA id: get_sha diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bfe80a115..510338d26 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -58,17 +58,17 @@ jobs: - name: Setup Node.js uses: actions/setup-node@v6 with: - node-version: '24.13.0' + node-version: "24.13.0" cache: "pnpm" cache-dependency-path: dashboard/pnpm-lock.yaml - name: Build dashboard dist shell: bash + working-directory: dashboard run: | - pnpm --dir dashboard install --frozen-lockfile - pnpm --dir dashboard run build - echo "${{ steps.tag.outputs.tag }}" > dashboard/dist/assets/version - cd dashboard + pnpm install --frozen-lockfile + pnpm run build + echo "${{ steps.tag.outputs.tag }}" > dist/assets/version zip -r "AstrBot-${{ steps.tag.outputs.tag }}-dashboard.zip" dist - name: Upload dashboard artifact @@ -139,7 +139,6 @@ jobs: name: Dashboard-${{ steps.tag.outputs.tag }} path: release-assets - - name: Resolve release notes id: notes shell: bash diff --git a/.gitignore b/.gitignore index 085a9fda6..6e3b7394b 100644 --- a/.gitignore +++ b/.gitignore @@ -51,6 +51,7 @@ chroma venv/* pytest.ini IFLOW.md +CLAUDE.md # genie_tts data CharacterModels/ @@ -89,3 +90,4 @@ node_modules *pinokio* dashboard/pnpm-lock.yaml .obsidian +dashboard/.codex diff --git a/astrbot/_internal/__init__.py b/astrbot/_internal/__init__.py deleted file mode 100644 index 9d40e6d56..000000000 --- a/astrbot/_internal/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Astbot内部实现 -外部模块请勿导入 - -""" diff --git a/astrbot/_internal/abc/_abp/base_astrbot_abp_client.py b/astrbot/_internal/abc/_abp/base_astrbot_abp_client.py deleted file mode 100644 index b8ffa20a9..000000000 --- a/astrbot/_internal/abc/_abp/base_astrbot_abp_client.py +++ /dev/null @@ -1,54 +0,0 @@ -"""ABP (AstrBot Protocol) client - in-process star communication.""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Any - - -class BaseAstrbotAbpClient(ABC): - """ABP client: in-process star (plugin) communication. - - Stars register themselves; client delegates calls to registered instances. - - Subclass must implement: - - connect() -> None - - register_star(name, instance) -> None - - unregister_star(name) -> None - - call_star_tool(star, tool, args) -> Any - - shutdown() -> None - """ - - @property - @abstractmethod - def connected(self) -> bool: ... - - @abstractmethod - async def connect(self) -> None: - """Lightweight: just sets connected=True.""" - ... - - @abstractmethod - def register_star(self, star_name: str, star_instance: Any) -> None: - """Add star to internal registry.""" - ... - - @abstractmethod - def unregister_star(self, star_name: str) -> None: - """Remove star from registry (idempotent).""" - ... - - @abstractmethod - async def call_star_tool( - self, - star_name: str, - tool_name: str, - arguments: dict[str, Any], - ) -> Any: - """Delegate to star_instance.call_tool(tool_name, arguments).""" - ... - - @abstractmethod - async def shutdown(self) -> None: - """Set connected=False, cancel pending requests.""" - ... diff --git a/astrbot/_internal/abc/_acp/base_astrbot_acp_client.py b/astrbot/_internal/abc/_acp/base_astrbot_acp_client.py deleted file mode 100644 index 91f5207a0..000000000 --- a/astrbot/_internal/abc/_acp/base_astrbot_acp_client.py +++ /dev/null @@ -1,64 +0,0 @@ -"""ACP (AstrBot Communication Protocol) client. - -Transport: TCP | Unix Socket -Messages: JSON with Content-Length header -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Any - - -class BaseAstrbotAcpClient(ABC): - """ACP client: connects to ACP servers via TCP or Unix socket. - - Subclass must implement: - - connect() -> None - - connect_to_server(host, port) -> None - - connect_to_unix_socket(path) -> None - - call_tool(server, tool, args) -> Any - - send_notification(method, params) -> None - - shutdown() -> None - """ - - @property - @abstractmethod - def connected(self) -> bool: ... - - @abstractmethod - async def connect(self) -> None: ... - - @abstractmethod - async def connect_to_server(self, host: str, port: int) -> None: - """Connect via TCP.""" - ... - - @abstractmethod - async def connect_to_unix_socket(self, socket_path: str) -> None: - """Connect via Unix domain socket.""" - ... - - @abstractmethod - async def call_tool( - self, - server_name: str, - tool_name: str, - arguments: dict[str, Any], - ) -> Any: - """Call tool on server, return result.""" - ... - - @abstractmethod - async def send_notification( - self, - method: str, - params: dict[str, Any], - ) -> None: - """Send one-way notification.""" - ... - - @abstractmethod - async def shutdown(self) -> None: - """Close connection, cancel pending requests.""" - ... diff --git a/astrbot/_internal/abc/_acp/base_astrbot_acp_server.py b/astrbot/_internal/abc/_acp/base_astrbot_acp_server.py deleted file mode 100644 index 698ae1427..000000000 --- a/astrbot/_internal/abc/_acp/base_astrbot_acp_server.py +++ /dev/null @@ -1,66 +0,0 @@ -"""ACP (AstrBot Communication Protocol) server. - -Transport: TCP listening socket -Messages: JSON with Content-Length header -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from collections.abc import Callable -from typing import Any - - -class BaseAstrbotAcpServer(ABC): - """ACP server: listens for client connections, exposes tools. - - Subclass must implement: - - start(host, port) -> None - - register_tool(name, handler) -> None - - register_notification_handler(name, handler) -> None - - broadcast_notification(method, params) -> None - - shutdown() -> None - """ - - @property - @abstractmethod - def running(self) -> bool: - """True if server is accepting connections.""" - ... - - @abstractmethod - async def start(self, host: str = "127.0.0.1", port: int = 8765) -> None: - """Bind and listen. Block until shutdown.""" - ... - - @abstractmethod - def register_tool( - self, - name: str, - handler: Callable[..., Any], - ) -> None: - """Register async tool handler (receives params dict, returns result).""" - ... - - @abstractmethod - def register_notification_handler( - self, - name: str, - handler: Callable[..., Any], - ) -> None: - """Register async notification handler (receives params dict).""" - ... - - @abstractmethod - async def broadcast_notification( - self, - method: str, - params: dict[str, Any], - ) -> None: - """Send notification to all connected clients.""" - ... - - @abstractmethod - async def shutdown(self) -> None: - """Stop accepting, close all client connections.""" - ... diff --git a/astrbot/_internal/abc/_lsp/base_astrbot_lsp_client.py b/astrbot/_internal/abc/_lsp/base_astrbot_lsp_client.py deleted file mode 100644 index a92338342..000000000 --- a/astrbot/_internal/abc/_lsp/base_astrbot_lsp_client.py +++ /dev/null @@ -1,106 +0,0 @@ -"""LSP (Language Server Protocol) client. - -Transport: stdio subprocess -Messages: JSON-RPC 2.0 with Content-Length header -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Any - - -class LspMessage: - """JSON-RPC 2.0 message.""" - - jsonrpc: str = "2.0" - id: int | str | None = None - method: str | None = None - params: dict[str, Any] | None = None - result: Any = None - error: dict[str, Any] | None = None - - -class LspRequest(LspMessage): - """Outgoing request.""" - - def __init__(self, method: str, params: dict[str, Any] | None = None) -> None: - self.id = id(self) - self.method = method - self.params = params - - -class LspResponse(LspMessage): - """Incoming response.""" - - -class LspNotification(LspMessage): - """Incoming notification (no id).""" - - -class BaseAstrbotLspClient(ABC): - """LSP client: connects to LSP servers via stdio subprocess. - - Subclass must implement: - - connect() -> None - - connect_to_server(command, workspace_uri) -> None - - send_request(method, params) -> dict - - send_notification(method, params) -> None - - shutdown() -> None - """ - - @property - @abstractmethod - def connected(self) -> bool: - """True if connected to an LSP server.""" - ... - - @abstractmethod - async def connect(self) -> None: - self._connected = False - ... - - @abstractmethod - async def connect_to_server( - self, - command: list[str], - workspace_uri: str, - ) -> None: - """Start LSP server subprocess and complete handshake. - - Steps: - 1. Spawn subprocess with stdin/stdout pipes - 2. Send initialize request - 3. Wait for response - 4. Send initialized notification - """ - ... - - @abstractmethod - async def send_request( - self, - method: str, - params: dict[str, Any] | None = None, - ) -> Any: - """Send JSON-RPC request and return result. - - Raises: - RuntimeError: not connected - Exception: server returned error - - """ - ... - - @abstractmethod - async def send_notification( - self, - method: str, - params: dict[str, Any] | None = None, - ) -> None: - """Send JSON-RPC notification (no response expected).""" - ... - - @abstractmethod - async def shutdown(self) -> None: - """Send shutdown, terminate subprocess, cleanup.""" - ... diff --git a/astrbot/_internal/abc/_mcp/base_astrbot_mcp_client.py b/astrbot/_internal/abc/_mcp/base_astrbot_mcp_client.py deleted file mode 100644 index e68fa4426..000000000 --- a/astrbot/_internal/abc/_mcp/base_astrbot_mcp_client.py +++ /dev/null @@ -1,89 +0,0 @@ -"""MCP (Model Context Protocol) client. - -Transport: stdio | SSE | streamable_http -Messages: JSON-RPC 2.0 -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Any, Literal, TypedDict - - -class McpServerConfig(TypedDict, total=False): - """MCP server configuration.""" - - # Stdio transport - command: str - args: list[str] - env: dict[str, str] - cwd: str - - # HTTP transport - url: str - headers: dict[str, str] - transport: Literal["sse", "streamable_http"] - - -class McpToolInfo(TypedDict): - """MCP tool descriptor.""" - - name: str - description: str - inputSchema: dict[str, Any] - - -class BaseAstrbotMcpClient(ABC): - """MCP client: connects to MCP servers for external tools. - - Subclass must implement: - - connect() -> None - - connect_to_server(config, name) -> None - - list_tools() -> list[McpToolInfo] - - call_tool(name, args, timeout) -> CallToolResult - - cleanup() -> None - """ - - session: Any # mcp.ClientSession - - @property - @abstractmethod - def connected(self) -> bool: ... - - @abstractmethod - async def connect(self) -> None: - """Initialize client session.""" - ... - - @abstractmethod - async def connect_to_server( - self, - config: McpServerConfig, - name: str, - ) -> None: - """Connect to MCP server. - - Stdio: {"command": "python", "args": ["server.py"], "env": {...}} - HTTP: {"url": "https://...", "transport": "sse"} - """ - ... - - @abstractmethod - async def list_tools(self) -> list[McpToolInfo]: - """Call tools/list and return tools.""" - ... - - @abstractmethod - async def call_tool( - self, - name: str, - arguments: dict[str, Any], - read_timeout_seconds: int = 60, - ) -> Any: - """Call tools/call with reconnection support.""" - ... - - @abstractmethod - async def cleanup(self) -> None: - """Close all server connections.""" - ... diff --git a/astrbot/_internal/abc/base_astrbot_gateway.py b/astrbot/_internal/abc/base_astrbot_gateway.py deleted file mode 100644 index 990622b08..000000000 --- a/astrbot/_internal/abc/base_astrbot_gateway.py +++ /dev/null @@ -1,71 +0,0 @@ -"""AstrBot Gateway - HTTP/WebSocket API server. - -Built on FastAPI, provides: -- HTTP REST API (stats, inspector, config) -- WebSocket for real-time events -- Static file serving (dashboard) -- Authentication (JWT/API key) -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod - - -class BaseAstrbotGateway(ABC): - """Gateway: HTTP/WebSocket server built on FastAPI. - - ┌─────────────────────────────────────────────────────────┐ - │ FastAPI App │ - ├─────────────────────────────────────────────────────────┤ - │ REST Endpoints WebSocket │ - │ ├─ GET /api/stats ├─ /ws (connection manager)│ - │ ├─ GET /api/inspector/* │ │ - │ ├─ GET /api/memory/* │ │ - │ └─ ... │ │ - │ │ - │ Middleware: CORS, Auth, Logging │ - └─────────────────────────────────────────────────────────┘ - │ - ▼ - ┌─────────────────────────┐ - │ Orchestrator │ - │ (owns protocol clients)│ - └─────────────────────────┘ - - Routes (typical): - GET / → Dashboard static files - GET /api/stats → System statistics - GET /api/inspector/stars → List registered stars - WS /ws → WebSocket for real-time events - - serve() Lifecycle: - 1. Create FastAPI app - 2. Register routes - 3. Start WebSocket manager - 4. Bind to host:port - 5. Run ASGI server (uvicorn/hypercorn) - 6. Block until shutdown - 7. Close all connections - - Subclass must implement: - - serve(): start server, block until shutdown - """ - - @abstractmethod - async def serve(self) -> None: - """Start gateway server - blocks until shutdown. - - Should: - 1. Create FastAPI app with routes - 2. Configure CORS, auth middleware - 3. Start WebSocket connection manager - 4. Bind to ASTRBOT_PORT (default 6185) - 5. Run ASGI server - 6. Handle graceful shutdown on SIGTERM/SIGINT - - Raises: - OSError: address already in use - - """ - ... diff --git a/astrbot/_internal/abc/base_astrbot_orchestrator.py b/astrbot/_internal/abc/base_astrbot_orchestrator.py deleted file mode 100644 index b297fbada..000000000 --- a/astrbot/_internal/abc/base_astrbot_orchestrator.py +++ /dev/null @@ -1,350 +0,0 @@ -"""AstrBot Orchestrator - core runtime lifecycle manager. - -Architecture -============ - - ┌─────────────────────────────────────────────────────┐ - │ Orchestrator │ - │ (owns lifecycle of all protocol clients + stars) │ - └─────────────────────────────────────────────────────┘ - │ - ┌──────────────┼──────────────┐ - ▼ ▼ ▼ - ┌─────────┐ ┌─────────┐ ┌─────────┐ - │ LSP │ │ MCP │ │ ACP │ - │ Client │ │ Client │ │ Client │ - └─────────┘ └─────────┘ └─────────┘ - │ │ │ - ▼ ▼ ▼ - LSP Servers MCP Servers ACP Services - - ┌─────────────────────────────────────────────────────┐ - │ ABP Client │ - │ (in-process star registry) │ - └─────────────────────────────────────────────────────┘ - │ - ▼ - ┌─────────┐ - │ Stars │ - │(Plugins) │ - └─────────┘ - - -Lifecycle State Machine -======================= - - States: - ┌─────────┐ - │ INIT │───► orchestrator created, clients not initialized - └────┬────┘ - │ start() - ▼ - ┌─────────┐ - │ RUNNING │◄─── run_loop() executing - └────┬────┘ - │ shutdown() - ▼ - ┌──────────┐ - │ SHUTDOWN │─── all clients closed, ready for GC - └──────────┘ - - Transitions: - INIT + start() ──► RUNNING - RUNNING + shutdown() ──► SHUTDOWN - - For each protocol client, the orchestrator: - 1. Creates instance in __init__ - 2. Calls connect() to initialize - 3. Calls protocol-specific setup (connect_to_server, etc) - 4. Manages via run_loop() heartbeat - 5. Calls shutdown() on final cleanup - - -Star Registration Flow -===================== - - orchestrator.register_star("my-star", MyStar()) - │ - ▼ - ┌───────────────────┐ - │ ABP Client │ - │ .register_star() │ - └───────────────────┘ - │ - ▼ - ┌───────────────────┐ - │ Internal dict │ - │ {"my-star": obj} │ - └───────────────────┘ - - -Message Routing (conceptual) -=========================== - - External Tool Call - │ - ▼ - ┌──────────────┐ list_tools() ┌──────────────┐ - │ MCP Client │────────────────────►│ MCP Server │ - └──────────────┘◄────────────────────└──────────────┘ - │ tool result - ▼ - ┌──────────────┐ call_tool() ┌──────────────┐ - │ ABP │────────────────────►│ Star │ - │ Client │◄────────────────────└──────────────┘ - └──────────────┘ tool result - │ - ▼ - Return to caller - - -run_loop() Responsibilities -=========================== - - while running: - │─ check LSP server health (ping/heartbeat) - │─ check MCP session status (reconnect if needed) - │─ check ACP client connections - │─ process any pending star notifications - │─ sleep(SLEEP_INTERVAL) - - -Shutdown Sequence -================== - - shutdown() - │ - ├─ set _running = False - │ - ├─ LSP.shutdown() - │ └─ send "shutdown" request - │ └─ terminate subprocess - │ - ├─ ACP.shutdown() - │ └─ close TCP/Unix connections - │ - ├─ ABP.shutdown() - │ └─ cancel pending requests - │ - └─ MCP.cleanup() - └─ close all sessions - └─ cleanup subprocesses - - -Exception Handling -================== - - Each protocol client should: - - Catch connection errors - - Attempt reconnection with exponential backoff - - Log errors but don't crash run_loop - - Raise on irrecoverable failures - - The orchestrator run_loop should: - - Catch CancelledError on shutdown - - Catch Exception and log (don't crash) - - Ensure cleanup runs in finally block -""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from astrbot._internal.protocols.abp.client import AstrbotAbpClient - from astrbot._internal.protocols.acp.client import AstrbotAcpClient - from astrbot._internal.protocols.lsp.client import AstrbotLspClient - - from astrbot._internal.protocols._mcp.client import McpClient - - -#: Default heartbeat interval for run_loop() -DEFAULT_SLEEP_INTERVAL: float = 5.0 - - -class BaseAstrbotOrchestrator(ABC): - """Core runtime: owns lifecycle of all protocol clients and stars. - - ┌────────────────────────────────────────────────────────────┐ - │ Protocol Clients (always present, never None after init) │ - ├────────────────────────────────────────────────────────────┤ - │ lsp: Language Server Protocol │ - │ Purpose: code completion, diagnostics, hover, etc │ - │ Transport: stdio subprocess │ - │ │ - │ mcp: Model Context Protocol │ - │ Purpose: external tool access │ - │ Transport: stdio | SSE | HTTP │ - │ │ - │ acp: AstrBot Communication Protocol │ - │ Purpose: inter-service communication │ - │ Transport: TCP | Unix Socket │ - │ │ - │ abp: AstrBot Protocol │ - │ Purpose: in-process star (plugin) communication │ - │ Transport: direct method calls │ - └────────────────────────────────────────────────────────────┘ - - ┌────────────────────────────────────────────────────────────┐ - │ Star Registry │ - ├────────────────────────────────────────────────────────────┤ - │ _stars: dict[str, Any] │ - │ Stars are plugins registered by name │ - │ ABP client delegates calls to registered stars │ - └────────────────────────────────────────────────────────────┘ - - Subclass must implement: - - __init__(): create all protocol client instances - - run_loop(): main event loop (block until shutdown) - - register_star(name, instance): add to registry + ABP - - unregister_star(name): remove from registry + ABP - - shutdown(): clean up all clients - """ - - #: LSP client for language intelligence - lsp: AstrbotLspClient - - #: MCP client for external tools - mcp: McpClient - - #: ACP client for inter-service communication - acp: AstrbotAcpClient - - #: ABP client for in-process star communication - abp: AstrbotAbpClient - - def __init__(self) -> None: - """Initialize orchestrator and all protocol clients. - - After __init__, all clients exist but are not connected. - Call start() or run_loop() to begin operation. - - Example: - class MyOrchestrator(BaseAstrbotOrchestrator): - def __init__(self): - self.lsp = AstrbotLspClient() - self.mcp = McpClient() - self.acp = AstrbotAcpClient() - self.abp = AstrbotAbpClient() - self._stars: dict[str, Any] = {} - self._running = False - - """ - self._stars: dict[str, Any] = {} - self._running: bool = False - - @property - def running(self) -> bool: - """True if run_loop() is executing.""" - return self._running - - @abstractmethod - async def start(self) -> None: - """Initialize all protocol clients. - - Called once before run_loop(). Should: - 1. Call lsp.connect() - 2. Call mcp.connect() - 3. Call acp.connect() - 4. Call abp.connect() - 5. Set _running = True - - Raises: - Exception: if any client fails to initialize - - """ - ... - - @abstractmethod - async def run_loop(self) -> None: - """Main event loop - blocks until shutdown. - - Execution: - self._running = True - try: - while self._running: - await self._heartbeat() - await anyio.sleep(DEFAULT_SLEEP_INTERVAL) - except asyncio.CancelledError: - pass # shutdown requested - finally: - self._running = False - - _heartbeat() responsibilities: - - Check LSP server health (optional ping) - - Check MCP session status, reconnect if needed - - Check ACP connections - - Process any pending star notifications - - Raises: - asyncio.CancelledError: when shutdown() called - - Note: - Subclass defines _heartbeat() for periodic tasks. - This method only handles the loop control. - - """ - ... - - @abstractmethod - async def register_star(self, name: str, star_instance: Any) -> None: - """Register a star (plugin) with the orchestrator. - - Args: - name: Unique identifier for the star - instance: Star plugin instance (must have .call_tool() method) - - Does: - self._stars[name] = star_instance - self.abp.register_star(name, star_instance) - - Raises: - ValueError: if name already registered - - """ - ... - - @abstractmethod - async def unregister_star(self, name: str) -> None: - """Unregister a star (plugin) from the orchestrator. - - Args: - name: Identifier of star to remove - - Does: - del self._stars[name] - self.abp.unregister_star(name) - - Note: - Idempotent - does nothing if name not found. - - """ - ... - - @abstractmethod - async def get_star(self, name: str) -> Any | None: - """Get registered star by name. Returns None if not found.""" - ... - - @abstractmethod - async def list_stars(self) -> list[str]: - """Return list of registered star names.""" - ... - - @abstractmethod - async def shutdown(self) -> None: - """Graceful shutdown of orchestrator and all clients. - - Execution order: - 1. self._running = False (stop run_loop) - 2. await lsp.shutdown() - 3. await acp.shutdown() - 4. await abp.shutdown() - 5. await mcp.cleanup() - - Does NOT unregister stars - caller should do that first. - - After shutdown, orchestrator is ready for garbage collection. - """ - ... diff --git a/astrbot/_internal/geteway/__init__.py b/astrbot/_internal/geteway/__init__.py deleted file mode 100644 index b88ac0e3b..000000000 --- a/astrbot/_internal/geteway/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Gateway module - FastAPI server for the dashboard backend.""" - -from .server import AstrbotGateway -from .ws_manager import WebSocketManager - -__all__ = ["AstrbotGateway", "WebSocketManager"] diff --git a/astrbot/_internal/geteway/deps.py b/astrbot/_internal/geteway/deps.py deleted file mode 100644 index fe2b8bb49..000000000 --- a/astrbot/_internal/geteway/deps.py +++ /dev/null @@ -1 +0,0 @@ -"""依赖注入""" diff --git a/astrbot/_internal/geteway/routes/inspector.py b/astrbot/_internal/geteway/routes/inspector.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/astrbot/_internal/geteway/routes/memory.py b/astrbot/_internal/geteway/routes/memory.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/astrbot/_internal/geteway/routes/stats.py b/astrbot/_internal/geteway/routes/stats.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/astrbot/_internal/geteway/server.py b/astrbot/_internal/geteway/server.py deleted file mode 100644 index 28f03686d..000000000 --- a/astrbot/_internal/geteway/server.py +++ /dev/null @@ -1,221 +0,0 @@ -"""AstrBot Gateway - FastAPI server for the dashboard backend. - -Provides REST API endpoints and WebSocket connections for the frontend dashboard. -The gateway acts as the communication bridge between the dashboard and the orchestrator. -""" - -from __future__ import annotations - -import json -from contextlib import asynccontextmanager -from typing import TYPE_CHECKING, Any - -from astrbot import logger -from astrbot._internal.abc.base_astrbot_gateway import BaseAstrbotGateway -from astrbot._internal.abc.base_astrbot_orchestrator import BaseAstrbotOrchestrator -from astrbot._internal.geteway.ws_manager import WebSocketManager - -if TYPE_CHECKING: - from fastapi import FastAPI, WebSocket, WebSocketDisconnect -else: - try: - from fastapi import FastAPI, WebSocket, WebSocketDisconnect - except ImportError: - logger.warning("FastAPI not installed, gateway unavailable.") - FastAPI = None - WebSocket = None - WebSocketDisconnect = None -from fastapi.middleware.cors import CORSMiddleware - -log = logger - - -class AstrbotGateway(BaseAstrbotGateway): - """FastAPI-based gateway server for AstrBot. - - Handles: - - REST API endpoints for configuration and stats - - WebSocket connections for real-time communication - - CORS middleware for dashboard access - """ - - def __init__(self, orchestrator: BaseAstrbotOrchestrator) -> None: - self.orchestrator = orchestrator - self.ws_manager = WebSocketManager() - self._app: FastAPI | None = None - self._host = "0.0.0.0" - self._port = 8765 - - async def serve(self) -> None: - """Start the gateway server. - - Creates and runs a FastAPI application with WebSocket support. - """ - if FastAPI is None: - raise RuntimeError("FastAPI is not installed") - log.info(f"Starting AstrBot Gateway on {self._host}:{self._port}") - - @asynccontextmanager - async def lifespan(app: FastAPI): - log.info("Gateway server started.") - yield - await self.ws_manager.broadcast({"type": "server_shutdown"}) - log.info("Gateway server stopped.") - - self._app = FastAPI( - title="AstrBot Gateway", - description="Backend API for AstrBot dashboard", - version="1.0.0", - lifespan=lifespan, - ) - self._app.add_middleware( - CORSMiddleware, - allow_origins=["*"], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], - ) - self._setup_routes() - import uvicorn - - config = uvicorn.Config( - self._app, - host=self._host, - port=self._port, - log_level="info", - ) - server = uvicorn.Server(config) - await server.serve() - - def _setup_routes(self) -> None: - """Set up API routes.""" - if self._app is None: - return - from fastapi import APIRouter - - @self._app.get("/health") - async def health(): - return {"status": "ok"} - - @self._app.websocket("/ws") - async def websocket_endpoint(ws: WebSocket): - await self.ws_manager.connect(ws) - try: - while True: - data = await ws.receive_text() - try: - message = json.loads(data) - response = await self._handle_ws_message(message) - if response: - await ws.send_json(response) - except json.JSONDecodeError: - await ws.send_json({"error": "Invalid JSON"}) - except WebSocketDisconnect: - self.ws_manager.disconnect(ws) - - stats_router = APIRouter(prefix="/api/stats", tags=["stats"]) - - @stats_router.get("/overview") - async def get_overview(): - return await self._get_stats_overview() - - self._app.include_router(stats_router) - inspector_router = APIRouter(prefix="/api/inspector", tags=["inspector"]) - - @inspector_router.get("/stars") - async def list_stars(): - return await self._list_stars() - - @inspector_router.get("/stars/{star_name}") - async def get_star(star_name: str): - return await self._get_star_detail(star_name) - - self._app.include_router(inspector_router) - memory_router = APIRouter(prefix="/api/memory", tags=["memory"]) - - @memory_router.get("/") - async def get_memory(): - return await self._get_memory_info() - - self._app.include_router(memory_router) - - async def _handle_ws_message( - self, - message: dict[str, Any], - ) -> dict[str, Any] | None: - """Handle an incoming WebSocket message. - - Args: - message: Parsed JSON message from the client - - Returns: - Response message to send back, or None for no response - - """ - msg_type = message.get("type") - data = message.get("data", {}) - if msg_type == "ping": - return {"type": "pong", "data": {}} - if msg_type == "call_tool": - return await self._handle_call_tool(data) - if msg_type == "get_stars": - return {"type": "stars_list", "data": await self._list_stars()} - return { - "type": "error", - "data": {"message": f"Unknown message type: {msg_type}"}, - } - - async def _handle_call_tool(self, data: dict[str, Any]) -> dict[str, Any]: - """Handle a tool call request via WebSocket.""" - star_name = data.get("star") - tool_name = data.get("tool") - arguments = data.get("arguments", {}) - if not star_name or not tool_name: - return { - "type": "tool_result", - "data": {"error": "Missing star or tool name"}, - } - try: - result = await self.orchestrator.abp.call_star_tool( - star_name, - tool_name, - arguments, - ) - return {"type": "tool_result", "data": {"result": result}} - except Exception as e: - return {"type": "tool_result", "data": {"error": str(e)}} - - async def _get_stats_overview(self) -> dict[str, Any]: - """Get overview statistics.""" - return { - "stars_count": len(self.orchestrator.abp._stars), - "lsp_connected": self.orchestrator.lsp._connected, - "mcp_sessions": getattr(self.orchestrator.mcp, "session", None) is not None, - "acp_clients": len(getattr(self.orchestrator.acp, "_clients", [])), - } - - async def _list_stars(self) -> list[dict[str, Any]]: - """List all registered stars.""" - stars = [] - for name in self.orchestrator.abp._stars: - stars.append({"name": name, "status": "active"}) - return stars - - async def _get_star_detail(self, star_name: str) -> dict[str, Any]: - """Get details of a specific star.""" - star = self.orchestrator.abp._stars.get(star_name) - if not star: - return {"error": f"Star '{star_name}' not found"} - return {"name": star_name, "status": "active"} - - async def _get_memory_info(self) -> dict[str, Any]: - """Get memory usage information.""" - import gc - - gc.collect() - return {"gc_objects": len(gc.get_objects()), "python_memory": "N/A"} - - def set_listen_address(self, host: str, port: int) -> None: - """Set the listen address for the gateway server.""" - self._host = host - self._port = port diff --git a/astrbot/_internal/geteway/ws_manager.py b/astrbot/_internal/geteway/ws_manager.py deleted file mode 100644 index b3c82697b..000000000 --- a/astrbot/_internal/geteway/ws_manager.py +++ /dev/null @@ -1,100 +0,0 @@ -"""WebSocket connection manager for the AstrBot gateway.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -import anyio - -from astrbot import logger - -if TYPE_CHECKING: - from fastapi import WebSocket -else: - try: - from fastapi import WebSocket - except ImportError: - logger.warning("FastAPI not installed, WebSocketManager unavailable.") - WebSocket = None -log = logger - - -class WebSocketManager: - """Manages all active WebSocket connections. - - Provides connection/disconnection handling and broadcast capabilities. - """ - - def __init__(self) -> None: - self._connections: set[WebSocket] = set() - self._lock = anyio.Lock() - - async def connect(self, websocket: WebSocket) -> None: - """Accept and register a new WebSocket connection.""" - await websocket.accept() - async with self._lock: - self._connections.add(websocket) - log.debug(f"WebSocket connected. Total: {len(self._connections)}") - - async def disconnect(self, websocket: WebSocket) -> None: - """Remove a WebSocket connection.""" - async with self._lock: - self._connections.discard(websocket) - log.debug(f"WebSocket disconnected. Total: {len(self._connections)}") - - async def send_json(self, websocket: WebSocket, data: dict[str, Any]) -> None: - """Send JSON data to a specific WebSocket. - - Args: - websocket: Target WebSocket connection - data: Data to send (must be JSON-serializable) - - """ - try: - await websocket.send_json(data) - except Exception as e: - log.warning(f"Failed to send to WebSocket: {e}") - await self.disconnect(websocket) - - async def broadcast(self, data: dict[str, Any]) -> None: - """Broadcast JSON data to all connected WebSockets. - - Args: - data: Data to broadcast (must be JSON-serializable) - - """ - async with self._lock: - connections = list(self._connections) - for conn in connections: - try: - await conn.send_json(data) - except Exception as e: - log.warning(f"Failed to broadcast to WebSocket: {e}") - async with self._lock: - self._connections.discard(conn) - - async def send_to( - self, - websocket: WebSocket, - message: str | dict[str, Any], - ) -> None: - """Send a message to a specific WebSocket. - - Args: - websocket: Target WebSocket connection - message: Message to send (string or dict) - - """ - try: - if isinstance(message, str): - await websocket.send_text(message) - else: - await websocket.send_json(message) - except Exception as e: - log.warning(f"Failed to send to WebSocket: {e}") - await self.disconnect(websocket) - - @property - def connection_count(self) -> int: - """Return the number of active connections.""" - return len(self._connections) diff --git a/astrbot/_internal/protocols/_abp/__init__.py b/astrbot/_internal/protocols/_abp/__init__.py deleted file mode 100644 index 54f74818f..000000000 --- a/astrbot/_internal/protocols/_abp/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""ABP module - AstrBot Protocol client implementation (built-in plugin protocol).""" - -from .client import AstrbotAbpClient - -__all__ = ["AstrbotAbpClient"] diff --git a/astrbot/_internal/protocols/_abp/client.py b/astrbot/_internal/protocols/_abp/client.py deleted file mode 100644 index 00cd7c9b9..000000000 --- a/astrbot/_internal/protocols/_abp/client.py +++ /dev/null @@ -1,94 +0,0 @@ -"""ABP (AstrBot Protocol) client implementation. - -ABP is the built-in plugin protocol where the orchestrator acts as client -connecting to internal stars (plugins) embedded in the runtime. -""" - -from __future__ import annotations - -from typing import Any - -from astrbot import logger -from astrbot._internal.abc._abp.base_astrbot_abp_client import BaseAstrbotAbpClient - -log = logger - - -class AstrbotAbpClient(BaseAstrbotAbpClient): - """ABP client for communicating with internal stars (built-in plugins). - - The orchestrator acts as the client, sending requests to and receiving - notifications from stars running within the same process. - """ - - def __init__(self) -> None: - self._connected = False - self._stars: dict[str, Any] = {} - # Use a simple dict for pending requests; we avoid asyncio.Future here. - self._pending_requests: dict[str, Any] = {} - self._request_id = 0 - - @property - def connected(self) -> bool: - """True if connected to stars registry.""" - return self._connected - - async def connect(self) -> None: - """Connect to internal stars registry.""" - log.debug("ABP client connecting to internal stars...") - self._connected = True - log.info("ABP client connected to internal stars registry.") - - async def call_star_tool( - self, - star_name: str, - tool_name: str, - arguments: dict[str, Any], - ) -> Any: - """Call a tool on a registered star. - - Args: - star_name: Name of the star (plugin) - tool_name: Name of the tool to call - arguments: Tool arguments - - Returns: - Tool call result - - """ - if not self._connected: - raise RuntimeError("ABP client is not connected") - - star = self._stars.get(star_name) - if not star: - raise ValueError(f"Star '{star_name}' not found") - - request_id = f"{self._request_id}" - self._request_id += 1 - - # No asyncio.Future used; store a placeholder entry for tracking if needed. - self._pending_requests[request_id] = None - - try: - # Call the star's tool handler - result = await star.call_tool(tool_name, arguments) - return result - finally: - self._pending_requests.pop(request_id, None) - - def register_star(self, star_name: str, star_instance: Any) -> None: - """Register a star (plugin) with the ABP client.""" - self._stars[star_name] = star_instance - log.debug(f"Star '{star_name}' registered with ABP client.") - - def unregister_star(self, star_name: str) -> None: - """Unregister a star from the ABP client.""" - self._stars.pop(star_name, None) - log.debug(f"Star '{star_name}' unregistered from ABP client.") - - async def shutdown(self) -> None: - """Shutdown the ABP client connection.""" - self._connected = False - # Clear any pending requests (no asyncio futures used in this implementation) - self._pending_requests.clear() - log.info("ABP client shut down.") diff --git a/astrbot/_internal/protocols/_acp/__init__.py b/astrbot/_internal/protocols/_acp/__init__.py deleted file mode 100644 index 853768409..000000000 --- a/astrbot/_internal/protocols/_acp/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""ACP module - AstrBot Communication Protocol client and server implementations.""" - -from .client import AstrbotAcpClient -from .server import AstrbotAcpServer - -__all__ = ["AstrbotAcpClient", "AstrbotAcpServer"] diff --git a/astrbot/_internal/protocols/_acp/client.py b/astrbot/_internal/protocols/_acp/client.py deleted file mode 100644 index 65d8a4948..000000000 --- a/astrbot/_internal/protocols/_acp/client.py +++ /dev/null @@ -1,222 +0,0 @@ -"""ACP (AstrBot Communication Protocol) client implementation. - -ACP is a client-server protocol for inter-service communication, -similar to MCP but designed specifically for AstrBot's architecture. -""" - -from __future__ import annotations - -import asyncio -import json -from typing import Any - -from astrbot import logger -from astrbot._internal.abc._acp.base_astrbot_acp_client import BaseAstrbotAcpClient - -log = logger - - -class AstrbotAcpClient(BaseAstrbotAcpClient): - """ACP client for communicating with ACP servers. - - The orchestrator acts as an ACP client, connecting to external - ACP-compatible services. - """ - - def __init__(self) -> None: - self._connected = False - self._reader: asyncio.StreamReader | None = None - self._writer: asyncio.StreamWriter | None = None - self._server_url: str | None = None - self._pending_requests: dict[str, asyncio.Future[dict[str, Any]]] = {} - self._request_id = 0 - self._reader_task: asyncio.Task[None] | None = None - - @property - def connected(self) -> bool: - """True if connected to an ACP server.""" - return self._connected - - async def connect(self) -> None: - """Connect to configured ACP servers. - - ACP servers can be accessed via TCP (host:port) or Unix socket. - """ - log.debug("ACP client connecting...") - # TODO: Load ACP server configurations - self._connected = True - log.info("ACP client initialized.") - - async def connect_to_server(self, host: str, port: int) -> None: - """Connect to an ACP server via TCP. - - Args: - host: Server hostname or IP - port: Server port - - """ - self._server_url = f"{host}:{port}" - self._reader, self._writer = await asyncio.open_connection(host, port) - self._connected = True - - # Start reading responses - self._reader_task = asyncio.create_task(self._read_messages()) - - log.info(f"ACP client connected to {self._server_url}") - - async def connect_to_unix_socket(self, socket_path: str) -> None: - """Connect to an ACP server via Unix socket. - - Args: - socket_path: Path to the Unix socket - - """ - self._server_url = f"unix://{socket_path}" - self._reader, self._writer = await asyncio.open_unix_connection(socket_path) - self._connected = True - - self._reader_task = asyncio.create_task(self._read_messages()) - - log.info(f"ACP client connected to {self._server_url}") - - async def _read_messages(self) -> None: - """Background task to read ACP messages.""" - if not self._reader: - return - - buffer = b"" - while self._connected: - try: - data = await self._reader.read(4096) - if not data: - break - buffer += data - - while True: - header_end = buffer.find(b"\n") - if header_end == -1: - break - - try: - header = json.loads(buffer[:header_end].decode("utf-8")) - except json.JSONDecodeError: - buffer = buffer[header_end + 1 :] - continue - - content_length = header.get("content-length", 0) - if ( - content_length == 0 - or len(buffer) < header_end + 1 + content_length - ): - break - - content = buffer[header_end + 1 : header_end + 1 + content_length] - buffer = buffer[header_end + 1 + content_length :] - - message = json.loads(content.decode("utf-8")) - - if "id" in message: - request_id = str(message["id"]) - future = self._pending_requests.pop(request_id, None) - if future and not future.done(): - if "error" in message: - future.set_exception(Exception(str(message["error"]))) - else: - future.set_result(message.get("result", {})) - else: - await self._handle_notification(message) - - except Exception as e: - if self._connected: - log.error(f"ACP read error: {e}") - break - - async def _handle_notification(self, notification: dict[str, Any]) -> None: - """Handle incoming ACP notifications.""" - method = notification.get("method", "") - log.debug(f"ACP notification: {method}") - - async def call_tool( - self, - server_name: str, - tool_name: str, - arguments: dict[str, Any], - ) -> Any: - """Call a tool on an ACP server. - - Args: - server_name: Name of the ACP server - tool_name: Name of the tool to call - arguments: Tool arguments - - Returns: - Tool call result - - """ - if not self._connected: - raise RuntimeError("ACP client is not connected") - - request_id = str(self._request_id) - self._request_id += 1 - - message = { - "jsonrpc": "2.0", - "id": request_id, - "method": f"{server_name}/{tool_name}", - "params": arguments, - } - - future: asyncio.Future[dict[str, Any]] = asyncio.Future() - self._pending_requests[request_id] = future - - await self._send_message(message) - return await future - - async def _send_message(self, message: dict[str, Any]) -> None: - """Send an ACP message.""" - if not self._writer: - raise RuntimeError("ACP client not connected") - - content = json.dumps(message) - header = json.dumps({"content-length": len(content)}) + "\n" - - self._writer.write((header + content).encode()) - await self._writer.drain() - - async def send_notification( - self, - method: str, - params: dict[str, Any] | None = None, - ) -> None: - """Send a one-way notification to the server.""" - message = { - "jsonrpc": "2.0", - "method": method, - "params": params or {}, - } - await self._send_message(message) - - async def shutdown(self) -> None: - """Shutdown the ACP client connection.""" - self._connected = False - - if self._reader_task: - self._reader_task.cancel() - try: - await self._reader_task - except asyncio.CancelledError: - pass - - if self._writer: - self._writer.close() - try: - await self._writer.wait_closed() - except Exception: - pass - - for future in self._pending_requests.values(): - if not future.done(): - future.cancel() - self._pending_requests.clear() - - log.info("ACP client shut down.") diff --git a/astrbot/_internal/protocols/_acp/server.py b/astrbot/_internal/protocols/_acp/server.py deleted file mode 100644 index 910de7009..000000000 --- a/astrbot/_internal/protocols/_acp/server.py +++ /dev/null @@ -1,225 +0,0 @@ -"""ACP (AstrBot Communication Protocol) server implementation. - -ACP servers listen for connections from ACP clients and provide -services/tools to the orchestrator. -""" - -from __future__ import annotations - -import asyncio -import json -from collections.abc import Callable -from typing import Any - -from astrbot import logger -from astrbot._internal.abc._acp.base_astrbot_acp_server import BaseAstrbotAcpServer - -log = logger - - -class AstrbotAcpServer(BaseAstrbotAcpServer): - """ACP server for accepting connections from ACP clients. - - ACP servers expose tools/notifications that can be called by clients. - """ - - def __init__(self) -> None: - self._running = False - self._host: str = "127.0.0.1" - self._port: int = 8765 - self._server: asyncio.Server | None = None - self._clients: set[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = set() - self._tool_handlers: dict[str, Callable[..., Any]] = {} - self._notification_handlers: dict[str, Callable[..., Any]] = {} - - def register_tool(self, name: str, handler: Callable[..., Any]) -> None: - """Register a tool handler. - - Args: - name: Tool name - handler: Async callable that handles tool calls - - """ - self._tool_handlers[name] = handler - log.debug(f"ACP server registered tool: {name}") - - def register_notification_handler( - self, - name: str, - handler: Callable[..., Any], - ) -> None: - """Register a notification handler. - - Args: - name: Notification method name - handler: Async callable that handles notifications - - """ - self._notification_handlers[name] = handler - log.debug(f"ACP server registered notification handler: {name}") - - async def start(self, host: str = "127.0.0.1", port: int = 8765) -> None: - """Start the ACP server. - - Args: - host: Host to bind to - port: Port to listen on - - """ - self._host = host - self._port = port - self._server = await asyncio.start_server( - self._handle_client, - host=host, - port=port, - ) - self._running = True - log.info(f"ACP server listening on {host}:{port}") - - async def _handle_client( - self, - reader: asyncio.StreamReader, - writer: asyncio.StreamWriter, - ) -> None: - """Handle an incoming ACP client connection.""" - addr = writer.get_extra_info("peername") - log.debug(f"ACP client connected: {addr}") - self._clients.add((reader, writer)) - - buffer = b"" - try: - while self._running: - try: - data = await reader.read(4096) - if not data: - break - buffer += data - - while True: - header_end = buffer.find(b"\n") - if header_end == -1: - break - - try: - header = json.loads(buffer[:header_end].decode("utf-8")) - except json.JSONDecodeError: - buffer = buffer[header_end + 1 :] - continue - - content_length = header.get("content-length", 0) - if ( - content_length == 0 - or len(buffer) < header_end + 1 + content_length - ): - break - - content = buffer[ - header_end + 1 : header_end + 1 + content_length - ] - buffer = buffer[header_end + 1 + content_length :] - - message = json.loads(content.decode("utf-8")) - response = await self._handle_message(message) - - if response: - content = json.dumps(response) - resp_header = ( - json.dumps({"content-length": len(content)}) + "\n" - ) - writer.write(resp_header.encode() + content.encode()) - await writer.drain() - - except Exception as e: - log.error(f"ACP client error ({addr}): {e}") - break - - finally: - self._clients.discard((reader, writer)) - writer.close() - try: - await writer.wait_closed() - except Exception: - pass - log.debug(f"ACP client disconnected: {addr}") - - async def _handle_message(self, message: dict[str, Any]) -> dict[str, Any] | None: - """Handle an incoming ACP message.""" - method = message.get("method", "") - msg_id = message.get("id") - params = message.get("params", {}) - - # Check if it's a notification (no id) or request (has id) - if msg_id is None: - # Notification - handler = self._notification_handlers.get(method) - if handler: - try: - await handler(params) - except Exception as e: - log.error(f"ACP notification handler error ({method}): {e}") - return None - - # Request - result = None - error = None - - handler = self._tool_handlers.get(method) - if handler: - try: - result = await handler(params) - except Exception as e: - error = str(e) - log.error(f"ACP tool handler error ({method}): {e}") - else: - error = f"Unknown method: {method}" - - response: dict[str, Any] = {"jsonrpc": "2.0", "id": msg_id} - if error: - response["error"] = {"code": -32601, "message": error} - else: - response["result"] = result - - return response - - async def broadcast_notification(self, method: str, params: dict[str, Any]) -> None: - """Broadcast a notification to all connected clients. - - Args: - method: Notification method name - params: Notification parameters - - """ - message = { - "jsonrpc": "2.0", - "method": method, - "params": params, - } - content = json.dumps(message) - header = json.dumps({"content-length": len(content)}) + "\n" - data = header.encode() + content.encode() - - for reader, writer in list(self._clients): - try: - writer.write(data) - await writer.drain() - except Exception as e: - log.warning(f"Failed to broadcast to client: {e}") - - async def shutdown(self) -> None: - """Shutdown the ACP server.""" - self._running = False - - if self._server: - self._server.close() - await self._server.wait_closed() - self._server = None - - for reader, writer in list(self._clients): - writer.close() - try: - await writer.wait_closed() - except Exception: - pass - self._clients.clear() - - log.info("ACP server shut down.") diff --git a/astrbot/_internal/protocols/_lsp/__init__.py b/astrbot/_internal/protocols/_lsp/__init__.py deleted file mode 100644 index f7708a27d..000000000 --- a/astrbot/_internal/protocols/_lsp/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""LSP module - Language Server Protocol client implementation.""" - -from .client import AstrbotLspClient - -__all__ = ["AstrbotLspClient"] diff --git a/astrbot/_internal/protocols/_lsp/client.py b/astrbot/_internal/protocols/_lsp/client.py deleted file mode 100644 index 27bbbdc23..000000000 --- a/astrbot/_internal/protocols/_lsp/client.py +++ /dev/null @@ -1,244 +0,0 @@ -"""LSP (Language Server Protocol) client implementation. - -The orchestrator acts as an LSP client, connecting to LSP servers -that provide language intelligence features (completions, diagnostics, etc.). -""" - -from __future__ import annotations - -import json -from typing import Any - -import anyio -from anyio.abc import ByteReceiveStream, ByteSendStream, Process - -from astrbot import logger -from astrbot._internal.abc._lsp.base_astrbot_lsp_client import BaseAstrbotLspClient - -log = logger - - -class AstrbotLspClient(BaseAstrbotLspClient): - """LSP client for communicating with LSP servers. - - Implements the Microsoft Language Server Protocol for connecting to - external language intelligence services. - """ - - def __init__(self) -> None: - self._connected = False - self._reader: ByteReceiveStream | None = None - self._writer: ByteSendStream | None = None - self._server_process: Process | None = None - self._pending_requests: dict[int, Any] = {} - self._request_id = 0 - self._server_command: list[str] | None = None - # anyio TaskGroup handle for background readers - self._task_group: Any | None = None - - @property - def connected(self) -> bool: - """True if connected to an LSP server.""" - return self._connected - - async def connect(self) -> None: - """Connect to configured LSP servers. - - LSP servers are typically stdio-based subprocesses. This method - establishes the communication channel. - """ - log.debug("LSP client connecting...") - # TODO: Load LSP server configurations and start subprocesses - # For now, mark as connected in idle mode - self._connected = True - log.info("LSP client initialized.") - - async def connect_to_server(self, command: list[str], workspace_uri: str) -> None: - """Connect to an LSP server subprocess. - - Args: - command: Command line to start the LSP server (e.g., ["python", "lsp_server.py"]) - workspace_uri: Root URI of the workspace to serve - - """ - log.debug(f"Starting LSP server: {' '.join(command)}") - - self._server_process = await anyio.open_process( - command, - stdin=-1, - stdout=-1, - stderr=-1, - ) - self._reader = self._server_process.stdout - self._writer = self._server_process.stdin - self._server_command = command - self._connected = True - - # Start reading responses in background using anyio TaskGroup - # Create and enter a TaskGroup so the reader runs until we close it at shutdown. - self._task_group = anyio.create_task_group() - await self._task_group.__aenter__() - self._task_group.start_soon(self._read_responses) - - # Send initialize request - await self.send_request( - "initialize", - { - "processId": None, - "rootUri": workspace_uri, - "capabilities": {}, - }, - ) - - # Send initialized notification - await self.send_notification("initialized", {}) - - log.info(f"LSP client connected to server: {command[0]}") - - async def send_request( - self, - method: str, - params: dict[str, Any] | None = None, - ) -> Any: - """Send an LSP request and wait for response.""" - if not self._writer: - raise RuntimeError("LSP client not connected") - - request_id = self._request_id - self._request_id += 1 - - message = { - "jsonrpc": "2.0", - "id": request_id, - "method": method, - "params": params or {}, - } - - # Use anyio.Event for request/response matching - response_event: anyio.Event = anyio.Event() - response_holder: dict[str, Any] = {} - - async def set_response(response: dict[str, Any]) -> None: - response_holder["response"] = response - response_event.set() - - self._pending_requests[request_id] = set_response - - content = json.dumps(message) - headers = f"Content-Length: {len(content)}\r\n\r\n" - await self._writer.send((headers + content).encode()) - - # Wait for response with timeout - with anyio.move_on_after(30): - await response_event.wait() - - if "response" in response_holder: - return response_holder["response"] - raise TimeoutError(f"LSP request {method} timed out") - - async def send_notification( - self, - method: str, - params: dict[str, Any] | None = None, - ) -> None: - """Send an LSP notification (no response expected).""" - if not self._writer: - raise RuntimeError("LSP client not connected") - - message = { - "jsonrpc": "2.0", - "method": method, - "params": params or {}, - } - - content = json.dumps(message) - headers = f"Content-Length: {len(content)}\r\n\r\n" - await self._writer.send((headers + content).encode()) - - async def _read_responses(self) -> None: - """Background task to read LSP responses.""" - if not self._reader: - return - - buffer = b"" - try: - while self._connected: - try: - data = await self._reader.receive() - if not data: - break - buffer += data - - while True: - # Parse Content-Length header - header_end = buffer.find(b"\r\n\r\n") - if header_end == -1: - break - - header = buffer[:header_end].decode("utf-8") - content_length = 0 - for line in header.split("\r\n"): - if line.startswith("Content-Length:"): - content_length = int(line.split(":")[1].strip()) - - if content_length == 0: - break - - total_length = header_end + 4 + content_length - if len(buffer) < total_length: - break - - content = buffer[header_end + 4 : total_length] - buffer = buffer[total_length:] - - response = json.loads(content.decode("utf-8")) - - # Handle response vs notification - if "id" in response: - request_id = response["id"] - handler = self._pending_requests.pop(request_id, None) - if handler: - await handler(response) - else: - # Notification (e.g., window/logMessage) - await self._handle_notification(response) - - except anyio.EndOfStream: - break - except anyio.get_cancelled_exc_class(): - # Task was cancelled via the TaskGroup cancel/exit during shutdown - pass - - async def _handle_notification(self, notification: dict[str, Any]) -> None: - """Handle incoming LSP notifications.""" - method = notification.get("method", "") - log.debug(f"LSP notification: {method}") - - async def shutdown(self) -> None: - """Shutdown the LSP client.""" - self._connected = False - - if self._task_group: - try: - # Exit the TaskGroup, which cancels background tasks started within it - await self._task_group.__aexit__(None, None, None) - except anyio.get_cancelled_exc_class(): - pass - self._task_group = None - - if self._server_process: - try: - await self.send_notification("shutdown", {}) - except Exception: - pass - - self._server_process.terminate() - try: - with anyio.move_on_after(5): - await self._server_process.wait() - except Exception: - self._server_process.kill() - self._server_process = None - - self._pending_requests.clear() - log.info("LSP client shut down.") diff --git a/astrbot/_internal/protocols/_mcp/__init__.py b/astrbot/_internal/protocols/_mcp/__init__.py deleted file mode 100644 index 7826f38f4..000000000 --- a/astrbot/_internal/protocols/_mcp/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -"""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", - "MCPInitError", - "MCPInitSummary", - "MCPInitTimeoutError", - "MCPShutdownTimeoutError", - "MCPTool", - "McpClient", - "get_mcp_config_path", - "load_mcp_config", - "save_mcp_config", -] diff --git a/astrbot/_internal/protocols/_mcp/client.py b/astrbot/_internal/protocols/_mcp/client.py deleted file mode 100644 index 390a792bf..000000000 --- a/astrbot/_internal/protocols/_mcp/client.py +++ /dev/null @@ -1,411 +0,0 @@ -"""MCP client implementation.""" - -import asyncio -import logging -import os -import sys -from contextlib import AsyncExitStack -from datetime import timedelta -from typing import Any - -from tenacity import ( - before_sleep_log, - retry, - retry_if_exception_type, - stop_after_attempt, - wait_exponential, -) - -from astrbot._internal.abc._mcp.base_astrbot_mcp_client import ( - BaseAstrbotMcpClient, - McpServerConfig, - McpToolInfo, -) -from astrbot.core.utils.log_pipe import LogPipe - -logger = logging.getLogger("astrbot") -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(BaseAstrbotMcpClient): - def __init__(self) -> None: - self.session: mcp.ClientSession | None = None - self.exit_stack = AsyncExitStack() - self._old_exit_stacks: list[AsyncExitStack] = [] - self.name: str | None = None - self.active: bool = True - self.tools: list[mcp.Tool] = [] - self.server_errlogs: list[str] = [] - self.running_event = anyio.Event() - self.process_pid: int | None = None - self._mcp_server_config: McpServerConfig | None = None - self._server_name: str | None = None - self._reconnect_lock = anyio.Lock() - self._reconnecting: bool = False - - async def connect(self) -> None: - """Initialize the MCP client connection. - - Note: Actual server connections are made via connect_to_server(). - This method prepares the client for use. - """ - logger.debug("MCP client initialized.") - - @property - def connected(self) -> bool: - """True if MCP client has an active session.""" - return self.session is not None - - async def list_tools(self) -> list[McpToolInfo]: - """List all tools from connected MCP servers.""" - if not self.session: - return [] - result = await self.list_tools_and_save() - tools: list[McpToolInfo] = [ - { - "name": tool.name, - "description": tool.description or "", - "inputSchema": tool.inputSchema, - } - for tool in result.tools - ] - return tools - - async def call_tool( - self, - name: str, - arguments: dict[str, Any], - read_timeout_seconds: int = 60, - ) -> Any: - """Call a tool on the MCP server with reconnection support.""" - return await self.call_tool_with_reconnect( - tool_name=name, - arguments=arguments, - read_timeout_seconds=timedelta(seconds=read_timeout_seconds), - ) - - @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, config: McpServerConfig, 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: - config: Configuration for the MCP server. See https://modelcontextprotocol.io/quickstart/server - - """ - self._mcp_server_config = config - self._server_name = name - self.process_pid = None - cfg = _prepare_config(dict(config)) - - def logging_callback( - msg: str | mcp.types.LoggingMessageNotificationParams, - ) -> None: - 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": - 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, - ) - 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, - ), - ) - 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, - ) - 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, - ), - ) - else: - cfg = _prepare_stdio_env(cfg) - server_params = mcp.StdioServerParameters(**cfg) - - def callback(msg: str | mcp.types.LoggingMessageNotificationParams) -> None: - 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, - ), - ), - ) - self.process_pid = self._extract_stdio_process_pid(stdio_transport) - 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: - 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}...", - ) - if self.exit_stack: - self._old_exit_stacks.append(self.exit_stack) - self.session = None - self.exit_stack = AsyncExitStack() - 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...", - ) - await self._reconnect() - raise - - return await _call_with_retry() - - async def cleanup(self) -> None: - """Clean up resources including old exit stacks from reconnections""" - try: - await self.exit_stack.aclose() - except Exception as e: - logger.debug(f"Error closing current exit stack: {e}") - self._old_exit_stacks.clear() - self.running_event.set() - self.process_pid = None diff --git a/astrbot/_internal/protocols/_mcp/config.py b/astrbot/_internal/protocols/_mcp/config.py deleted file mode 100644 index 7faeb3477..000000000 --- a/astrbot/_internal/protocols/_mcp/config.py +++ /dev/null @@ -1,55 +0,0 @@ -"""MCP configuration management.""" - -import json -import os - -from astrbot.core.utils.astrbot_path import get_astrbot_data_path - -DEFAULT_MCP_CONFIG: dict[str, dict[str, object]] = {"mcpServers": {}} - - -def get_mcp_config_path() -> str: - """Get the path to the MCP configuration file.""" - data_dir = get_astrbot_data_path() - return os.path.join(data_dir, "mcp_server.json") - - -def load_mcp_config() -> dict: - """Load MCP configuration from file. - - Returns: - MCP configuration dict. If file doesn't exist, returns default config. - - """ - config_path = get_mcp_config_path() - if not os.path.exists(config_path): - # Create default config if not exists - os.makedirs(os.path.dirname(config_path), exist_ok=True) - with open(config_path, "w", encoding="utf-8") as f: - json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4) - return DEFAULT_MCP_CONFIG - - try: - with open(config_path, encoding="utf-8") as f: - return json.load(f) - except Exception: - return DEFAULT_MCP_CONFIG - - -def save_mcp_config(config: dict) -> bool: - """Save MCP configuration to file. - - Args: - config: MCP configuration dict to save. - - Returns: - True if successful, False otherwise. - - """ - config_path = get_mcp_config_path() - try: - with open(config_path, "w", encoding="utf-8") as f: - json.dump(config, f, ensure_ascii=False, indent=4) - return True - except Exception: - return False diff --git a/astrbot/_internal/protocols/_mcp/tool.py b/astrbot/_internal/protocols/_mcp/tool.py deleted file mode 100644 index b059ac6e4..000000000 --- a/astrbot/_internal/protocols/_mcp/tool.py +++ /dev/null @@ -1,49 +0,0 @@ -"""MCP tool wrapper.""" - -from __future__ import annotations - -from datetime import timedelta -from typing import TYPE_CHECKING, Any - -try: - import mcp as _mcp -except (ModuleNotFoundError, ImportError): - _mcp: Any = None - -from mcp.types import Tool as MCPTool_T - -from astrbot._internal.tools.base import FunctionTool - -if TYPE_CHECKING: - from astrbot._internal.protocols._mcp.client import McpClient - - -class MCPTool(FunctionTool): - """A function tool that calls an MCP service.""" - - def __init__( - self, - mcp_tool: MCPTool_T, - mcp_client: McpClient, - mcp_server_name: str, - **kwargs: Any, - ) -> 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, **kwargs: Any) -> Any: - """Call the MCP tool with the given arguments.""" - # Note: For actual usage, context.tool_call_timeout is needed - # but for simplicity we use a default timeout here - return await self.mcp_client.call_tool_with_reconnect( - tool_name=self.mcp_tool.name, - arguments=kwargs, - read_timeout_seconds=timedelta(seconds=60), - ) diff --git a/astrbot/_internal/runtime/__init__.py b/astrbot/_internal/runtime/__init__.py deleted file mode 100644 index 38d1843cd..000000000 --- a/astrbot/_internal/runtime/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from astrbot._internal.runtime.__main__ import bootstrap - -__all__ = ["bootstrap"] diff --git a/astrbot/_internal/runtime/__main__.py b/astrbot/_internal/runtime/__main__.py deleted file mode 100644 index 120195161..000000000 --- a/astrbot/_internal/runtime/__main__.py +++ /dev/null @@ -1,24 +0,0 @@ -from __future__ import annotations - -import anyio - -from astrbot._internal.abc.base_astrbot_gateway import BaseAstrbotGateway -from astrbot._internal.abc.base_astrbot_orchestrator import BaseAstrbotOrchestrator -from astrbot._internal.geteway.server import AstrbotGateway -from astrbot._internal.runtime.orchestrator import AstrbotOrchestrator - - -async def bootstrap(): - orchestrator: BaseAstrbotOrchestrator = AstrbotOrchestrator() - gw: BaseAstrbotGateway = AstrbotGateway(orchestrator) - - # anyio 的结构化并发 - async with anyio.create_task_group() as tg: - tg.start_soon(orchestrator.lsp.connect) # 启动 LSP client - tg.start_soon(orchestrator.mcp.connect) # 启动 MCP client - tg.start_soon(orchestrator.acp.connect) # 启动 ACP client - tg.start_soon(orchestrator.abp.connect) # 启动 ABP client - await anyio.sleep(0.5) - tg.start_soon(orchestrator.run_loop) # 启动编排器循环 - - tg.start_soon(gw.serve) # 面板后端服务 diff --git a/astrbot/_internal/runtime/orchestrator.py b/astrbot/_internal/runtime/orchestrator.py deleted file mode 100644 index e54249628..000000000 --- a/astrbot/_internal/runtime/orchestrator.py +++ /dev/null @@ -1,158 +0,0 @@ -"""AstrBot Orchestrator - core runtime that coordinates all protocol clients. - -The orchestrator manages the lifecycle of LSP, MCP, ACP, and ABP clients, -and runs the main event loop that dispatches messages between components. -""" - -from __future__ import annotations - -from typing import Any - -import anyio - -from astrbot import logger -from astrbot._internal.abc.base_astrbot_orchestrator import BaseAstrbotOrchestrator -from astrbot._internal.protocols._abp.client import AstrbotAbpClient -from astrbot._internal.protocols._acp.client import AstrbotAcpClient -from astrbot._internal.protocols._lsp.client import AstrbotLspClient -from astrbot._internal.protocols._mcp.client import McpClient -from astrbot._internal.stars import RuntimeStatusStar - -log = logger - - -class AstrbotOrchestrator(BaseAstrbotOrchestrator): - """Core runtime orchestrator for AstrBot. - - Manages: - - LSP client: Language Server Protocol for editor integrations - - MCP client: Model Context Protocol for external tool servers - - ACP client: AstrBot Communication Protocol for inter-service communication - - ABP client: AstrBot Protocol for built-in star (plugin) communication - """ - - def __init__(self) -> None: - # Initialize protocol clients (use concrete types for full method access) - self.lsp = AstrbotLspClient() - self.mcp = McpClient() - self.acp = AstrbotAcpClient() - self.abp = AstrbotAbpClient() - - self._running = False - self._stars: dict[str, Any] = {} - self._message_count: int = 0 - self._last_activity_timestamp: float | None = None - - # Auto-register RuntimeStatusStar - self._runtime_status_star = RuntimeStatusStar() - self._runtime_status_star.set_orchestrator(self) - self._stars["runtime-status-star"] = self._runtime_status_star - self.abp.register_star("runtime-status-star", self._runtime_status_star) - - log.debug("AstrbotOrchestrator initialized.") - - async def start(self) -> None: - """Initialize all protocol clients. - - Calls connect() on all protocol clients to prepare them for use. - """ - log.info("Starting AstrbotOrchestrator...") - - await self.lsp.connect() - await self.mcp.connect() - await self.acp.connect() - await self.abp.connect() - - self._running = True - log.info("AstrbotOrchestrator started.") - - async def run_loop(self) -> None: - """Main orchestrator event loop. - - This loop runs continuously, handling: - - Periodic health checks of protocol clients - - Message routing between protocols - - Star (plugin) lifecycle management - """ - self._running = True - log.info("AstrbotOrchestrator run loop started.") - - stop_event = anyio.Event() - - def set_stop() -> None: - stop_event.set() - - # Store the callback for cleanup - self._stop_callback = set_stop - - try: - while self._running: - # TODO: Periodic tasks: - # - Check LSP server health - # - Check MCP session status - # - Check ACP client connections - # - Process any pending star notifications - - # Wait for 5 seconds or until shutdown is called - with anyio.move_on_after(5): - await stop_event.wait() - - except anyio.get_cancelled_exc_class(): - log.info("Orchestrator run loop cancelled.") - finally: - self._running = False - self._stop_callback = None - log.info("AstrbotOrchestrator run loop stopped.") - - async def register_star(self, name: str, star_instance: Any) -> None: - """Register a star (plugin) with the orchestrator. - - Args: - name: Unique name for the star - star_instance: Star plugin instance - - """ - self._stars[name] = star_instance - self.abp.register_star(name, star_instance) - log.info(f"Star '{name}' registered.") - - async def unregister_star(self, name: str) -> None: - """Unregister a star (plugin) from the orchestrator. - - Args: - name: Name of the star to unregister - - """ - self._stars.pop(name, None) - self.abp.unregister_star(name) - log.info(f"Star '{name}' unregistered.") - - async def get_star(self, name: str) -> Any | None: - """Get a registered star by name.""" - return self._stars.get(name) - - async def list_stars(self) -> list[str]: - """List all registered star names.""" - return list(self._stars.keys()) - - def record_activity(self) -> None: - """Record a message activity for stats tracking.""" - self._message_count += 1 - import time - - self._last_activity_timestamp = time.time() - - async def shutdown(self) -> None: - """Shutdown the orchestrator and all protocol clients.""" - log.info("Shutting down AstrbotOrchestrator...") - self._running = False - - # Shutdown all protocol clients - await self.lsp.shutdown() - await self.acp.shutdown() - await self.abp.shutdown() - - # MCP cleanup - await self.mcp.cleanup() - - log.info("AstrbotOrchestrator shut down.") diff --git a/astrbot/_internal/skills/__init__.py b/astrbot/_internal/skills/__init__.py deleted file mode 100644 index e36af0aed..000000000 --- a/astrbot/_internal/skills/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Internal skills module - re-exports from core.skills.skill_manager.""" - -from astrbot.core.skills.skill_manager import ( - SkillInfo, - SkillManager, - build_skills_prompt, -) - -__all__ = [ - "SkillInfo", - "SkillManager", - "build_skills_prompt", -] diff --git a/astrbot/_internal/stars/__init__.py b/astrbot/_internal/stars/__init__.py deleted file mode 100644 index 1cc5c2193..000000000 --- a/astrbot/_internal/stars/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Stars (built-in plugins) for AstrBot runtime.""" - -from astrbot._internal.stars.runtime_status_star import RuntimeStatusStar - -__all__ = ["RuntimeStatusStar"] diff --git a/astrbot/_internal/stars/runtime_status_star.py b/astrbot/_internal/stars/runtime_status_star.py deleted file mode 100644 index ed33ba8cb..000000000 --- a/astrbot/_internal/stars/runtime_status_star.py +++ /dev/null @@ -1,125 +0,0 @@ -"""RuntimeStatusStar - ABP plugin that exposes core runtime internal state. - -This star provides tools for querying: -- Runtime status (running state, uptime) -- Protocol client status (LSP, MCP, ACP, ABP) -- Registered stars registry -- Message counts and metrics -""" - -from __future__ import annotations - -import time -from dataclasses import dataclass, field -from datetime import datetime, timezone -from typing import Any - - -@dataclass -class RuntimeStatusStar: - """ABP star that exposes core runtime internal state as callable tools. - - Tools provided: - - get_runtime_status: Returns running state and uptime - - get_protocol_status: Returns LSP, MCP, ACP, ABP status - - get_star_registry: Returns registered star names - - get_stats: Returns message counts and metrics - """ - - name: str = "runtime-status-star" - description: str = "ABP plugin that exposes core runtime internal state" - - _start_time: float = field(default_factory=time.time, init=False) - _orchestrator: Any = field(default=None, init=False) - - def set_orchestrator(self, orchestrator: Any) -> None: - """Set the orchestrator reference for status queries.""" - self._orchestrator = orchestrator - - async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any: - """Handle tool calls from ABP client. - - Args: - tool_name: Name of the tool to call - arguments: Tool arguments - - Returns: - Tool result - - """ - if tool_name == "get_runtime_status": - return self._get_runtime_status() - if tool_name == "get_protocol_status": - return await self._get_protocol_status() - if tool_name == "get_star_registry": - return await self._get_star_registry() - if tool_name == "get_stats": - return self._get_stats() - raise ValueError(f"Unknown tool: {tool_name}") - - def _get_runtime_status(self) -> dict[str, Any]: - """Get overall runtime state.""" - running = ( - getattr(self._orchestrator, "running", False) - if self._orchestrator - else False - ) - uptime_seconds = time.time() - self._start_time - return { - "running": running, - "uptime_seconds": uptime_seconds, - } - - async def _get_protocol_status(self) -> dict[str, Any]: - """Get status of each protocol client.""" - if not self._orchestrator: - return { - "lsp": {"connected": False, "name": "lsp-client"}, - "mcp": {"connected": False, "name": "mcp-client"}, - "acp": {"connected": False, "name": "acp-client"}, - "abp": {"connected": False, "name": "abp-client"}, - } - - return { - "lsp": { - "connected": getattr(self._orchestrator.lsp, "connected", False), - "name": "lsp-client", - }, - "mcp": { - "connected": getattr(self._orchestrator.mcp, "connected", False), - "name": "mcp-client", - }, - "acp": { - "connected": getattr(self._orchestrator.acp, "connected", False), - "name": "acp-client", - }, - "abp": { - "connected": getattr(self._orchestrator.abp, "connected", False), - "name": "abp-client", - }, - } - - async def _get_star_registry(self) -> dict[str, Any]: - """Get list of registered stars.""" - if not self._orchestrator: - return {"stars": []} - - stars = await self._orchestrator.list_stars() - return {"stars": stars} - - def _get_stats(self) -> dict[str, Any]: - """Get message counts and metrics.""" - result: dict[str, Any] = { - "uptime_seconds": time.time() - self._start_time, - } - if self._orchestrator: - result["total_messages"] = getattr(self._orchestrator, "_message_count", 0) - last_ts = getattr(self._orchestrator, "_last_activity_timestamp", None) - if last_ts is not None: - result["last_activity"] = datetime.fromtimestamp( - last_ts, - tz=timezone.utc, - ).isoformat() - else: - result["last_activity"] = None - return result diff --git a/astrbot/_internal/tools/__init__.py b/astrbot/_internal/tools/__init__.py deleted file mode 100644 index 434182911..000000000 --- a/astrbot/_internal/tools/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Internal tools module for AstrBot runtime.""" - -from .base import FunctionTool, ToolSet - -__all__ = ["FunctionTool", "ToolSet"] diff --git a/astrbot/_internal/tools/base.py b/astrbot/_internal/tools/base.py deleted file mode 100644 index a2f019adb..000000000 --- a/astrbot/_internal/tools/base.py +++ /dev/null @@ -1,338 +0,0 @@ -"""Base tool classes for AstrBot internal runtime. - -This module provides the FunctionTool base class used by MCP tools -in the new internal architecture. -""" - -import copy -from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator -from dataclasses import dataclass, field -from typing import Any - -from pydantic import model_validator - -ParametersType = dict[str, Any] - - -@dataclass -class ToolSchema: - """A class representing the schema of a tool for function calling.""" - - name: str - """The name of the tool.""" - - description: str - """The description of the tool.""" - - parameters: ParametersType = field(default_factory=dict) - """The parameters of the tool, in JSON Schema format.""" - - active: bool = True - """Whether the tool is active.""" - - @model_validator(mode="after") - def validate_parameters(self) -> "ToolSchema": - """Validate the parameters JSON schema.""" - import jsonschema - - jsonschema.validate( - self.parameters, - jsonschema.Draft202012Validator.META_SCHEMA, - ) - return self - - -@dataclass -class FunctionTool(ToolSchema): - """A callable tool, for function calling.""" - - handler: Callable[..., Awaitable[str | None] | AsyncGenerator[Any, 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 - """ - - 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 = "mcp" - """ - 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, **kwargs: Any) -> Any: - """Run the tool with the given arguments. The handler field has priority.""" - raise NotImplementedError( - "FunctionTool.call() must be implemented by subclasses or set a handler.", - ) - - -class ToolSet: - """A collection of FunctionTools grouped under a namespace. - - ToolSets allow organizing related tools together. The LLM sees tools - as "namespace/tool_name" when calling. - """ - - def __init__(self, namespace: str, tools: list[ToolSchema] | None = None) -> None: - self.namespace = namespace - self._tools: dict[str, ToolSchema] = {} - if tools: - for tool in tools: - self.add(tool) - - def add(self, tool: ToolSchema) -> None: - """Add a tool to the set.""" - existing = self._tools.get(tool.name) - if existing is not None and not getattr(tool, "active", True) and getattr(existing, "active", True): - return # Don't overwrite active with inactive - self._tools[tool.name] = tool - - def add_tool(self, tool: ToolSchema) -> None: - """Add a tool to the set (alias for add()).""" - self.add(tool) - - def remove(self, name: str) -> ToolSchema | None: - """Remove and return a tool by name.""" - return self._tools.pop(name, None) - - def remove_tool(self, name: str) -> None: - """Remove a tool by its name.""" - self._tools.pop(name, None) - - def get(self, name: str) -> ToolSchema | None: - """Get a tool by name.""" - return self._tools.get(name) - - def get_tool(self, name: str) -> ToolSchema | None: - """Get a tool by name (alias for get).""" - return self.get(name) - - def list_tools(self) -> list[ToolSchema]: - """List all tools in this set.""" - return list(self._tools.values()) - - def __iter__(self) -> Iterator[ToolSchema]: - return iter(self._tools.values()) - - def __len__(self) -> int: - return len(self._tools) - - def __bool__(self) -> bool: - return bool(self._tools) - - def __repr__(self) -> str: - return f"ToolSet(namespace={self.namespace!r}, tools={self.list_tools()!r})" - - def __str__(self) -> str: - return f"ToolSet({self.namespace}, {len(self)} tools)" - - def names(self) -> list[str]: - """Get names of all tools in this set.""" - return [tool.name for tool in self.tools] - - def empty(self) -> bool: - """Check if the tool set is empty.""" - return len(self) == 0 - - def merge(self, other: "ToolSet") -> None: - """Merge another ToolSet into this one.""" - for tool in other.tools: - self.add(tool) - - def normalize(self) -> None: - """Sort tools by name for deterministic serialization.""" - self._tools = dict(sorted(self._tools.items(), key=lambda x: x[0])) - - def get_light_tool_set(self) -> "ToolSet": - """Return a light tool set with only name/description.""" - light_tools: list[ToolSchema] = [] - for tool in self.tools: - if hasattr(tool, "active") and not tool.active: - continue - light_tools.append( - FunctionTool( - name=tool.name, - description=tool.description, - parameters={"type": "object", "properties": {}}, - handler=None, - ), - ) - return ToolSet("default", light_tools) - - def get_param_only_tool_set(self) -> "ToolSet": - """Return a tool set with name/parameters only (no description).""" - param_tools: list[ToolSchema] = [] - 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("default", param_tools) - - @property - def tools(self) -> list[ToolSchema]: - """List all tools in this set.""" - return list(self._tools.values()) - - def openai_schema( - self, - omit_empty_parameter_field: bool = False, - ) -> list[dict[str, Any]]: - """Convert tools to OpenAI API function calling schema format.""" - result: list[dict[str, Any]] = [] - for tool in self._tools.values(): - func_dict: dict[str, Any] = {"name": tool.name} - if tool.description: - func_dict["description"] = tool.description - - if tool.parameters is not None: - if ( - tool.parameters.get("properties") - ) or not omit_empty_parameter_field: - func_dict["parameters"] = tool.parameters - - func_def: dict[str, Any] = { - "type": "function", - "function": func_dict, - } - - 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: dict[str, Any] = {"type": "object"} - if tool.parameters: - input_schema["properties"] = tool.parameters.get("properties", {}) - input_schema["required"] = tool.parameters.get("required", []) - tool_def: dict[str, Any] = {"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: - 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 = {} - origin_type = schema.get("type") - target_type = origin_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"] - 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: - result["items"] = {"type": "string"} - - return result - - tools_list = [] - 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_list.append(d) - - declarations: dict[str, Any] = {} - if tools_list: - declarations["function_declarations"] = tools_list - return declarations - - def get_func_desc_openai_style(self, omit_empty_parameter_field: bool = False): - """Get tools in OpenAI function calling style (deprecated).""" - return self.openai_schema(omit_empty_parameter_field) - - def get_func_desc_anthropic_style(self): - """Get tools in Anthropic style (deprecated).""" - return self.anthropic_schema() - - def get_func_desc_google_genai_style(self): - """Get tools in Google GenAI style (deprecated).""" - return self.google_schema() diff --git a/astrbot/_internal/tools/builtin.py b/astrbot/_internal/tools/builtin.py deleted file mode 100644 index 791dd94af..000000000 --- a/astrbot/_internal/tools/builtin.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Builtin tools for AstrBot - re-exports from core.tools for backward compatibility. - -This module re-exports the builtin tools (cron, send_message, kb_query) from -the deprecated core.tools module for backward compatibility. - -TODO: These tools should be fully migrated to _internal and core.tools -should be removed once all consumers update their imports. -""" - -from __future__ import annotations - -# Re-export cron tools -from astrbot.core.tools.cron_tools import ( - CREATE_CRON_JOB_TOOL, - DELETE_CRON_JOB_TOOL, - LIST_CRON_JOBS_TOOL, - CreateActiveCronTool, - DeleteCronJobTool, - ListCronJobsTool, -) - -# Re-export knowledge_base_query tool -from astrbot.core.tools.kb_query import ( - KNOWLEDGE_BASE_QUERY_TOOL, - KnowledgeBaseQueryTool, -) - -# Re-export send_message tool -from astrbot.core.tools.send_message import ( - SEND_MESSAGE_TO_USER_TOOL, - SendMessageToUserTool, -) - -__all__ = [ - # Cron tools - "CREATE_CRON_JOB_TOOL", - "DELETE_CRON_JOB_TOOL", - "KNOWLEDGE_BASE_QUERY_TOOL", - "LIST_CRON_JOBS_TOOL", - "SEND_MESSAGE_TO_USER_TOOL", - # Classes - "CreateActiveCronTool", - "DeleteCronJobTool", - "KnowledgeBaseQueryTool", - "ListCronJobsTool", - "SendMessageToUserTool", -] diff --git a/astrbot/_internal/tools/registry.py b/astrbot/_internal/tools/registry.py deleted file mode 100644 index f97917bb0..000000000 --- a/astrbot/_internal/tools/registry.py +++ /dev/null @@ -1,1015 +0,0 @@ -"""Tools registry for AstrBot internal runtime. - -This module provides the canonical FunctionToolManager implementation -for tool registration and MCP server management. -""" - -from __future__ import annotations - -import asyncio -import copy -import json -import os -import threading -import urllib.parse -from collections.abc import Mapping -from dataclasses import dataclass -from types import MappingProxyType -from typing import Any - -import aiohttp - -from astrbot import logger - -# Re-export MCP-related constants and exceptions from protocols -from astrbot._internal.protocols._mcp import ( - DEFAULT_MCP_CONFIG as _DEFAULT_MCP_CONFIG, -) -from astrbot._internal.protocols._mcp import ( - MCPAllServicesFailedError, - MCPInitError, - MCPInitSummary, - MCPInitTimeoutError, - MCPShutdownTimeoutError, -) -from astrbot._internal.protocols._mcp.client import ( - McpClient, - _quick_test_mcp_connection, -) -from astrbot._internal.protocols._mcp.tool import MCPTool -from astrbot._internal.tools.base import FunctionTool, ToolSchema, ToolSet - -__all__ = [ - "DEFAULT_MCP_CONFIG", - "ENABLE_MCP_TIMEOUT_ENV", - "FuncCall", - "FunctionTool", - "FunctionToolManager", - "MCPAllServicesFailedError", - "MCPInitError", - "MCPInitSummary", - "MCPInitTimeoutError", - "MCPShutdownTimeoutError", - "ToolSet", -] - - -DEFAULT_MCP_CONFIG: dict[str, Any] = _DEFAULT_MCP_CONFIG -ENABLE_MCP_TIMEOUT_ENV = "ASTRBOT_MCP_TIMEOUT_ENABLED" -MCP_INIT_TIMEOUT_ENV = "ASTRBOT_MCP_INIT_TIMEOUT" - -DEFAULT_MCP_INIT_TIMEOUT_SECONDS = 180.0 -DEFAULT_ENABLE_MCP_TIMEOUT_SECONDS = 180.0 -MAX_MCP_TIMEOUT_SECONDS = 300.0 - - -@dataclass -class _MCPServerRuntime: - """Runtime state for a single MCP server.""" - - name: str - client: McpClient - shutdown_event: asyncio.Event - lifecycle_task: asyncio.Task[None] - - -class _MCPClientDictView(Mapping[str, McpClient]): - """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) -> McpClient: - 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"environment variable {env_name}" - if timeout is None: - timeout = os.getenv(env_name, str(default)) - else: - source = "explicit timeout argument" - - try: - timeout_value = float(timeout) - except (TypeError, ValueError): - logger.warning( - f"Timeout configuration ({source})={timeout!r} is invalid, using default {default:g} seconds.", - ) - return default - - if timeout_value <= 0: - logger.warning( - f"Timeout configuration ({source})={timeout_value:g} must be greater than 0, using default {default:g} seconds.", - ) - return default - - if timeout_value > MAX_MCP_TIMEOUT_SECONDS: - logger.warning( - f"Timeout configuration ({source})={timeout_value:g} is too large, limited to maximum " - f"{MAX_MCP_TIMEOUT_SECONDS:g} seconds to avoid long waits.", - ) - return MAX_MCP_TIMEOUT_SECONDS - - return timeout_value - - -# Alias for backward compatibility -FuncTool = FunctionTool - - -def _get_sp() -> Any: - """Get the SharedPreferences instance, avoiding circular imports.""" - from astrbot.core import sp - - return sp - - -def _get_astrbot_data_path() -> str: - """Get the AstrBot data directory path.""" - from astrbot.core.utils.astrbot_path import get_astrbot_data_path - - return get_astrbot_data_path() - - -# Builtin tools stubs - these functions don't exist in the new architecture -def _ensure_builtin_tools_loaded() -> None: - """No-op in new architecture.""" - - -def _get_builtin_tool_class(name: str) -> type[FunctionTool] | None: - """No-op in new architecture.""" - return None - - -def _get_builtin_tool_name(tool_cls: type[FunctionTool]) -> str | None: - """No-op in new architecture.""" - return None - - -def _iter_builtin_tool_classes() -> list[type[FunctionTool]]: - """No-op in new architecture.""" - return [] - - -class FunctionToolManager: - """Central registry for all function tools.""" - - def __init__(self) -> None: - self.func_list: list[FuncTool] = [] - """All tools including MCP tools and plugin tools, except AstrBot builtin tools.""" - self.builtin_func_list: dict[type[FuncTool], FuncTool] = {} - """All AstrBot builtin tools, keyed by their class. Values are instantiated tool objects, created on demand.""" - - self._mcp_server_runtime: dict[str, _MCPServerRuntime] = {} - """MCP runtime metadata, keyed by server name. Updated atomically on MCP lifecycle changes.""" - 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, McpClient]: - """Read-only compatibility view for external callers that still read mcp_client_dict. - - Note: Mutating this mapping is unsupported and will raise TypeError. - """ - return self._mcp_client_dict_view - - @property - def mcp_server_runtime_view(self) -> Mapping[str, _MCPServerRuntime]: - """Read-only view of MCP runtime metadata for external callers.""" - return self._mcp_server_runtime_view - - @property - def mcp_server_runtime(self) -> Mapping[str, _MCPServerRuntime]: - """Backward-compatible read-only view (deprecated). Do not mutate. - - Note: Mutations are not supported and will raise TypeError. - """ - 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: Any, - ) -> FuncTool: - """Create a FuncTool from a specification.""" - params = { - "type": "object", # hard-coded here - "properties": {}, - } - for param in func_args: - p = copy.deepcopy(param) - p.pop("name", None) - params["properties"][param["name"]] = p - return FuncTool( - name=name, - parameters=params, - description=desc, - handler=handler, - ) - - def add_func( - self, - name: str, - func_args: list, - desc: str, - handler: Any, - ) -> None: - """Add a function tool. - - Args: - name: Function name - func_args: Function arguments list, format: [{"type": "string", "name": "arg_name", "description": "arg_description"}, ...] - desc: Function description - handler: Handler function - - """ - # Check if the tool has been added before - 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"Added function tool: {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 remove(self, name: str) -> bool: - """Remove a tool by name. Returns True if found.""" - for i, f in enumerate(self.func_list): - if f.name == name: - self.func_list.pop(i) - return True - return False - - def add(self, tool: ToolSchema) -> None: - """Add a tool to the registry.""" - self.func_list.append(tool) - - def get_func(self, name: str | type[FuncTool]) -> FuncTool | None: - # Prefer returning active tools (later loaded ones override earlier ones) - # Using getattr(..., True) to match ToolSet.add_tool behavior: tools without active attribute are considered active - if isinstance(name, str): - for f in reversed(self.func_list): - if f.name == name and getattr(f, "active", True): - return f - # Fallback: get the last tool with the same name - for f in reversed(self.func_list): - if f.name == name: - return f - # Try builtin tools - try: - builtin_tool = self.get_builtin_tool(name) - if builtin_tool is not None and getattr(builtin_tool, "active", True): - return builtin_tool - return builtin_tool - except (KeyError, TypeError): - return None - return None - - def get_builtin_tool(self, tool: str | type[FuncTool]) -> FuncTool: - """Get a builtin tool by name or class.""" - _ensure_builtin_tools_loaded() - - if isinstance(tool, str): - tool_cls = _get_builtin_tool_class(tool) - if tool_cls is None: - raise KeyError(f"Builtin tool {tool} is not registered.") - elif isinstance(tool, type) and issubclass(tool, FunctionTool): - tool_cls = tool - if _get_builtin_tool_name(tool_cls) is None: - raise KeyError( - f"Builtin tool class {tool_cls.__module__}.{tool_cls.__name__} is not registered.", - ) - else: - raise TypeError("tool must be a builtin tool name or FunctionTool class.") - - cached_tool = self.builtin_func_list.get(tool_cls) - if cached_tool is not None: - return cached_tool - - builtin_tool = tool_cls() # type: ignore - self.builtin_func_list[tool_cls] = builtin_tool - return builtin_tool - - def iter_builtin_tools(self) -> list[FuncTool]: - """Iterate over all builtin tools.""" - _ensure_builtin_tools_loaded() - return [ - self.get_builtin_tool(tool_cls) for tool_cls in _iter_builtin_tool_classes() - ] - - def is_builtin_tool(self, name: str) -> bool: - """Check if a tool name is a builtin tool.""" - _ensure_builtin_tools_loaded() - return _get_builtin_tool_class(name) is not None - - def get_full_tool_set(self) -> ToolSet: - """Get the full tool set. - - Uses ToolSet.add_tool for population. For tools with the same name, - deduplication rules are: - - Prefer tools with active=True - - When active status is the same, later loaded tools override earlier ones - - Therefore, a later loaded inactive tool will not override an already active tool; - MCP tools can still override disabled builtin tools when needed. - """ - tool_set = ToolSet(namespace="default") - for tool in self.func_list: - tool_set.add_tool(tool) - return tool_set - - def register_internal_tools(self) -> None: - """Register built-in computer tools (shell, python, browser, neo).""" - # Import here to avoid circular imports - try: - from astrbot.core.computer.computer_tool_provider import get_all_tools - - for tool in get_all_tools(): - if self.get_func(tool.name) is None: - self.add(tool) # type: ignore[arg-type] - except ImportError: - logger.debug("Computer tools not available") - - @staticmethod - def _log_safe_mcp_debug_config(cfg: dict) -> None: - """Log sanitized MCP debug config info.""" - # Only log sanitized summary to avoid leaking sensitive info in command/args/url - 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" Command executable: {executable}, argument count: {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" Host: {scheme}://{host}{port}") - - async def init_mcp_clients( - self, - raise_on_all_failed: bool = False, - ) -> MCPInitSummary: - """Initialize MCP clients from mcp_server.json config file. - - The config file format: - ``` - { - "mcpServers": { - "weather": { - "command": "uv", - "args": [ - "--directory", - "/ABSOLUTE/PATH/TO/PARENT/FOLDER/weather", - "run", - "weather.py" - ] - } - } - ... - } - ``` - - Timeout behavior: - - Initialization timeout uses environment variable ASTRBOT_MCP_INIT_TIMEOUT or default. - - Dynamic enable timeout uses ASTRBOT_MCP_ENABLE_TIMEOUT (independent from init timeout). - """ - data_dir = _get_astrbot_data_path() - - mcp_json_file = os.path.join(data_dir, "mcp_server.json") - if not os.path.exists(mcp_json_file): - # Config file doesn't exist, create default - with open(mcp_json_file, "w", encoding="utf-8") as f: - json.dump(DEFAULT_MCP_CONFIG, f, ensure_ascii=False, indent=4) - logger.info(f"MCP config file not found, created default config at {mcp_json_file}") - return MCPInitSummary(total=0, success=0, failed=[]) - - with open(mcp_json_file, encoding="utf-8") as f: - mcp_server_json_obj: dict[str, dict] = json.load(f)["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"Waiting for {len(active_configs)} MCP services to initialize...") - - init_tasks = [ - asyncio.create_task( - self._start_mcp_server( - name=name, - cfg=cfg, - shutdown_event=shutdown_event, - 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, - timeout: float, - ) -> None: - """Initialize MCP server with timeout and register task/event together. - - This method is idempotent. If the server is already running, the existing - runtime is kept and the new config is ignored. - """ - 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={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: McpClient | None = None - try: - mcp_client = await asyncio.wait_for( - self._init_mcp_client(name, cfg), - timeout=timeout, - ) - except asyncio.TimeoutError as exc: - raise MCPInitTimeoutError( - f"Connected to MCP server {name} timeout ({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], - timeout: float, - *, - strict: bool = True, - ) -> list[str]: - """Shutdown runtimes and wait for lifecycle tasks to complete.""" - 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=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, timeout) - logger.warning( - "MCP server shutdown timeout (%s seconds), the following servers were not fully closed: %s", - f"{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: McpClient, - name: str, - ) -> None: - """Safely clean up a single MCP client, avoiding cleanup exceptions interrupting main flow.""" - 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) -> McpClient: - """Initialize a single MCP client.""" - 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 any previous tools from this MCP server - self.func_list = [ - f - for f in self.func_list - if not (isinstance(f, MCPTool) and f.mcp_server_name == name) - ] - - # Convert MCP tools to FuncTool and add to func_list - 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: - """Shut down and clean up an MCP client.""" - async with self._runtime_lock: - runtime = self._mcp_server_runtime.get(name) - if runtime: - client = runtime.client - # Close MCP connection - await self._cleanup_mcp_client_safely(client, name) - # Remove associated FuncTools - self.func_list = [ - f - for f in self.func_list - if not (isinstance(f, 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 - - # Runtime missing but stale tools may still exist after failed flows. - self.func_list = [ - f - for f in self.func_list - if not (isinstance(f, 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, - ) -> tuple[bool, str]: - """Test MCP server connection. - - Returns: - Tuple of (success, message) - - """ - success, error_msg = await _quick_test_mcp_connection(config) - if not success: - return False, 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] - return True, f"Connection successful, tools: {tool_names}" - except Exception as e: - return False, str(e) - finally: - logger.debug("Cleaning up MCP client after testing connection.") - await mcp_client.cleanup() - - async def enable_mcp_server( - self, - name: str, - config: dict, - shutdown_event: asyncio.Event | None = None, - timeout: float | str | None = None, - ) -> None: - """Enable a new MCP server and initialize it. - - Args: - name: The name of the MCP server. - config: Configuration for the MCP server. - shutdown_event: Event to signal when the MCP client should shut down. - timeout: Timeout in seconds for initialization. - Uses ASTRBOT_MCP_ENABLE_TIMEOUT by default (separate from init timeout). - - Raises: - MCPInitTimeoutError: If initialization does not complete within timeout. - Exception: If there is an error during initialization. - - """ - if timeout is None: - timeout_value = self._enable_timeout_default - else: - timeout_value = _resolve_timeout( - timeout=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, - timeout=timeout_value, - ) - - async def disable_mcp_server( - self, - name: str | None = None, - timeout: float = 10, - ) -> None: - """Disable an MCP server by its name. - - Args: - name: The name of the MCP server to disable. If None, ALL MCP servers will be disabled. - timeout: Timeout in seconds for shutdown. - - Raises: - MCPShutdownTimeoutError: If shutdown does not complete within timeout. - Only raised when disabling a specific server (name is not 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], timeout, strict=True) - else: - async with self._runtime_lock: - runtimes = list(self._mcp_server_runtime.values()) - await self._shutdown_runtimes(runtimes, 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( - "Detected different MCP initialization timeout and dynamic enable timeout: " - "initialization uses %s seconds, dynamic enable uses %s seconds. " - "Set them to the same value for consistency.", - f"{init_timeout:g}", - f"{enable_timeout:g}", - ) - self._timeout_mismatch_warned = True - - def get_func_desc_openai_style(self, omit_empty_parameter_field: bool = False) -> list: - """Get tools in OpenAI API style for all active tools.""" - tools = [f for f in self.func_list if getattr(f, "active", True)] - toolset = ToolSet(tools) - return toolset.openai_schema( - omit_empty_parameter_field=omit_empty_parameter_field, - ) - - def get_func_desc_anthropic_style(self) -> list: - """Get tools in Anthropic API style for all active tools.""" - tools = [f for f in self.func_list if getattr(f, "active", True)] - toolset = ToolSet(tools) - return toolset.anthropic_schema() - - def get_func_desc_google_genai_style(self) -> dict: - """Get tools in Google GenAI API style for all active tools.""" - tools = [f for f in self.func_list if getattr(f, "active", True)] - toolset = ToolSet(tools) - return toolset.google_schema() - - def deactivate_llm_tool(self, name: str) -> bool: - """Deactivate a registered function tool. - - Returns: - False if not found - - """ - func_tool = self.get_func(name) - if func_tool is not None: - func_tool.active = False - - sp = _get_sp() - 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 | None = None) -> bool: - """Activate a registered function tool. - - Args: - name: Tool name - star_map: Optional star_map for plugin dependency checking (ignored in new architecture) - - """ - func_tool = self.get_func(name) - if func_tool is not None: - if star_map is not None and func_tool.handler_module_path in star_map: - if not star_map[func_tool.handler_module_path].activated: - raise ValueError( - f"The plugin {star_map[func_tool.handler_module_path].name} " - f"that this function tool belongs to has been disabled. " - f"Please enable the plugin in the management panel first.", - ) - - func_tool.active = True - - sp = _get_sp() - 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: - """Get the MCP config file path.""" - data_dir = _get_astrbot_data_path() - return os.path.join(data_dir, "mcp_server.json") - - def load_mcp_config(self) -> dict[str, Any]: - """Load MCP configuration from file.""" - if not os.path.exists(self.mcp_config_path): - # Config file doesn't exist, create default - 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 - - try: - with open(self.mcp_config_path, encoding="utf-8") as f: - return json.load(f) - except Exception as e: - logger.error(f"Failed to load MCP config: {e}") - return DEFAULT_MCP_CONFIG - - 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"Failed to save MCP config: {e}") - return False - - async def sync_modelscope_mcp_servers(self, access_token: str) -> None: - """Sync MCP servers from ModelScope platform.""" - base_url = "https://www.modelscope.cn/openapi/v1" - url = f"{base_url}/mcp/servers/operational" - headers = { - "Authorization": f"Bearer {access_token.strip()}", - "Content-Type": "application/json", - } - - try: - async with aiohttp.ClientSession() as session: - async with session.get(url, headers=headers) as response: - if response.status == 200: - data = await response.json() - mcp_server_list = data.get("data", {}).get( - "mcp_server_list", - [], - ) - local_mcp_config = self.load_mcp_config() - - synced_count = 0 - for server in mcp_server_list: - server_name = server["name"] - operational_urls = server.get("operational_urls", []) - if not operational_urls: - continue - url_info = operational_urls[0] - server_url = url_info.get("url") - if not server_url: - continue - # Add to config (same name will override) - local_mcp_config["mcpServers"][server_name] = { - "url": server_url, - "transport": "sse", - "active": True, - "provider": "modelscope", - } - synced_count += 1 - - if synced_count > 0: - self.save_mcp_config(local_mcp_config) - tasks = [] - for server in mcp_server_list: - name = server["name"] - tasks.append( - self.enable_mcp_server( - name=name, - config=local_mcp_config["mcpServers"][name], - ), - ) - await asyncio.gather(*tasks) - logger.info( - f"Synced {synced_count} MCP servers from ModelScope", - ) - else: - logger.warning("No available ModelScope MCP servers found") - else: - raise Exception( - f"ModelScope API request failed: HTTP {response.status}", - ) - - except aiohttp.ClientError as e: - raise Exception(f"Network connection error: {e!s}") - except Exception as e: - raise Exception(f"Error syncing ModelScope MCP servers: {e!s}") - - def __str__(self) -> str: - return str(self.func_list) - - def __repr__(self) -> str: - return str(self.func_list) - - -# Alias for backward compatibility -FuncCall = FunctionToolManager diff --git a/astrbot/api/__init__.py b/astrbot/api/__init__.py index e9f27c690..5c50690e3 100644 --- a/astrbot/api/__init__.py +++ b/astrbot/api/__init__.py @@ -12,9 +12,6 @@ Modules: from astrbot import logger -# Tool API -from astrbot._internal.tools.base import FunctionTool, ToolSet - # MCP API from astrbot.api.mcp import ( MCPClient, diff --git a/astrbot/builtin_stars/builtin_commands/commands/__init__.py b/astrbot/builtin_stars/builtin_commands/commands/__init__.py index c44c55e84..8b0019f77 100644 --- a/astrbot/builtin_stars/builtin_commands/commands/__init__.py +++ b/astrbot/builtin_stars/builtin_commands/commands/__init__.py @@ -3,6 +3,7 @@ from .admin import AdminCommands from .conversation import ConversationCommands from .help import HelpCommand +from .provider import ProviderCommands from .setunset import SetUnsetCommands from .sid import SIDCommand @@ -10,6 +11,8 @@ __all__ = [ "AdminCommands", "ConversationCommands", "HelpCommand", + "ProviderCommands", + "SetUnsetCommands", "SIDCommand", "SetUnsetCommands", ] diff --git a/astrbot/builtin_stars/builtin_commands/commands/provider.py b/astrbot/builtin_stars/builtin_commands/commands/provider.py index 74887a9d4..971d6ca8a 100644 --- a/astrbot/builtin_stars/builtin_commands/commands/provider.py +++ b/astrbot/builtin_stars/builtin_commands/commands/provider.py @@ -1,301 +1,25 @@ from __future__ import annotations import asyncio -import time -from collections.abc import Sequence -from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal, Protocol, TypeAlias, TypedDict from astrbot import logger from astrbot.api import star from astrbot.api.event import AstrMessageEvent, MessageEventResult -from astrbot.core.provider.entities import ProviderMeta, ProviderType +from astrbot.core.provider.entities import ProviderType from astrbot.core.utils.error_redaction import safe_error -if TYPE_CHECKING: - from astrbot.core.provider.provider import Provider - - -MODEL_LIST_CACHE_TTL_SECONDS_DEFAULT = 30.0 -MODEL_LOOKUP_MAX_CONCURRENCY_DEFAULT = 4 -MODEL_LOOKUP_MAX_CONCURRENCY_UPPER_BOUND = 16 -MODEL_LIST_CACHE_TTL_KEY = "model_list_cache_ttl_seconds" -MODEL_LOOKUP_MAX_CONCURRENCY_KEY = "model_lookup_max_concurrency" -MODEL_CACHE_MAX_ENTRIES = 512 - - -@dataclass(frozen=True) -class _ModelLookupConfig: - umo: str | None - cache_ttl_seconds: float - max_concurrency: int - - -class ListedProvider(Protocol): - def meta(self) -> ProviderMeta: ... - - async def test(self) -> None: ... - - -class _ProviderDisplayEntry(TypedDict): - type: Literal["llm", "tts", "stt"] - info: str - mark: str - provider: ListedProvider - - -ReachabilityCheckResult: TypeAlias = tuple[bool, str | None, str | None] | BaseException - - -class _ModelCache: - def __init__(self) -> None: - self._store: dict[tuple[str, str | None], tuple[float, list[str]]] = {} - - def get(self, provider_id: str, umo: str | None, ttl: float) -> list[str] | None: - if ttl <= 0: - return None - entry = self._store.get((provider_id, umo)) - if not entry: - return None - timestamp, models = entry - if time.monotonic() - timestamp > ttl: - self._store.pop((provider_id, umo), None) - return None - return models - - def set( - self, - provider_id: str, - umo: str | None, - models: list[str], - ttl: float, - ) -> None: - if ttl <= 0: - return - self._store[(provider_id, umo)] = (time.monotonic(), list(models)) - self._evict_if_needed() - - def _evict_if_needed(self) -> None: - if len(self._store) <= MODEL_CACHE_MAX_ENTRIES: - return - # Drop oldest entries first when cache grows too large. - overflow = len(self._store) - MODEL_CACHE_MAX_ENTRIES - for key, _ in sorted( - self._store.items(), - key=lambda item: item[1][0], - )[:overflow]: - self._store.pop(key, None) - - def invalidate( - self, - provider_id: str | None = None, - *, - umo: str | None = None, - ) -> None: - if provider_id is None: - self._store.clear() - return - if umo is not None: - self._store.pop((provider_id, umo), None) - return - stale_keys = [ - cache_key for cache_key in self._store if cache_key[0] == provider_id - ] - for cache_key in stale_keys: - self._store.pop(cache_key, None) - class ProviderCommands: def __init__(self, context: star.Context) -> None: self.context = context - self._model_cache = _ModelCache() - self._register_provider_change_hook() - - def _register_provider_change_hook(self) -> None: - set_change_callback = getattr( - self.context.provider_manager, - "set_provider_change_callback", - None, - ) - if callable(set_change_callback): - set_change_callback(self._on_provider_manager_changed) - return - register_change_hook = getattr( - self.context.provider_manager, - "register_provider_change_hook", - None, - ) - if callable(register_change_hook): - register_change_hook(self._on_provider_manager_changed) - - def invalidate_provider_models_cache( - self, - provider_id: str | None = None, - *, - umo: str | None = None, - ) -> None: - """Public hook for cache invalidation on external provider config changes.""" - self._model_cache.invalidate(provider_id, umo=umo) - - def _on_provider_manager_changed( - self, - provider_id: str, - provider_type: ProviderType, - umo: str | None, - ) -> None: - if provider_type == ProviderType.CHAT_COMPLETION: - self.invalidate_provider_models_cache(provider_id, umo=umo) - - def _get_provider_settings(self, umo: str | None) -> dict: - if not umo: - return {} - try: - return self.context.get_config(umo).get("provider_settings", {}) or {} - except Exception as e: - logger.debug( - "读取 provider_settings 失败,使用默认值: %s", - safe_error("", e), - ) - return {} - - def _get_model_cache_ttl(self, umo: str | None) -> float: - settings = self._get_provider_settings(umo) - raw = settings.get( - MODEL_LIST_CACHE_TTL_KEY, - MODEL_LIST_CACHE_TTL_SECONDS_DEFAULT, - ) - try: - return max(float(raw), 0.0) - except Exception as e: - logger.debug( - "读取 %s 失败,回退默认值 %r: %s", - MODEL_LIST_CACHE_TTL_KEY, - MODEL_LIST_CACHE_TTL_SECONDS_DEFAULT, - safe_error("", e), - ) - return MODEL_LIST_CACHE_TTL_SECONDS_DEFAULT - - def _get_model_lookup_concurrency(self, umo: str | None) -> int: - settings = self._get_provider_settings(umo) - raw = settings.get( - MODEL_LOOKUP_MAX_CONCURRENCY_KEY, - MODEL_LOOKUP_MAX_CONCURRENCY_DEFAULT, - ) - try: - value = int(raw) - except Exception as e: - logger.debug( - "读取 %s 失败,回退默认值 %r: %s", - MODEL_LOOKUP_MAX_CONCURRENCY_KEY, - MODEL_LOOKUP_MAX_CONCURRENCY_DEFAULT, - safe_error("", e), - ) - value = MODEL_LOOKUP_MAX_CONCURRENCY_DEFAULT - return min(max(value, 1), MODEL_LOOKUP_MAX_CONCURRENCY_UPPER_BOUND) - - def _get_model_lookup_config(self, umo: str | None) -> _ModelLookupConfig: - return _ModelLookupConfig( - umo=umo, - cache_ttl_seconds=self._get_model_cache_ttl(umo), - max_concurrency=self._get_model_lookup_concurrency(umo), - ) - - def _resolve_model_name( - self, - model_name: str, - models: Sequence[str], - ) -> str | None: - """Resolve model name with precedence: - exact > case-insensitive > provider-qualified suffix. - """ - requested = model_name.strip() - if not requested: - return None - - requested_norm = requested.casefold() - - # exact / case-insensitive match - for candidate in models: - if candidate == requested or candidate.casefold() == requested_norm: - return candidate - - # provider-qualified suffix match: - # e.g. candidate `openai/gpt-4o` should match requested `gpt-4o`. - for candidate in models: - cand_norm = candidate.casefold() - if cand_norm.endswith(f"/{requested_norm}") or cand_norm.endswith( - f":{requested_norm}", - ): - return candidate - - return None - - def _apply_model( - self, - prov: Provider, - model_name: str, - *, - umo: str | None = None, - ) -> str: - prov.set_model(model_name) - self.invalidate_provider_models_cache(prov.meta().id, umo=umo) - return f"切换模型成功｡当前提供商: [{prov.meta().id}] 当前模型: [{prov.get_model()}]" - - async def _get_provider_models( - self, - provider: Provider, - *, - config: _ModelLookupConfig, - use_cache: bool = True, - ) -> list[str]: - provider_id = provider.meta().id - ttl_seconds = config.cache_ttl_seconds - umo = config.umo - if use_cache: - cached = self._model_cache.get(provider_id, umo, ttl_seconds) - if cached is not None: - return cached - - models = list(await provider.get_models()) - if use_cache: - self._model_cache.set(provider_id, umo, models, ttl_seconds) - return models - - async def _get_models_or_reply_error( - self, - message: AstrMessageEvent, - prov: Provider, - config: _ModelLookupConfig, - *, - error_prefix: str, - disable_t2i: bool = False, - warning_log: str | None = None, - ) -> list[str] | None: - try: - return await self._get_provider_models(prov, config=config) - except asyncio.CancelledError: - raise - except Exception as e: - if warning_log is not None: - logger.warning( - warning_log, - prov.meta().id, - safe_error("", e), - ) - result = MessageEventResult().message(safe_error(error_prefix, e)) - if disable_t2i: - result = result.use_t2i(False) - message.set_result(result) - return None def _log_reachability_failure( self, - provider: ListedProvider, + provider, provider_capability_type: ProviderType | None, err_code: str, err_reason: str, ) -> None: - """记录不可达原因到日志｡""" meta = provider.meta() logger.warning( "Provider reachability check failed: id=%s type=%s code=%s reason=%s", @@ -305,11 +29,7 @@ class ProviderCommands: err_reason, ) - async def _test_provider_capability( - self, - provider: ListedProvider, - ) -> tuple[bool, str | None, str | None]: - """测试单个 provider 的可用性""" + async def _test_provider_capability(self, provider): meta = provider.meta() provider_capability_type = meta.provider_type @@ -320,96 +40,73 @@ class ProviderCommands: err_code = "TEST_FAILED" err_reason = safe_error("", e) self._log_reachability_failure( - provider, - provider_capability_type, - err_code, - err_reason, + provider, provider_capability_type, err_code, err_reason ) return False, err_code, err_reason - async def _find_provider_for_model( + async def _build_provider_display_data( self, - model_name: str, - *, - exclude_provider_id: str | None = None, - config: _ModelLookupConfig, - use_cache: bool = True, - ) -> tuple[Provider | None, str | None]: - all_providers = [] - for provider in self.context.get_all_providers(): - provider_meta = provider.meta() - if provider_meta.provider_type != ProviderType.CHAT_COMPLETION: - continue - if ( - exclude_provider_id is not None - and provider_meta.id == exclude_provider_id - ): - continue - all_providers.append(provider) - if not all_providers: - return None, None + providers, + provider_type: str, + reachability_check_enabled: bool, + ) -> list[dict]: + if not providers: + return [] - semaphore = asyncio.Semaphore(config.max_concurrency) - - async def fetch_models( - provider: Provider, - ) -> tuple[Provider, list[str] | None, str | None]: - async with semaphore: - try: - models = await self._get_provider_models( - provider, - config=config, - use_cache=use_cache, - ) - return provider, models, None - except asyncio.CancelledError: - raise - except Exception as e: - err = safe_error("", e) - logger.debug( - "跨提供商查找模型 %s 获取 %s 模型列表失败: %s", - model_name, - provider.meta().id, - err, - ) - return provider, None, err - - results = await asyncio.gather( - *(fetch_models(provider) for provider in all_providers), - ) - failed_provider_errors: list[tuple[str, str]] = [] - for provider, models, err in results: - if err is not None: - failed_provider_errors.append((provider.meta().id, err)) - continue - if models is None: - continue - - matched_model_name = self._resolve_model_name(model_name, models) - if matched_model_name is not None: - return provider, matched_model_name - - if failed_provider_errors and len(failed_provider_errors) == len(all_providers): - failed_ids = ",".join( - provider_id for provider_id, _ in failed_provider_errors + if reachability_check_enabled: + check_results = await asyncio.gather( + *[self._test_provider_capability(provider) for provider in providers], + return_exceptions=True, ) - logger.error( - "跨提供商查找模型 %s 时,所有 %d 个提供商的 get_models() 均失败: %s｡请检查配置或网络", - model_name, - len(all_providers), - failed_ids, + else: + check_results = [None for _ in providers] + + display_data = [] + for provider, reachable in zip(providers, check_results): + meta = provider.meta() + id_ = meta.id + error_code = None + + if isinstance(reachable, asyncio.CancelledError): + raise reachable + if isinstance(reachable, Exception): + self._log_reachability_failure( + provider, + None, + reachable.__class__.__name__, + safe_error("", reachable), + ) + reachable_flag = False + error_code = reachable.__class__.__name__ + elif isinstance(reachable, tuple): + reachable_flag, error_code, _ = reachable + else: + reachable_flag = reachable + + if provider_type == "llm": + info = f"{id_} ({meta.model})" + else: + info = f"{id_}" + + if reachable_flag is True: + mark = " ✅" + elif reachable_flag is False: + if error_code: + mark = f" ❌(errcode: {error_code})" + else: + mark = " ❌" + else: + mark = "" + + display_data.append( + { + "info": info, + "mark": mark, + "provider": provider, + } ) - elif failed_provider_errors: - logger.debug( - "跨提供商查找模型 %s 时有 %d 个提供商获取模型失败: %s", - model_name, - len(failed_provider_errors), - ",".join( - f"{provider_id}({error})" - for provider_id, error in failed_provider_errors - ), - ) - return None, None + + return display_data async def provider( self, @@ -423,356 +120,129 @@ class ProviderCommands: reachability_check_enabled = cfg.get("reachability_check", True) if idx is None: - parts = ["## 载入的 LLM 提供商\n"] + parts = ["## LLM Providers\n"] - # 获取所有类型的提供商 llms = list(self.context.get_all_providers()) ttss = self.context.get_all_tts_providers() stts = self.context.get_all_stt_providers() - # 构造待检测列表: [(provider, type_label), ...] - all_providers: list[ - tuple[ListedProvider, Literal["llm", "tts", "stt"]] - ] = [] - all_providers.extend([(p, "llm") for p in llms]) - all_providers.extend([(p, "tts") for p in ttss]) - all_providers.extend([(p, "stt") for p in stts]) - - # 并发测试连通性 - if reachability_check_enabled: - if all_providers: - await event.send( - MessageEventResult().message( - "正在进行提供商可达性测试,请稍候...", - ), - ) - check_results: list[ReachabilityCheckResult] = await asyncio.gather( - *[self._test_provider_capability(p) for p, _ in all_providers], - return_exceptions=True, - ) - else: - # 用 None 表示未检测 - check_results = [None for _ in all_providers] - - # 整合结果 - display_data: list[_ProviderDisplayEntry] = [] - for (p, p_type), reachable in zip(all_providers, check_results): - meta = p.meta() - id_ = meta.id - error_code = None - reachable_flag: bool | None - - if isinstance(reachable, asyncio.CancelledError): - raise reachable - if isinstance(reachable, Exception): - # 异常情况下兜底处理,避免单个 provider 导致列表失败 - self._log_reachability_failure( - p, - None, - reachable.__class__.__name__, - safe_error("", reachable), - ) - reachable_flag = False - error_code = reachable.__class__.__name__ - elif isinstance(reachable, tuple): - reachable_flag, error_code, _ = reachable - else: - reachable_flag = None - - # 根据类型构建显示名称 - if p_type == "llm": - info = f"{id_} ({meta.model})" - else: - info = f"{id_}" - - # 确定状态标记 - if reachable_flag is True: - mark = " ✅" - elif reachable_flag is False: - if error_code: - mark = f" ❌(错误码: {error_code})" - else: - mark = " ❌" - else: - mark = "" # 不支持检测时不显示标记 - - display_data.append( - { - "type": p_type, - "info": info, - "mark": mark, - "provider": p, - }, + if reachability_check_enabled and (llms or ttss or stts): + await event.send( + MessageEventResult().message("👀 Testing provider reachability...") ) - # 分组输出 - # 1. LLM - llm_data = [d for d in display_data if d["type"] == "llm"] + llm_data, tts_data, stt_data = await asyncio.gather( + self._build_provider_display_data( + llms, + "llm", + reachability_check_enabled, + ), + self._build_provider_display_data( + ttss, + "tts", + reachability_check_enabled, + ), + self._build_provider_display_data( + stts, + "stt", + reachability_check_enabled, + ), + ) + + provider_using = self.context.get_using_provider(umo=umo) for i, d in enumerate(llm_data): line = f"{i + 1}. {d['info']}{d['mark']}" - provider_using = self.context.get_using_provider(umo=umo) if ( provider_using and provider_using.meta().id == d["provider"].meta().id ): - line += " (当前使用)" + line += " 👈" parts.append(line + "\n") - # 2. TTS - tts_data = [d for d in display_data if d["type"] == "tts"] if tts_data: - parts.append("\n## 载入的 TTS 提供商\n") + parts.append("\n## TTS Providers\n") + tts_using = self.context.get_using_tts_provider(umo=umo) for i, d in enumerate(tts_data): line = f"{i + 1}. {d['info']}{d['mark']}" - tts_using = self.context.get_using_tts_provider(umo=umo) if tts_using and tts_using.meta().id == d["provider"].meta().id: - line += " (当前使用)" + line += " 👈" parts.append(line + "\n") - # 3. STT - stt_data = [d for d in display_data if d["type"] == "stt"] if stt_data: - parts.append("\n## 载入的 STT 提供商\n") + parts.append("\n## STT Providers\n") + stt_using = self.context.get_using_stt_provider(umo=umo) for i, d in enumerate(stt_data): line = f"{i + 1}. {d['info']}{d['mark']}" - stt_using = self.context.get_using_stt_provider(umo=umo) if stt_using and stt_using.meta().id == d["provider"].meta().id: - line += " (当前使用)" + line += " 👈" parts.append(line + "\n") - parts.append("\n使用 /provider <序号> 切换 LLM 提供商｡") + parts.append("\nUse /provider to switch LLM providers.") ret = "".join(parts) if ttss: - ret += "\n使用 /provider tts <序号> 切换 TTS 提供商｡" + ret += "\nUse /provider tts to switch TTS providers." if stts: - ret += "\n使用 /provider stt <序号> 切换 STT 提供商｡" - if not reachability_check_enabled: - ret += "\n已跳过提供商可达性检测,如需检测请在配置文件中开启｡" + ret += "\nUse /provider stt to switch STT providers." event.set_result(MessageEventResult().message(ret)) elif idx == "tts": if idx2 is None: - event.set_result(MessageEventResult().message("请输入序号｡")) + event.set_result( + MessageEventResult().message("Please enter the index.") + ) return if idx2 > len(self.context.get_all_tts_providers()) or idx2 < 1: - event.set_result(MessageEventResult().message("无效的提供商序号｡")) + event.set_result( + MessageEventResult().message("❌ Invalid provider index.") + ) return - tts_provider = self.context.get_all_tts_providers()[idx2 - 1] - id_ = tts_provider.meta().id + provider = self.context.get_all_tts_providers()[idx2 - 1] + id_ = provider.meta().id await self.context.provider_manager.set_provider( provider_id=id_, provider_type=ProviderType.TEXT_TO_SPEECH, umo=umo, ) - event.set_result(MessageEventResult().message(f"成功切换到 {id_}｡")) + event.set_result( + MessageEventResult().message(f"✅ Successfully switched to {id_}.") + ) elif idx == "stt": if idx2 is None: - event.set_result(MessageEventResult().message("请输入序号｡")) + event.set_result( + MessageEventResult().message("Please enter the index.") + ) return if idx2 > len(self.context.get_all_stt_providers()) or idx2 < 1: - event.set_result(MessageEventResult().message("无效的提供商序号｡")) + event.set_result( + MessageEventResult().message("❌ Invalid provider index.") + ) return - stt_provider = self.context.get_all_stt_providers()[idx2 - 1] - id_ = stt_provider.meta().id + provider = self.context.get_all_stt_providers()[idx2 - 1] + id_ = provider.meta().id await self.context.provider_manager.set_provider( provider_id=id_, provider_type=ProviderType.SPEECH_TO_TEXT, umo=umo, ) - event.set_result(MessageEventResult().message(f"成功切换到 {id_}｡")) + event.set_result( + MessageEventResult().message(f"✅ Successfully switched to {id_}.") + ) elif isinstance(idx, int): if idx > len(self.context.get_all_providers()) or idx < 1: - event.set_result(MessageEventResult().message("无效的提供商序号｡")) + event.set_result( + MessageEventResult().message("❌ Invalid provider index.") + ) return - llm_provider = self.context.get_all_providers()[idx - 1] - id_ = llm_provider.meta().id + provider = self.context.get_all_providers()[idx - 1] + id_ = provider.meta().id await self.context.provider_manager.set_provider( provider_id=id_, provider_type=ProviderType.CHAT_COMPLETION, umo=umo, ) - event.set_result(MessageEventResult().message(f"成功切换到 {id_}｡")) + event.set_result( + MessageEventResult().message(f"✅ Successfully switched to {id_}.") + ) else: - event.set_result(MessageEventResult().message("无效的参数｡")) - - async def _switch_model_by_name( - self, - message: AstrMessageEvent, - model_name: str, - prov: Provider, - ) -> None: - model_name = model_name.strip() - if not model_name: - message.set_result(MessageEventResult().message("模型名不能为空｡")) - return - - umo = message.unified_msg_origin - config = self._get_model_lookup_config(umo) - curr_provider_id = prov.meta().id - - models = await self._get_models_or_reply_error( - message, - prov, - config, - error_prefix="获取当前提供商模型列表失败: ", - warning_log="获取当前提供商 %s 模型列表失败,停止跨提供商查找: %s", - ) - if models is None: - return - - matched_model_name = self._resolve_model_name(model_name, models) - if matched_model_name is not None: - message.set_result( - MessageEventResult().message( - self._apply_model(prov, matched_model_name, umo=umo), - ), - ) - return - - target_prov, matched_target_model_name = await self._find_provider_for_model( - model_name, - exclude_provider_id=curr_provider_id, - config=config, - ) - - if target_prov is None or matched_target_model_name is None: - message.set_result( - MessageEventResult().message( - f"模型 [{model_name}] 未在任何已配置的提供商中找到,或所有提供商模型列表获取失败,请检查配置或网络后重试｡", - ), - ) - return - - target_id = target_prov.meta().id - try: - await self.context.provider_manager.set_provider( - provider_id=target_id, - provider_type=ProviderType.CHAT_COMPLETION, - umo=umo, - ) - self._apply_model(target_prov, matched_target_model_name, umo=umo) - message.set_result( - MessageEventResult().message( - f"检测到模型 [{matched_target_model_name}] 属于提供商 [{target_id}],已自动切换提供商并设置模型｡", - ), - ) - except asyncio.CancelledError: - raise - except Exception as e: - message.set_result( - MessageEventResult().message( - safe_error("跨提供商切换并设置模型失败: ", e), - ), - ) - - async def model_ls( - self, - message: AstrMessageEvent, - idx_or_name: int | str | None = None, - ) -> None: - """查看或者切换模型""" - prov = self.context.get_using_provider(message.unified_msg_origin) - if not prov: - message.set_result( - MessageEventResult().message("未找到任何 LLM 提供商｡请先配置｡"), - ) - return - config = self._get_model_lookup_config(message.unified_msg_origin) - - if idx_or_name is None: - models = await self._get_models_or_reply_error( - message, - prov, - config, - error_prefix="获取模型列表失败: ", - disable_t2i=True, - ) - if models is None: - return - parts = ["下面列出了此模型提供商可用模型:"] - for i, model in enumerate(models, 1): - parts.append(f"\n{i}. {model}") - - curr_model = prov.get_model() or "无" - parts.append(f"\n当前模型: [{curr_model}]") - parts.append( - "\nTips: 使用 /model <模型名/编号> 切换模型｡输入模型名时可自动跨提供商查找并切换;跨提供商也可使用 /provider 切换｡", - ) - - ret = "".join(parts) - message.set_result(MessageEventResult().message(ret).use_t2i(False)) - elif isinstance(idx_or_name, int): - models = await self._get_models_or_reply_error( - message, - prov, - config, - error_prefix="获取模型列表失败: ", - ) - if models is None: - return - if idx_or_name > len(models) or idx_or_name < 1: - message.set_result(MessageEventResult().message("模型序号错误｡")) - else: - try: - new_model = models[idx_or_name - 1] - message.set_result( - MessageEventResult().message( - self._apply_model( - prov, - new_model, - umo=message.unified_msg_origin, - ), - ), - ) - except Exception as e: - message.set_result( - MessageEventResult().message( - safe_error("切换模型未知错误: ", e), - ), - ) - return - else: - await self._switch_model_by_name(message, idx_or_name, prov) - - async def key(self, message: AstrMessageEvent, index: int | None = None) -> None: - prov = self.context.get_using_provider(message.unified_msg_origin) - if not prov: - message.set_result( - MessageEventResult().message("未找到任何 LLM 提供商｡请先配置｡"), - ) - return - - if index is None: - keys_data = prov.get_keys() - curr_key = prov.get_current_key() - parts = ["Key:"] - for i, k in enumerate(keys_data, 1): - parts.append(f"\n{i}. {k[:8]}") - - parts.append(f"\n当前 Key: {curr_key[:8]}") - parts.append("\n当前模型: " + prov.get_model()) - parts.append("\n使用 /key 切换 Key｡") - - ret = "".join(parts) - message.set_result(MessageEventResult().message(ret).use_t2i(False)) - else: - keys_data = prov.get_keys() - if index > len(keys_data) or index < 1: - message.set_result(MessageEventResult().message("Key 序号错误｡")) - else: - try: - new_key = keys_data[index - 1] - prov.set_key(new_key) - self.invalidate_provider_models_cache( - prov.meta().id, - umo=message.unified_msg_origin, - ) - message.set_result(MessageEventResult().message("切换 Key 成功｡")) - except Exception as e: - message.set_result( - MessageEventResult().message( - safe_error("切换 Key 未知错误: ", e), - ), - ) - return + event.set_result(MessageEventResult().message("❌ Invalid parameter.")) diff --git a/astrbot/builtin_stars/builtin_commands/main.py b/astrbot/builtin_stars/builtin_commands/main.py index e416678a1..fb8d04f25 100644 --- a/astrbot/builtin_stars/builtin_commands/main.py +++ b/astrbot/builtin_stars/builtin_commands/main.py @@ -6,9 +6,6 @@ from .commands import ( AlterCmdCommands, ConversationCommands, HelpCommand, - LLMCommands, - PersonaCommands, - PluginCommands, ProviderCommands, SetUnsetCommands, SIDCommand, @@ -26,9 +23,8 @@ class Main(star.Star): self.plugin_c = PluginCommands(self.context) self.admin_c = AdminCommands(self.context) self.conversation_c = ConversationCommands(self.context) + self.help_c = HelpCommand(self.context) self.provider_c = ProviderCommands(self.context) - self.persona_c = PersonaCommands(self.context) - self.alter_cmd_c = AlterCmdCommands(self.context) self.setunset_c = SetUnsetCommands(self.context) self.t2i_c = T2ICommand(self.context) self.tts_c = TTSCommand(self.context) @@ -163,41 +159,15 @@ class Main(star.Star): await self.conversation_c.new_conv(message) @filter.permission_type(filter.PermissionType.ADMIN) - @filter.command("groupnew") - async def groupnew_conv(self, message: AstrMessageEvent, sid: str) -> None: - """创建新群聊对话""" - await self.conversation_c.groupnew_conv(message, sid) - - @filter.command("switch") - async def switch_conv( + @filter.command("provider") + async def provider( self, - message: AstrMessageEvent, - index: int | None = None, + event: AstrMessageEvent, + idx: str | int | None = None, + idx2: int | None = None, ) -> None: - """通过 /ls 前面的序号切换对话""" - await self.conversation_c.switch_conv(message, index) - - @filter.command("rename") - async def rename_conv(self, message: AstrMessageEvent, new_name: str) -> None: - """重命名对话""" - await self.conversation_c.rename_conv(message, new_name) - - @filter.command("del") - async def del_conv(self, message: AstrMessageEvent) -> None: - """删除当前对话""" - await self.conversation_c.del_conv(message) - - @filter.permission_type(filter.PermissionType.ADMIN) - @filter.command("key") - async def key(self, message: AstrMessageEvent, index: int | None = None) -> None: - """查看或者切换 Key""" - await self.provider_c.key(message, index) - - @filter.permission_type(filter.PermissionType.ADMIN) - @filter.command("persona") - async def persona(self, message: AstrMessageEvent) -> None: - """查看或者切换 Persona""" - await self.persona_c.persona(message) + """View or switch LLM Provider""" + await self.provider_c.provider(event, idx, idx2) @filter.permission_type(filter.PermissionType.ADMIN) @filter.command("dashboard_update") diff --git a/astrbot/cli/commands/cmd_plug.py b/astrbot/cli/commands/cmd_plug.py index ec5b6668d..e7a7978d0 100644 --- a/astrbot/cli/commands/cmd_plug.py +++ b/astrbot/cli/commands/cmd_plug.py @@ -75,7 +75,7 @@ def new(name: str) -> None: # Rewrite README.md with open(plug_path / "README.md", "w", encoding="utf-8") as f: f.write( - f"# {name}\n\n{desc}\n\n# Support\n\n[Documentation](https://astrbot.app)\n", + f"# {name}\n\n{desc}\n\n# Support\n\n[Documentation](https://docs.astrbot.app)\n" ) # Rewrite main.py diff --git a/astrbot/core/agent/mcp_client.py b/astrbot/core/agent/mcp_client.py index 7f03c1c8a..98e455182 100644 --- a/astrbot/core/agent/mcp_client.py +++ b/astrbot/core/agent/mcp_client.py @@ -14,13 +14,15 @@ This file exists solely for backward compatibility and will be removed in a futu """ import asyncio +import copy import logging import os import sys import warnings from contextlib import AsyncExitStack from datetime import timedelta -from typing import Any, Generic, TextIO +from pathlib import Path, PureWindowsPath +from typing import Any, Generic from tenacity import ( before_sleep_log, @@ -179,6 +181,61 @@ async def _quick_test_mcp_connection(config: dict) -> tuple[bool, str]: return False, f"{e!s}" +def _normalize_mcp_input_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Normalize common non-standard MCP JSON Schema variants. + + Some MCP servers incorrectly mark required properties with a boolean + `required: true` on the property schema itself. Draft 2020-12 requires the + parent object to declare `required` as an array of property names instead. + We lift those booleans to the parent object so the schema remains usable + without disabling validation entirely. + """ + + def _normalize(node: Any) -> Any: + if isinstance(node, list): + return [_normalize(item) for item in node] + + if not isinstance(node, dict): + return node + + normalized = {key: _normalize(value) for key, value in node.items()} + + properties = normalized.get("properties") + if isinstance(properties, dict): + original_properties = ( + node.get("properties") + if isinstance(node.get("properties"), dict) + else {} + ) + required = normalized.get("required") + required_list = required[:] if isinstance(required, list) else [] + + for prop_name, prop_schema in properties.items(): + if not isinstance(prop_schema, dict): + continue + + original_prop_schema = original_properties.get(prop_name, {}) + prop_required = ( + original_prop_schema.get("required") + if isinstance(original_prop_schema, dict) + else None + ) + if isinstance(prop_required, bool): + if prop_schema.get("required") is prop_required: + prop_schema.pop("required", None) + if prop_required: + required_list.append(prop_name) + + if required_list: + normalized["required"] = list(dict.fromkeys(required_list)) + elif isinstance(required, list): + normalized.pop("required", None) + + return normalized + + return _normalize(copy.deepcopy(schema)) + + class MCPClient: def __init__(self) -> None: # Initialize session and client objects @@ -494,7 +551,7 @@ class MCPTool(FunctionTool, Generic[TContext]): super().__init__( name=mcp_tool.name, description=mcp_tool.description or "", - parameters=mcp_tool.inputSchema, + parameters=_normalize_mcp_input_schema(mcp_tool.inputSchema), ) self.mcp_tool = mcp_tool self.mcp_client = mcp_client diff --git a/astrbot/core/agent/message.py b/astrbot/core/agent/message.py index 4f46ff6ec..448c30022 100644 --- a/astrbot/core/agent/message.py +++ b/astrbot/core/agent/message.py @@ -8,6 +8,7 @@ from pydantic import ( BaseModel, GetCoreSchemaHandler, PrivateAttr, + ValidationError, model_serializer, model_validator, ) @@ -167,6 +168,15 @@ class ToolCallPart(BaseModel): """A part of the arguments of the tool call.""" +class CheckpointData(BaseModel): + """Internal checkpoint data for linking LLM turns to platform history.""" + + id: str + + +CHECKPOINT_ROLE = "_checkpoint" + + class Message(BaseModel): """A message in a conversation.""" @@ -175,9 +185,10 @@ class Message(BaseModel): "user", "assistant", "tool", + "_checkpoint", ] - content: str | list[ContentPart] | None = None + content: str | list[ContentPart] | CheckpointData | None = None """The content of the message.""" tool_calls: list[ToolCall] | list[dict] | None = None @@ -187,9 +198,18 @@ class Message(BaseModel): """The ID of the tool call.""" _no_save: bool = PrivateAttr(default=False) + _checkpoint_after: CheckpointData | None = PrivateAttr(default=None) @model_validator(mode="after") def check_content_required(self): + if self.role == CHECKPOINT_ROLE: + if not isinstance(self.content, CheckpointData): + raise ValueError("checkpoint message content must be CheckpointData") + return self + + if isinstance(self.content, CheckpointData): + raise ValueError("CheckpointData is only allowed for role='_checkpoint'") + # assistant + tool_calls is not None: allow content to be None if self.role == "assistant" and self.tool_calls is not None: return self @@ -233,3 +253,87 @@ class SystemMessageSegment(Message): """A message segment from the system.""" role: Literal["system"] = "system" + + +class CheckpointMessageSegment(Message): + """Internal checkpoint segment for persisted conversation history.""" + + role: Literal["_checkpoint"] = "_checkpoint" + content: CheckpointData | None = None + + +def is_checkpoint_message(message: Message | dict) -> bool: + """Return whether a message is an internal checkpoint.""" + if isinstance(message, Message): + return message.role == CHECKPOINT_ROLE + return isinstance(message, dict) and message.get("role") == CHECKPOINT_ROLE + + +def get_checkpoint_id(message: Message | dict) -> str | None: + """Return the checkpoint id from an internal checkpoint message.""" + if not is_checkpoint_message(message): + return None + + content = ( + message.content if isinstance(message, Message) else message.get("content") + ) + if isinstance(content, CheckpointData): + return content.id + if isinstance(content, dict): + checkpoint_id = content.get("id") + return ( + checkpoint_id if isinstance(checkpoint_id, str) and checkpoint_id else None + ) + return None + + +def strip_checkpoint_messages(history: list[dict]) -> list[dict]: + """Remove internal checkpoint messages from provider-facing history.""" + return [message for message in history if not is_checkpoint_message(message)] + + +def _get_checkpoint_data(message: Message | dict) -> CheckpointData | None: + if not is_checkpoint_message(message): + return None + + content = ( + message.content if isinstance(message, Message) else message.get("content") + ) + if isinstance(content, CheckpointData): + return content + if isinstance(content, dict): + try: + return CheckpointData.model_validate(content) + except ValidationError: + return None + return None + + +def bind_checkpoint_messages(history: list[dict]) -> list[Message]: + """Load persisted history and bind checkpoint segments to prior messages.""" + messages: list[Message] = [] + for item in history: + if is_checkpoint_message(item): + checkpoint = _get_checkpoint_data(item) + if checkpoint is not None and messages: + messages[-1]._checkpoint_after = checkpoint + continue + + message = Message.model_validate(item) + if item.get("_no_save"): + message._no_save = True + messages.append(message) + + return messages + + +def dump_messages_with_checkpoints(messages: list[Message]) -> list[dict]: + """Dump runtime messages and reinsert bound checkpoint segments.""" + dumped: list[dict] = [] + for message in messages: + dumped.append(message.model_dump()) + if message._checkpoint_after is not None: + dumped.append( + CheckpointMessageSegment(content=message._checkpoint_after).model_dump() + ) + return dumped diff --git a/astrbot/core/agent/runners/coze/coze_agent_runner.py b/astrbot/core/agent/runners/coze/coze_agent_runner.py index b12bb51e9..2f24556ef 100644 --- a/astrbot/core/agent/runners/coze/coze_agent_runner.py +++ b/astrbot/core/agent/runners/coze/coze_agent_runner.py @@ -18,6 +18,11 @@ from astrbot.core.provider.entities import ( ) from astrbot.core.provider.provider import Provider +from ...hooks import BaseAgentRunHooks +from ...message import is_checkpoint_message +from ...response import AgentResponseData +from ...run_context import ContextWrapper, TContext +from ..base import AgentResponse, AgentState, BaseAgentRunner from .coze_api_client import CozeAPIClient if sys.version_info >= (3, 12): @@ -160,6 +165,8 @@ class CozeAgentRunner(BaseAgentRunner[TContext]): # 处理历史上下文 if not self.auto_save_history and contexts: for ctx in contexts: + if is_checkpoint_message(ctx): + continue if isinstance(ctx, dict) and "role" in ctx and "content" in ctx: # 处理上下文中的图片 content = ctx["content"] diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index de224bf38..4b47c5ff5 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -5,8 +5,8 @@ import time import traceback from collections.abc import AsyncGenerator, AsyncIterator from contextlib import suppress -from dataclasses import dataclass, field -from typing import Any, Literal, TypeVar +from dataclasses import dataclass, field, replace +from pathlib import Path from mcp.types import ( BlobResourceContents, @@ -51,8 +51,28 @@ from astrbot.core.provider.entities import ( ProviderRequest, ToolCallsResult, ) +from astrbot.core.provider.modalities import ( + log_context_sanitize_stats, + sanitize_contexts_by_modalities, +) from astrbot.core.provider.provider import Provider +from ..context.compressor import ContextCompressor +from ..context.config import ContextConfig +from ..context.manager import ContextManager +from ..context.token_counter import EstimateTokenCounter, TokenCounter +from ..hooks import BaseAgentRunHooks +from ..message import ( + AssistantMessageSegment, + Message, + ToolCallMessageSegment, + bind_checkpoint_messages, +) +from ..response import AgentResponseData, AgentStats +from ..run_context import ContextWrapper, TContext +from ..tool_executor import BaseFunctionToolExecutor +from .base import AgentResponse, AgentState, BaseAgentRunner + if sys.version_info >= (3, 12): from typing import override else: @@ -256,16 +276,16 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): # MODIFIE the req.func_tool to use light tool schemas self.req.func_tool = light_set - messages = [] # append existing messages in the run context - for msg in request.contexts: - m = Message.model_validate(msg) - if isinstance(msg, dict) and msg.get("_no_save"): - m._no_save = True - messages.append(m) - if request.prompt is not None: - assembled_context = await request.assemble_context() - messages.append(Message.model_validate(assembled_context)) + messages = bind_checkpoint_messages(request.contexts or []) + if ( + request.prompt is not None + or request.image_urls + or request.audio_urls + or request.extra_user_content_parts + ): + m = await self._assemble_request_context_for_provider(request) + messages.append(Message.model_validate(m)) if request.system_prompt: messages.insert( 0, @@ -276,15 +296,148 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): self.stats = AgentStats() self.stats.start_time = time.time() + def _read_tool_hint(self) -> str: + if self.read_tool is not None: + return f"`{self.read_tool.name}`" + return "the available file-read tool" + + async def _assemble_request_context_for_provider( + self, + request: ProviderRequest, + ) -> dict[str, T.Any]: + modalities = self.provider.provider_config.get("modalities", None) + if not isinstance(modalities, list): + return await request.assemble_context() + + supports_image = "image" in modalities + supports_audio = "audio" in modalities + if supports_image and supports_audio: + return await request.assemble_context() + + adjusted_request = replace( + request, + image_urls=request.image_urls if supports_image else [], + audio_urls=request.audio_urls if supports_audio else [], + ) + context = await adjusted_request.assemble_context() + content = context.get("content") + if isinstance(content, str): + content_blocks: list[dict[str, T.Any]] = [{"type": "text", "text": content}] + elif isinstance(content, list): + content_blocks = content + else: + content_blocks = [] + + if not supports_image: + for _ in request.image_urls: + content_blocks.append({"type": "text", "text": "[Image]"}) + if not supports_audio: + for _ in request.audio_urls: + content_blocks.append({"type": "text", "text": "[Audio]"}) + + return {"role": "user", "content": content_blocks} + + async def _write_tool_result_overflow_file( + self, + *, + tool_call_id: str, + content: str, + ) -> str: + if self.tool_result_overflow_dir is None: + raise ValueError("tool_result_overflow_dir is not configured") + + overflow_dir = Path(self.tool_result_overflow_dir).resolve(strict=False) + safe_tool_call_id = ( + "".join( + ch if ch.isalnum() or ch in {"-", "_", "."} else "_" + for ch in tool_call_id + ).strip("._") + or "tool_call" + ) + file_name = f"{safe_tool_call_id}_{uuid.uuid4().hex[:8]}.txt" + overflow_path = overflow_dir / file_name + + def _run() -> str: + overflow_dir.mkdir(parents=True, exist_ok=True) + overflow_path.write_text(content, encoding="utf-8") + return str(overflow_path) + + return await asyncio.to_thread(_run) + + async def _materialize_large_tool_result( + self, + *, + tool_call_id: str, + content: str, + ) -> str: + if self.tool_result_overflow_dir is None or self.read_tool is None: + return content + + estimated_tokens = self._tool_result_token_counter.count_tokens( + [Message(role="tool", content=content, tool_call_id=tool_call_id)] + ) + if estimated_tokens <= self.TOOL_RESULT_MAX_ESTIMATED_TOKENS: + return content + + preview = self._truncate_tool_result_preview(content, tool_call_id=tool_call_id) + try: + overflow_path = await self._write_tool_result_overflow_file( + tool_call_id=tool_call_id, + content=content, + ) + except Exception as exc: + logger.warning( + "Failed to spill oversized tool result for %s: %s", + tool_call_id, + exc, + exc_info=True, + ) + error_notice = ( + "Tool output exceeded the inline result limit " + f"({estimated_tokens} estimated tokens > " + f"{self.TOOL_RESULT_MAX_ESTIMATED_TOKENS}) and could not be written " + f"to `{self.tool_result_overflow_dir}`: {exc}" + ) + if not preview: + return error_notice + return f"{preview}\n\n{error_notice}" + + notice = self.TOOL_RESULT_OVERFLOW_NOTICE_TEMPLATE.format( + overflow_path=overflow_path, + read_tool_hint=self._read_tool_hint(), + ) + if not preview: + return notice + return f"{preview}\n\n{notice}" + + def _truncate_tool_result_preview( + self, + content: str, + *, + tool_call_id: str, + ) -> str: + preview = content + while preview: + estimated_tokens = self._tool_result_token_counter.count_tokens( + [Message(role="tool", content=preview, tool_call_id=tool_call_id)] + ) + if estimated_tokens <= self.TOOL_RESULT_PREVIEW_MAX_ESTIMATED_TOKENS: + return preview + next_len = len(preview) // 2 + if next_len <= 0: + break + preview = preview[:next_len] + return preview + async def _iter_llm_responses( self, *, include_model: bool = True, ) -> AsyncGenerator[LLMResponse, None]: """Yields chunks *and* a final LLMResponse.""" - payload: dict[str, Any] = { - "contexts": self.run_context.messages, # list[Message] - "func_tool": self.req.func_tool, + payload = { + "contexts": self._sanitize_contexts_for_provider(self.run_context.messages), + "func_tool": self._func_tool_for_provider(), "session_id": self.req.session_id, "extra_user_content_parts": self.req.extra_user_content_parts, # list[ContentPart] "abort_signal": self._abort_signal, @@ -428,6 +581,35 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): completion_text="All available chat models are unavailable.", ) + def _sanitize_contexts_for_provider( + self, + contexts: list[Message] | list[dict[str, T.Any]], + ) -> list[Message] | list[dict[str, T.Any]]: + if not self._should_fix_modalities_for_provider(): + return contexts + sanitized_contexts, stats = sanitize_contexts_by_modalities( + contexts, + self.provider.provider_config.get("modalities", None), + ) + log_context_sanitize_stats(stats) + return sanitized_contexts + + def _should_fix_modalities_for_provider(self) -> bool: + modalities = self.provider.provider_config.get("modalities", None) + return isinstance(modalities, list) + + def _func_tool_for_provider(self) -> ToolSet | None: + if not self.req.func_tool: + return None + modalities = self.provider.provider_config.get("modalities", None) + if isinstance(modalities, list) and "tool_use" not in modalities: + logger.debug( + "Provider %s does not support tool_use, clearing tools for request.", + self.provider, + ) + return None + return self.req.func_tool + def _simple_print_message_role(self, tag: str = ""): roles = [] for message in self.run_context.messages: @@ -537,7 +719,6 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): async for llm_response in self._iter_llm_responses_with_fallback(): if llm_response.is_chunk: - # update ttft if self.stats.time_to_first_token == 0: self.stats.time_to_first_token = time.time() - self.stats.start_time @@ -553,7 +734,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): chain=MessageChain().message(llm_response.completion_text), ), ) - elif llm_response.reasoning_content: + if llm_response.reasoning_content: yield AgentResponse( type="streaming_delta", data=AgentResponseData( @@ -654,11 +835,47 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): chain=MessageChain().message(llm_resp.completion_text), ), ) + if llm_resp.reasoning_content: + yield AgentResponse( + type="llm_result", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_resp.reasoning_content, + ), + ), + ) # 如果有工具调用,还需处理工具调用 if llm_resp.tools_call_name: if self.tool_schema_mode == "lazy_load": llm_resp, _ = await self._resolve_tool_exec(llm_resp) + if not llm_resp.tools_call_name: + logger.warning( + "skills_like tool re-query returned no tool calls; fallback to assistant response." + ) + if llm_resp.result_chain: + yield AgentResponse( + type="llm_result", + data=AgentResponseData(chain=llm_resp.result_chain), + ) + elif llm_resp.completion_text: + yield AgentResponse( + type="llm_result", + data=AgentResponseData( + chain=MessageChain().message(llm_resp.completion_text), + ), + ) + if llm_resp.reasoning_content: + yield AgentResponse( + type="llm_result", + data=AgentResponseData( + chain=MessageChain(type="reasoning").message( + llm_resp.reasoning_content, + ), + ), + ) + await self._complete_with_assistant_response(llm_resp) + return tool_call_result_blocks = [] cached_images = [] # Collect cached images for LLM visibility @@ -855,9 +1072,11 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): ): # in 'lazy_load' mode, raw.func_tool is light schema, does not have handler # so we need to get the tool from the raw tool set - func_tool = self._lazy_load_raw_tool_set.get_tool(func_tool_name) + func_tool = self._skill_like_raw_tool_set.get_tool(func_tool_name) + available_tools = self._skill_like_raw_tool_set.names() else: func_tool = req.func_tool.get_tool(func_tool_name) + available_tools = req.func_tool.names() logger.info(f"使用工具:{func_tool_name},参数:{func_tool_args}") @@ -865,7 +1084,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): logger.warning(f"未找到指定的工具: {func_tool_name},将跳过｡") _append_tool_call_result( func_tool_id, - f"error: Tool {func_tool_name} not found.", + f"error: Tool {func_tool_name} not found. Available tools are: {', '.join(available_tools)}", ) continue @@ -1122,7 +1341,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): if param_subset.tools and tool_names: contexts = self._build_tool_requery_context(tool_names) requery_resp = await self.provider.text_chat( - contexts=contexts, + contexts=self._sanitize_contexts_for_provider(contexts), func_tool=param_subset, model=self.req.model, session_id=self.req.session_id, @@ -1147,7 +1366,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]): extra_instruction=self.SKILLS_LIKE_REQUERY_REPAIR_INSTRUCTION, ) repair_resp = await self.provider.text_chat( - contexts=repair_contexts, + contexts=self._sanitize_contexts_for_provider(repair_contexts), func_tool=param_subset, model=self.req.model, session_id=self.req.session_id, diff --git a/astrbot/core/astr_agent_run_util.py b/astrbot/core/astr_agent_run_util.py index de53e305a..8100be85b 100644 --- a/astrbot/core/astr_agent_run_util.py +++ b/astrbot/core/astr_agent_run_util.py @@ -91,19 +91,29 @@ def _build_tool_result_status_message( return status_msg -def _extract_final_streaming_chain(msg_chain: MessageChain) -> MessageChain | None: - if not msg_chain.chain: +def _should_buffer_llm_result( + buffer_intermediate_messages: bool, + stream_to_general: bool, + agent_runner: AgentRunner, +) -> bool: + return ( + buffer_intermediate_messages + and not stream_to_general + and not agent_runner.streaming + ) + + +def _merge_buffered_llm_chains( + buffered_llm_chains: list[MessageChain], +) -> MessageChain | None: + if not buffered_llm_chains: return None - final_chain: list[BaseMessageComponent] = [] - for comp in msg_chain.chain: - if isinstance(comp, Plain): - continue - final_chain.append(comp) - - if not final_chain: - return None - return MessageChain(chain=final_chain, type=msg_chain.type) + merged_chain = MessageChain() + for chain in buffered_llm_chains: + merged_chain.chain.extend(chain.chain) + buffered_llm_chains.clear() + return merged_chain async def run_agent( @@ -113,10 +123,17 @@ async def run_agent( show_tool_call_result: bool = False, stream_to_general: bool = False, show_reasoning: bool = False, + buffer_intermediate_messages: bool = False, ) -> AsyncGenerator[MessageChain | None, None]: step_idx = 0 astr_event = agent_runner.run_context.context.event tool_name_by_call_id: dict[str, str] = {} + buffered_llm_chains: list[MessageChain] = [] + can_buffer_llm_result = _should_buffer_llm_result( + buffer_intermediate_messages, + stream_to_general, + agent_runner, + ) while step_idx < max_step + 1: step_idx += 1 @@ -145,6 +162,17 @@ async def run_agent( agent_runner.request_stop() if resp.type == "aborted": + if can_buffer_llm_result: + merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) + if merged_chain: + astr_event.set_result( + MessageEventResult( + chain=merged_chain.chain, + result_content_type=ResultContentType.LLM_RESULT, + ), + ) + yield merged_chain + astr_event.clear_result() if not stop_watcher.done(): stop_watcher.cancel() try: @@ -217,6 +245,10 @@ async def run_agent( continue if stream_to_general or not agent_runner.streaming: + if can_buffer_llm_result and resp.type == "llm_result": + buffered_llm_chains.append(resp.data["chain"]) + continue + content_typ = ( ResultContentType.LLM_RESULT if resp.type == "llm_result" @@ -228,7 +260,7 @@ async def run_agent( result_content_type=content_typ, ), ) - yield + yield resp.data["chain"] astr_event.clear_result() elif resp.type == "streaming_delta": chain = resp.data["chain"] @@ -236,11 +268,19 @@ async def run_agent( # display the reasoning content only when configured continue yield resp.data["chain"] # MessageChain - elif resp.type == "llm_result": - if final_chain := _extract_final_streaming_chain( - resp.data["chain"], - ): - yield final_chain + + if can_buffer_llm_result and agent_runner.done(): + merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) + if merged_chain: + astr_event.set_result( + MessageEventResult( + chain=merged_chain.chain, + result_content_type=ResultContentType.LLM_RESULT, + ), + ) + yield merged_chain + astr_event.clear_result() + if not stop_watcher.done(): stop_watcher.cancel() try: @@ -314,6 +354,7 @@ async def run_live_agent( show_tool_use: bool = True, show_tool_call_result: bool = False, show_reasoning: bool = False, + buffer_intermediate_messages: bool = False, ) -> AsyncGenerator[MessageChain | None, None]: """Live Mode 的 Agent 运行器,支持流式 TTS @@ -338,6 +379,7 @@ async def run_live_agent( show_tool_call_result=show_tool_call_result, stream_to_general=False, show_reasoning=show_reasoning, + buffer_intermediate_messages=buffer_intermediate_messages, ): yield chain return @@ -370,7 +412,8 @@ async def run_live_agent( show_tool_use, show_tool_call_result, show_reasoning, - ), + buffer_intermediate_messages, + ) ) # 2. 启动 TTS 任务:负责从 text_queue 读取文本并生成音频到 audio_queue @@ -456,6 +499,7 @@ async def _run_agent_feeder( show_tool_use: bool, show_tool_call_result: bool, show_reasoning: bool, + buffer_intermediate_messages: bool, ) -> None: """运行 Agent 并将文本输出分句放入队列""" buffer = "" @@ -467,6 +511,7 @@ async def _run_agent_feeder( show_tool_call_result=show_tool_call_result, stream_to_general=False, show_reasoning=show_reasoning, + buffer_intermediate_messages=buffer_intermediate_messages, ): if chain is None: continue diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index 9e317300e..815e07bc0 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -21,7 +21,7 @@ from astrbot.core.astr_agent_hooks import MAIN_AGENT_HOOKS from astrbot.core.astr_agent_run_util import AgentRunner from astrbot.core.astr_agent_tool_exec import FunctionToolExecutor from astrbot.core.conversation_mgr import Conversation -from astrbot.core.message.components import File, Image, Record, Reply +from astrbot.core.message.components import File, Image, Record, Reply, Video from astrbot.core.persona_error_reply import ( extract_persona_custom_error_message_from_persona, set_persona_custom_error_message_on_event, @@ -449,6 +449,33 @@ def _append_quoted_audio_attachment(req: ProviderRequest, audio_path: str) -> No ) +async def _append_video_attachment( + req: ProviderRequest, + video: Video, + *, + quoted: bool = False, +) -> None: + try: + video_path = await video.convert_to_file_path() + except Exception as exc: # noqa: BLE001 + if quoted: + logger.error("Error processing quoted video attachment: %s", exc) + else: + logger.error("Error processing video attachment: %s", exc) + return + + video_name = os.path.basename(video_path) + if quoted: + text = ( + f"[Video Attachment in quoted message: " + f"name {video_name}, path {video_path}]" + ) + else: + text = f"[Video Attachment: name {video_name}, path {video_path}]" + + req.extra_user_content_parts.append(TextPart(text=text)) + + def _get_quoted_message_parser_settings( provider_settings: dict[str, object] | None, ) -> QuotedMessageParserSettings: @@ -742,143 +769,8 @@ async def _decorate_llm_request( _append_system_reminders(event, req, cfg, tz) -def _modalities_fix(provider: Provider, req: ProviderRequest) -> None: - if req.image_urls: - provider_cfg = provider.provider_config.get("modalities", ["image"]) - if "image" not in provider_cfg: - logger.debug( - "Provider %s does not support image, using placeholder.", - provider, - ) - image_count = len(req.image_urls) - placeholder = " ".join(["[Image]"] * image_count) - if req.prompt: - req.prompt = f"{placeholder} {req.prompt}" - else: - req.prompt = placeholder - req.image_urls = [] - if req.audio_urls: - provider_cfg = provider.provider_config.get("modalities", ["audio"]) - if "audio" not in provider_cfg: - logger.debug( - "Provider %s does not support audio, using placeholder.", - provider, - ) - audio_count = len(req.audio_urls) - placeholder = " ".join(["[Audio]"] * audio_count) - if req.prompt: - req.prompt = f"{placeholder} {req.prompt}" - else: - req.prompt = placeholder - req.audio_urls = [] - if req.func_tool: - provider_cfg = provider.provider_config.get("modalities", ["tool_use"]) - if "tool_use" not in provider_cfg: - logger.debug( - "Provider %s does not support tool_use, clearing tools.", - provider, - ) - req.func_tool = None - - -def _sanitize_context_by_modalities( - config: MainAgentBuildConfig, - provider: Provider, - req: ProviderRequest, -) -> None: - if not config.sanitize_context_by_modalities: - return - if not isinstance(req.contexts, list) or not req.contexts: - return - modalities = provider.provider_config.get("modalities", None) - if not modalities or not isinstance(modalities, list): - return - supports_image = bool("image" in modalities) - supports_audio = bool("audio" in modalities) - supports_tool_use = bool("tool_use" in modalities) - if supports_image and supports_audio and supports_tool_use: - return - sanitized_contexts: list[dict] = [] - removed_image_blocks = 0 - removed_audio_blocks = 0 - removed_tool_messages = 0 - removed_tool_calls = 0 - for msg in req.contexts: - if not isinstance(msg, dict): - continue - role = msg.get("role") - if not role: - continue - new_msg = msg - if not supports_tool_use: - if role == "tool": - removed_tool_messages += 1 - continue - if role == "assistant" and "tool_calls" in new_msg: - if "tool_calls" in new_msg: - removed_tool_calls += 1 - new_msg.pop("tool_calls", None) - new_msg.pop("tool_call_id", None) - - if not supports_image or not supports_audio: - content = new_msg.get("content") - if isinstance(content, list): - filtered_parts: list = [] - removed_any_multimodal = False - for part in content: - if isinstance(part, dict): - part_type = str(part.get("type", "")).lower() - if not supports_image and part_type in {"image_url", "image"}: - removed_any_multimodal = True - removed_image_blocks += 1 - continue - if not supports_audio and part_type in { - "audio_url", - "input_audio", - }: - removed_any_multimodal = True - removed_audio_blocks += 1 - continue - filtered_parts.append(part) - if removed_any_multimodal: - new_msg["content"] = filtered_parts - if role == "assistant": - content = new_msg.get("content") - has_tool_calls = bool(new_msg.get("tool_calls")) - if not has_tool_calls: - if not content: - continue - if isinstance(content, str) and (not content.strip()): - continue - sanitized_contexts.append(new_msg) - - if ( - removed_image_blocks - or removed_audio_blocks - or removed_tool_messages - or removed_tool_calls - ): - logger.debug( - "sanitize_context_by_modalities applied: " - "removed_image_blocks=%s, removed_audio_blocks=%s, " - "removed_tool_messages=%s, removed_tool_calls=%s", - removed_image_blocks, - removed_audio_blocks, - removed_tool_messages, - removed_tool_calls, - ) - req.contexts = sanitized_contexts - - -def _model_outputs_image(provider: Provider, req: ProviderRequest) -> bool: - model = req.model or provider.get_model() - if not model: - return False - model_info = LLM_METADATAS.get(model) - if not model_info: - return False - output_modalities = model_info.get("modalities", {}).get("output", []) - return "image" in output_modalities +def _plugin_tool_fix(event: AstrMessageEvent, req: ProviderRequest) -> None: + """根据事件中的插件设置，过滤请求中的工具列表。 def _should_disable_streaming_for_webchat_output( @@ -1100,6 +992,9 @@ async def build_main_agent( text=f"[File Attachment: name {file_name}, path {file_path}]", ), ) + elif isinstance(comp, Video): + await _append_video_attachment(req, comp) + # quoted message attachments reply_comps = [ comp for comp in event.message_obj.message if isinstance(comp, Reply) ] @@ -1134,6 +1029,11 @@ async def build_main_agent( text=f"[File Attachment in quoted message: name {file_name}, path {file_path}]", ), ) + elif isinstance(reply_comp, Video): + await _append_video_attachment(req, reply_comp, quoted=True) + + # Fallback quoted image extraction for reply-id-only payloads, or when + # embedded reply chain only contains placeholders (e.g. [Forward Message], [Image]). if not has_embedded_image: try: fallback_images = normalize_and_dedupe_strings( @@ -1184,6 +1084,17 @@ async def build_main_agent( event.set_extra("provider_request", req) if isinstance(req.contexts, str): req.contexts = json.loads(req.contexts) + thread_selected_text = event.get_extra("thread_selected_text") + if isinstance(thread_selected_text, str) and thread_selected_text.strip(): + req.extra_user_content_parts.append( + TextPart( + text=( + "The user is asking in a side thread about this selected " + "excerpt from the previous assistant answer:\n" + f"{thread_selected_text.strip()}" + ) + ) + ) req.image_urls = normalize_and_dedupe_strings(req.image_urls) req.audio_urls = normalize_and_dedupe_strings(req.audio_urls) @@ -1202,9 +1113,10 @@ async def build_main_agent( await _apply_kb(event, req, plugin_context, config) if not req.session_id: req.session_id = event.unified_msg_origin - _modalities_fix(provider, req) + _plugin_tool_fix(event, req) - _sanitize_context_by_modalities(config, provider, req) + await _apply_web_search_tools(event, req, plugin_context) + if config.llm_safety_mode: _apply_llm_safety_mode(config, req) if config.tool_providers: diff --git a/astrbot/core/backup/constants.py b/astrbot/core/backup/constants.py index 5ff7eb3ec..43af108e7 100644 --- a/astrbot/core/backup/constants.py +++ b/astrbot/core/backup/constants.py @@ -18,6 +18,7 @@ from astrbot.core.db.po import ( PlatformStat, Preference, SessionProjectRelation, + WebChatThread, ) from astrbot.core.knowledge_base.models import ( KBDocument, @@ -46,6 +47,7 @@ MAIN_DB_MODELS: dict[str, type[SQLModel]] = { "preferences": Preference, "platform_message_history": PlatformMessageHistory, "platform_sessions": PlatformSession, + "webchat_threads": WebChatThread, "chatui_projects": ChatUIProject, "session_project_relations": SessionProjectRelation, "attachments": Attachment, diff --git a/astrbot/core/backup/importer.py b/astrbot/core/backup/importer.py index d78e2cc69..084498817 100644 --- a/astrbot/core/backup/importer.py +++ b/astrbot/core/backup/importer.py @@ -61,6 +61,20 @@ def _get_major_version(version_str: str) -> str: return "0.0" +def _validate_path_within(target_path: Path, base_dir: Path) -> bool: + """Validate that target_path is within base_dir after resolving symlinks. + + Prevents path traversal attacks (CWE-22) by ensuring the resolved + target path is relative to the resolved base directory. + """ + try: + resolved = target_path.resolve(strict=False) + base_resolved = base_dir.resolve(strict=False) + return resolved.is_relative_to(base_resolved) + except (OSError, ValueError): + return False + + CMD_CONFIG_FILE_PATH = os.path.join(get_astrbot_data_path(), "cmd_config.json") KB_PATH = get_astrbot_knowledge_base_path() DEFAULT_PLATFORM_STATS_INVALID_COUNT_WARN_LIMIT = 5 @@ -777,14 +791,13 @@ class AstrBotImporter: try: rel_path = name[len(media_prefix) :] target_path = kb_dir / rel_path - await anyio.Path(target_path.parent).mkdir( - parents=True, - exist_ok=True, - ) - with zf.open(name) as src: - content = src.read() - async with await anyio.open_file(target_path, "wb") as dst: - await dst.write(content) + # Validate path is within kb directory (CWE-22) + if not _validate_path_within(target_path, kb_dir): + logger.warning(f"媒体文件路径越界，已跳过: {target_path}") + continue + target_path.parent.mkdir(parents=True, exist_ok=True) + with zf.open(name) as src, open(target_path, "wb") as dst: + dst.write(src.read()) except Exception as e: result.add_warning(f"导入媒体文件 {name} 失败: {e}") @@ -844,14 +857,14 @@ class AstrBotImporter: else: target_path = attachments_dir / os.path.basename(name) - await anyio.Path(target_path.parent).mkdir( - parents=True, - exist_ok=True, - ) - with zf.open(name) as src: - content = src.read() - async with await anyio.open_file(target_path, "wb") as dst: - await dst.write(content) + # Validate path is within attachments directory (CWE-22) + if not _validate_path_within(target_path, attachments_dir): + logger.warning(f"附件路径越界，已跳过: {target_path}") + continue + + target_path.parent.mkdir(parents=True, exist_ok=True) + with zf.open(name) as src, open(target_path, "wb") as dst: + dst.write(src.read()) count += 1 except Exception as e: logger.warning(f"导入附件 {name} 失败: {e}") @@ -927,10 +940,11 @@ class AstrBotImporter: continue target_path = target_dir / rel_path - await anyio.Path(target_path.parent).mkdir( - parents=True, - exist_ok=True, - ) + # Validate path is within target directory (CWE-22) + if not _validate_path_within(target_path, target_dir): + result.add_warning(f"文件路径越界，已跳过: {name}") + continue + target_path.parent.mkdir(parents=True, exist_ok=True) with zf.open(name) as src: content = src.read() diff --git a/astrbot/core/computer/file_read_utils.py b/astrbot/core/computer/file_read_utils.py index d5a8cf72a..a90196dcd 100644 --- a/astrbot/core/computer/file_read_utils.py +++ b/astrbot/core/computer/file_read_utils.py @@ -73,7 +73,7 @@ class FileProbe: @dataclass(frozen=True) class ParsedDocument: - kind: Literal["docx", "pdf"] + kind: Literal["docx", "epub", "pdf"] file_bytes: bytes text: str @@ -91,7 +91,7 @@ print( json.dumps( {{ "size_bytes": path.stat().st_size, - "sample_b64": base64.b64encode(sample).decode("ascii"), + "sample_b64": base64.b64encode(sample).decode("utf-8"), }} ) ) @@ -140,7 +140,7 @@ print( json.dumps( {{ "size_bytes": len(data), - "base64": base64.b64encode(data).decode("ascii"), + "base64": base64.b64encode(data).decode("utf-8"), }} ) ) @@ -278,7 +278,7 @@ async def _probe_local_file(path: str) -> dict[str, str | int]: sample = file_obj.read(_FILE_SNIFF_BYTES) return { "size_bytes": file_path.stat().st_size, - "sample_b64": base64.b64encode(sample).decode("ascii"), + "sample_b64": base64.b64encode(sample).decode("utf-8"), } return await to_thread(_run) @@ -289,7 +289,7 @@ async def _read_local_image_base64(path: str) -> dict[str, str | int]: data = Path(path).read_bytes() return { "size_bytes": len(data), - "base64": base64.b64encode(data).decode("ascii"), + "base64": base64.b64encode(data).decode("utf-8"), } return await to_thread(_run) @@ -319,7 +319,7 @@ async def _compress_image_bytes_to_base64(data: bytes) -> dict[str, str | int]: return { "size_bytes": len(compressed_bytes), - "base64": base64.b64encode(compressed_bytes).decode("ascii"), + "base64": base64.b64encode(compressed_bytes).decode("utf-8"), "mime_type": "image/jpeg", } @@ -371,6 +371,18 @@ def _is_docx_bytes(file_bytes: bytes) -> bool: return any(name.startswith("word/") for name in names) +def _is_epub_bytes(file_bytes: bytes) -> bool: + try: + with zipfile.ZipFile(io.BytesIO(file_bytes)) as archive: + names = set(archive.namelist()) + with archive.open("mimetype") as mimetype_file: + mimetype = mimetype_file.read(64).decode("utf-8").strip() + except (KeyError, OSError, UnicodeDecodeError, zipfile.BadZipFile): + return False + + return mimetype == "application/epub+zip" and "META-INF/container.xml" in names + + async def _parse_local_docx_text(file_bytes: bytes, file_name: str) -> str: from astrbot.core.knowledge_base.parsers.markitdown_parser import ( MarkitdownParser, @@ -387,23 +399,48 @@ async def _parse_local_pdf_text(file_bytes: bytes, file_name: str) -> str: return result.text +async def _parse_local_epub_text(file_bytes: bytes, file_name: str) -> str: + from astrbot.core.knowledge_base.parsers.epub_parser import EpubParser + + result = await EpubParser().parse(file_bytes, file_name) + return result.text + + async def _parse_local_supported_document( path: str, sample: bytes, ) -> ParsedDocument | None: file_name = Path(path).name + suffix = Path(path).suffix.lower() if _looks_like_pdf(path, sample): file_bytes = await _read_local_file_bytes(path) text = await _parse_local_pdf_text(file_bytes, file_name) return ParsedDocument(kind="pdf", file_bytes=file_bytes, text=text) - if Path(path).suffix.lower() == ".docx" or _looks_like_zip_container(sample): + if suffix == ".epub": + file_bytes = await _read_local_file_bytes(path) + if not _is_epub_bytes(file_bytes): + return None + text = await _parse_local_epub_text(file_bytes, file_name) + return ParsedDocument(kind="epub", file_bytes=file_bytes, text=text) + + if suffix == ".docx": file_bytes = await _read_local_file_bytes(path) if not _is_docx_bytes(file_bytes): return None text = await _parse_local_docx_text(file_bytes, file_name) return ParsedDocument(kind="docx", file_bytes=file_bytes, text=text) + if _looks_like_zip_container(sample): + file_bytes = await _read_local_file_bytes(path) + if _is_epub_bytes(file_bytes): + text = await _parse_local_epub_text(file_bytes, file_name) + return ParsedDocument(kind="epub", file_bytes=file_bytes, text=text) + if _is_docx_bytes(file_bytes): + text = await _parse_local_docx_text(file_bytes, file_name) + return ParsedDocument(kind="docx", file_bytes=file_bytes, text=text) + return None + return None @@ -659,14 +696,14 @@ async def read_file_tool_result( return "Error reading file: image payload is empty." raw_bytes = base64.b64decode(raw_base64_data) compressed_payload = await _compress_image_bytes_to_base64(raw_bytes) - base64_data = str(compressed_payload.get("base64", "") or "") - if not base64_data: + compressed_base64_data = str(compressed_payload.get("base64", "") or "") + if not compressed_base64_data: return "Error reading file: compressed image payload is empty." return mcp.types.CallToolResult( content=[ mcp.types.ImageContent( type="image", - data=base64_data, + data=compressed_base64_data, mimeType=str( compressed_payload.get("mime_type", "") or "image/jpeg", ), diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 76a2b846a..c46644c85 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -137,6 +137,7 @@ DEFAULT_CONFIG: dict[str, Any] = { "streaming_response": False, "show_tool_use_status": False, "show_tool_call_result": False, + "buffer_intermediate_messages": False, "sanitize_context_by_modalities": False, "max_quoted_fallback_images": 20, "quoted_message_parser": { @@ -2741,6 +2742,9 @@ CONFIG_METADATA_2: Any = { "show_tool_call_result": { "type": "bool", }, + "buffer_intermediate_messages": { + "type": "bool", + }, "unsupported_streaming_strategy": { "type": "string", }, @@ -3492,6 +3496,15 @@ CONFIG_METADATA_3 = { "provider_settings.show_tool_use_status": True, }, }, + "provider_settings.buffer_intermediate_messages": { + "description": "合并 Agent 中间消息", + "type": "bool", + "hint": "开启后，非流式模式下多步工具调用过程中产生的中间文本将缓冲，待 Agent 完成后合并为一条回复发送。", + "condition": { + "provider_settings.agent_runner_type": "local", + "provider_settings.streaming_response": False, + }, + }, "provider_settings.sanitize_context_by_modalities": { "description": "按模型能力清理历史上下文", "type": "bool", diff --git a/astrbot/core/db/__init__.py b/astrbot/core/db/__init__.py index c7d9b09c4..6528951c3 100644 --- a/astrbot/core/db/__init__.py +++ b/astrbot/core/db/__init__.py @@ -24,6 +24,7 @@ from astrbot.core.db.po import ( ProviderStat, SessionProjectRelation, Stats, + WebChatThread, ) @@ -206,10 +207,26 @@ class BaseDatabase(abc.ABC): content: dict, sender_id: str | None = None, sender_name: str | None = None, + llm_checkpoint_id: str | None = None, ) -> PlatformMessageHistory: """Insert a new platform message history record.""" ... + @abc.abstractmethod + async def update_platform_message_history( + self, + message_id: int, + content: dict | None = None, + llm_checkpoint_id: str | None = None, + ) -> None: + """Update a platform message history record.""" + ... + + @abc.abstractmethod + async def delete_platform_message_history_by_id(self, message_id: int) -> None: + """Delete a platform message history record by its ID.""" + ... + @abc.abstractmethod async def delete_platform_message_offset( self, @@ -290,6 +307,68 @@ class BaseDatabase(abc.ABC): """Get a platform message history record by its ID.""" ... + @abc.abstractmethod + async def create_webchat_thread( + self, + creator: str, + parent_session_id: str, + parent_message_id: int, + base_checkpoint_id: str, + selected_text: str, + ) -> WebChatThread: + """Create a WebChat side thread.""" + ... + + @abc.abstractmethod + async def get_webchat_thread_by_id( + self, + thread_id: str, + ) -> WebChatThread | None: + """Get a WebChat side thread by thread_id.""" + ... + + @abc.abstractmethod + async def get_webchat_threads_by_parent_session( + self, + parent_session_id: str, + creator: str | None = None, + ) -> list[WebChatThread]: + """Get side threads for a parent WebChat session.""" + ... + + @abc.abstractmethod + async def get_webchat_thread_by_parent_message_and_text( + self, + parent_session_id: str, + parent_message_id: int, + selected_text: str, + creator: str | None = None, + ) -> WebChatThread | None: + """Get an existing side thread for the same selected text.""" + ... + + @abc.abstractmethod + async def delete_webchat_thread(self, thread_id: str) -> None: + """Delete a WebChat side thread.""" + ... + + @abc.abstractmethod + async def delete_webchat_threads_by_parent_session( + self, + parent_session_id: str, + ) -> list[str]: + """Delete side threads for a parent WebChat session.""" + ... + + @abc.abstractmethod + async def delete_webchat_threads_by_parent_message_ids( + self, + parent_session_id: str, + parent_message_ids: list[int], + ) -> list[str]: + """Delete side threads linked to parent message IDs.""" + ... + @abc.abstractmethod async def insert_attachment( self, diff --git a/astrbot/core/db/po.py b/astrbot/core/db/po.py index 10e61c868..73c05bd76 100644 --- a/astrbot/core/db/po.py +++ b/astrbot/core/db/po.py @@ -244,6 +244,37 @@ class PlatformMessageHistory(TimestampMixin, SQLModel, table=True): default=None, ) # Name of the sender in the platform content: dict = Field(sa_type=JSON, nullable=False) # a message chain list + llm_checkpoint_id: str | None = Field(default=None, index=True) + + +class WebChatThread(TimestampMixin, SQLModel, table=True): + """A side thread created from a selected WebChat assistant response.""" + + __tablename__: str = "webchat_threads" + + id: int | None = Field( + primary_key=True, + sa_column_kwargs={"autoincrement": True}, + default=None, + ) + thread_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + ) + creator: str = Field(nullable=False, index=True) + parent_session_id: str = Field(nullable=False, index=True) + parent_message_id: int = Field(nullable=False, index=True) + base_checkpoint_id: str = Field(nullable=False, index=True) + selected_text: str = Field(sa_type=Text, nullable=False) + + __table_args__ = ( + UniqueConstraint( + "thread_id", + name="uix_webchat_thread_id", + ), + ) class PlatformSession(TimestampMixin, SQLModel, table=True): diff --git a/astrbot/core/db/sqlite.py b/astrbot/core/db/sqlite.py index 2e2eed682..d88a5f68a 100644 --- a/astrbot/core/db/sqlite.py +++ b/astrbot/core/db/sqlite.py @@ -26,6 +26,7 @@ from astrbot.core.db.po import ( ProviderStat, SessionProjectRelation, SQLModel, + WebChatThread, ) from astrbot.core.db.po import Platform as DeprecatedPlatformStat from astrbot.core.db.po import Stats as DeprecatedStats @@ -55,6 +56,7 @@ class SQLiteDatabase(BaseDatabase): await self._ensure_persona_folder_columns(conn) await self._ensure_persona_skills_column(conn) await self._ensure_persona_custom_error_message_column(conn) + await self._ensure_platform_message_history_checkpoint_column(conn) await conn.commit() async def _ensure_persona_folder_columns(self, conn) -> None: @@ -96,6 +98,30 @@ class SQLiteDatabase(BaseDatabase): text("ALTER TABLE personas ADD COLUMN custom_error_message TEXT"), ) + async def _ensure_platform_message_history_checkpoint_column(self, conn) -> None: + """Ensure platform_message_history has llm_checkpoint_id.""" + result = await conn.execute(text("PRAGMA table_info(platform_message_history)")) + columns = {row[1] for row in result.fetchall()} + + if "llm_checkpoint_id" not in columns: + await conn.execute( + text( + "ALTER TABLE platform_message_history " + "ADD COLUMN llm_checkpoint_id VARCHAR DEFAULT NULL" + ) + ) + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS " + "ix_platform_message_history_llm_checkpoint_id " + "ON platform_message_history (llm_checkpoint_id)" + ) + ) + + # ==== + # Platform Statistics + # ==== + async def insert_platform_stats( self, platform_id, @@ -436,18 +462,57 @@ class SQLiteDatabase(BaseDatabase): content, sender_id=None, sender_name=None, + llm_checkpoint_id=None, ): """Insert a new platform message history record.""" - async with self.get_db() as session, session.begin(): - new_history = PlatformMessageHistory( - platform_id=platform_id, - user_id=user_id, - content=content, - sender_id=sender_id, - sender_name=sender_name, - ) - session.add(new_history) - return new_history + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + new_history = PlatformMessageHistory( + platform_id=platform_id, + user_id=user_id, + content=content, + sender_id=sender_id, + sender_name=sender_name, + llm_checkpoint_id=llm_checkpoint_id, + ) + session.add(new_history) + return new_history + + async def update_platform_message_history( + self, + message_id: int, + content: dict | None = None, + llm_checkpoint_id: str | None = None, + ) -> None: + """Update a platform message history record.""" + values = {} + if content is not None: + values["content"] = content + if llm_checkpoint_id is not None: + values["llm_checkpoint_id"] = llm_checkpoint_id + if not values: + return + + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + await session.execute( + update(PlatformMessageHistory) + .where(PlatformMessageHistory.id == message_id) + .values(**values) + ) + + async def delete_platform_message_history_by_id(self, message_id: int) -> None: + """Delete a platform message history record by ID.""" + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + await session.execute( + delete(PlatformMessageHistory).where( + PlatformMessageHistory.id == message_id + ) + ) async def delete_platform_message_offset( self, @@ -594,6 +659,136 @@ class SQLiteDatabase(BaseDatabase): result = await session.execute(query) return result.scalar_one_or_none() + async def create_webchat_thread( + self, + creator: str, + parent_session_id: str, + parent_message_id: int, + base_checkpoint_id: str, + selected_text: str, + ) -> WebChatThread: + """Create a WebChat side thread.""" + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + thread = WebChatThread( + creator=creator, + parent_session_id=parent_session_id, + parent_message_id=parent_message_id, + base_checkpoint_id=base_checkpoint_id, + selected_text=selected_text, + ) + session.add(thread) + await session.flush() + await session.refresh(thread) + return thread + + async def get_webchat_thread_by_id( + self, + thread_id: str, + ) -> WebChatThread | None: + """Get a WebChat side thread by thread_id.""" + async with self.get_db() as session: + session: AsyncSession + result = await session.execute( + select(WebChatThread).where(WebChatThread.thread_id == thread_id) + ) + return result.scalar_one_or_none() + + async def get_webchat_threads_by_parent_session( + self, + parent_session_id: str, + creator: str | None = None, + ) -> list[WebChatThread]: + """Get side threads for a parent WebChat session.""" + async with self.get_db() as session: + session: AsyncSession + query = select(WebChatThread).where( + WebChatThread.parent_session_id == parent_session_id + ) + if creator is not None: + query = query.where(WebChatThread.creator == creator) + query = query.order_by(WebChatThread.created_at) + result = await session.execute(query) + return list(result.scalars().all()) + + async def get_webchat_thread_by_parent_message_and_text( + self, + parent_session_id: str, + parent_message_id: int, + selected_text: str, + creator: str | None = None, + ) -> WebChatThread | None: + """Get an existing side thread for the same selected text.""" + async with self.get_db() as session: + session: AsyncSession + query = select(WebChatThread).where( + WebChatThread.parent_session_id == parent_session_id, + WebChatThread.parent_message_id == parent_message_id, + WebChatThread.selected_text == selected_text, + ) + if creator is not None: + query = query.where(WebChatThread.creator == creator) + result = await session.execute(query) + return result.scalar_one_or_none() + + async def delete_webchat_thread(self, thread_id: str) -> None: + """Delete a WebChat side thread.""" + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + await session.execute( + delete(WebChatThread).where(WebChatThread.thread_id == thread_id) + ) + + async def delete_webchat_threads_by_parent_session( + self, + parent_session_id: str, + ) -> list[str]: + """Delete side threads for a parent WebChat session.""" + threads = await self.get_webchat_threads_by_parent_session(parent_session_id) + thread_ids = [thread.thread_id for thread in threads] + if not thread_ids: + return [] + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + await session.execute( + delete(WebChatThread).where( + col(WebChatThread.thread_id).in_(thread_ids) + ) + ) + return thread_ids + + async def delete_webchat_threads_by_parent_message_ids( + self, + parent_session_id: str, + parent_message_ids: list[int], + ) -> list[str]: + """Delete side threads linked to parent message IDs.""" + if not parent_message_ids: + return [] + async with self.get_db() as session: + session: AsyncSession + result = await session.execute( + select(WebChatThread.thread_id).where( + WebChatThread.parent_session_id == parent_session_id, + col(WebChatThread.parent_message_id).in_(parent_message_ids), + ) + ) + thread_ids = list(result.scalars().all()) + if not thread_ids: + return [] + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + await session.execute( + delete(WebChatThread).where( + col(WebChatThread.thread_id).in_(thread_ids) + ) + ) + return thread_ids + async def insert_attachment(self, path, type, mime_type): """Insert a new attachment record.""" async with self.get_db() as session, session.begin(): diff --git a/astrbot/core/db/vec_db/faiss_impl/document_storage.py b/astrbot/core/db/vec_db/faiss_impl/document_storage.py index 8eb0f48af..83be01d6d 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -102,36 +102,29 @@ class DocumentStorage: async def _initialize_fts5(self, executor) -> None: try: - try: - await executor.execute( - text( - f""" - CREATE VIRTUAL TABLE IF NOT EXISTS {FTS_TABLE_NAME} - USING fts5( - search_text, - content='', - contentless_delete=1, - tokenize='unicode61' - ) - """, - ), + await self._create_fts5_table(executor, if_not_exists=True) + + is_valid_fts5, has_contentless_delete = await self._inspect_fts5_table( + executor, + ) + if not is_valid_fts5: + logger.warning( + f"Detected incompatible legacy table `{FTS_TABLE_NAME}` in " + f"{self.db_path}; recreating FTS5 table.", ) - self._fts_contentless_delete = True - except Exception: - await executor.execute( - text( - f""" - CREATE VIRTUAL TABLE IF NOT EXISTS {FTS_TABLE_NAME} - USING fts5( - search_text, - content='', - tokenize='unicode61' - ) - """, - ), + await executor.execute(text(f"DROP TABLE IF EXISTS {FTS_TABLE_NAME}")) + await self._create_fts5_table(executor, if_not_exists=False) + + is_valid_fts5, has_contentless_delete = await self._inspect_fts5_table( + executor, ) - self._fts_contentless_delete = False + if not is_valid_fts5: + raise RuntimeError( + f"Failed to create a valid FTS5 table `{FTS_TABLE_NAME}`", + ) + self.fts5_available = True + self._fts_contentless_delete = has_contentless_delete except Exception as e: self.fts5_available = False self._fts_contentless_delete = False @@ -140,6 +133,69 @@ class DocumentStorage: f"falling back to in-memory BM25 sparse retrieval: {e}", ) + async def _create_fts5_table(self, executor, if_not_exists: bool) -> None: + create_clause = ( + "CREATE VIRTUAL TABLE IF NOT EXISTS" + if if_not_exists + else "CREATE VIRTUAL TABLE" + ) + try: + await executor.execute( + text( + f""" + {create_clause} {FTS_TABLE_NAME} + USING fts5( + search_text, + content='', + contentless_delete=1, + tokenize='unicode61' + ) + """, + ), + ) + except Exception: + await executor.execute( + text( + f""" + {create_clause} {FTS_TABLE_NAME} + USING fts5( + search_text, + content='', + tokenize='unicode61' + ) + """, + ), + ) + + async def _inspect_fts5_table(self, executor) -> tuple[bool, bool]: + schema_result = await executor.execute( + text( + """ + SELECT sql + FROM sqlite_master + WHERE type='table' AND name=:table_name + """, + ), + {"table_name": FTS_TABLE_NAME}, + ) + create_sql = schema_result.scalar_one_or_none() + if not create_sql: + return False, False + + normalized_sql = create_sql.lower() + if "virtual table" not in normalized_sql or "using fts5" not in normalized_sql: + return False, False + + pragma_result = await executor.execute( + text(f"PRAGMA table_info({FTS_TABLE_NAME})"), + ) + columns = {row[1] for row in pragma_result.fetchall()} + if "search_text" not in columns: + return False, False + + normalized_sql_no_whitespace = "".join(normalized_sql.split()) + return True, "contentless_delete=1" in normalized_sql_no_whitespace + async def connect(self) -> None: """Connect to the SQLite database.""" if self.engine is None: diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py index 52691ec53..a255c895b 100644 --- a/astrbot/core/knowledge_base/kb_helper.py +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -102,6 +102,10 @@ Text chunk to process: return [chunk] +def _compact_chunks(chunks: list[str]) -> list[str]: + return [chunk.strip() for chunk in chunks if chunk and chunk.strip()] + + class KBHelper: vec_db: FaissVecDB | None kb: KnowledgeBase @@ -250,8 +254,8 @@ class KBHelper: saved_media: list[KBMedia] = [] if pre_chunked_text is not None: - # 如果提供了预分块文本,直接使用 - chunks_text = pre_chunked_text + # 如果提供了预分块文本，直接使用 + chunks_text = _compact_chunks(pre_chunked_text) file_size = sum(len(chunk) for chunk in chunks_text) logger.info(f"使用预分块文本进行上传,共 {len(chunks_text)} 个块｡") else: @@ -294,11 +298,41 @@ class KBHelper: if progress_callback: await progress_callback("chunking", 0, 100) - chunks_text = await self.chunker.chunk( - text_content, - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, - ) + try: + chunks_text = await self.chunker.chunk( + text_content, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + chunks_text = _compact_chunks(chunks_text) + except KnowledgeBaseUploadError: + raise + except Exception as exc: + raise KnowledgeBaseUploadError( + stage="chunking", + user_message=( + "分块失败：文档内容在切分文本块时发生错误。" + "请稍后重试，或调整分块参数后再次上传。" + ), + details={"file_name": file_name}, + ) from exc + + if not chunks_text or not any(chunk.strip() for chunk in chunks_text): + if pre_chunked_text is not None: + raise KnowledgeBaseUploadError( + stage="validation", + user_message=("预分块文本为空，未提供任何可索引文本块。"), + details={"file_name": file_name}, + ) + else: + raise KnowledgeBaseUploadError( + stage="chunking", + user_message=( + "分块失败：文档内容为空，未生成任何可索引文本块。" + ), + details={"file_name": file_name}, + ) + contents = [] metadatas = [] for idx, chunk_text in enumerate(chunks_text): @@ -657,6 +691,8 @@ class KBHelper: elif isinstance(result, list): final_chunks.extend(result) + final_chunks = _compact_chunks(final_chunks) + logger.info( f"文本修复完成: {len(initial_chunks)} 个原始块 -> {len(final_chunks)} 个最终块｡", ) diff --git a/astrbot/core/knowledge_base/parsers/__init__.py b/astrbot/core/knowledge_base/parsers/__init__.py index 184f2fd41..f3b5967cf 100644 --- a/astrbot/core/knowledge_base/parsers/__init__.py +++ b/astrbot/core/knowledge_base/parsers/__init__.py @@ -1,11 +1,13 @@ """文档解析器模块""" from .base import BaseParser, MediaItem, ParseResult +from .epub_parser import EpubParser from .pdf_parser import PDFParser from .text_parser import TextParser __all__ = [ "BaseParser", + "EpubParser", "MediaItem", "PDFParser", "ParseResult", diff --git a/astrbot/core/knowledge_base/parsers/epub_parser.py b/astrbot/core/knowledge_base/parsers/epub_parser.py new file mode 100644 index 000000000..159f6c736 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/epub_parser.py @@ -0,0 +1,162 @@ +"""EPUB document parser.""" + +import html +import re + +from astrbot.core.knowledge_base.parsers.base import BaseParser, ParseResult + +_KEYS = ( + "Title|Author|Creator|Language|Publisher|Date|Modified|Identifier|ISBN|Description|" + "Subject|Rights|Source|Series|标题|书名|作者|语言|出版社|日期|出版日期|标识符|简介|描述|" + "主题|版权|来源|系列|タイトル|書名|著者|言語|出版社|日付|識別子|説明|件名|権利|ソース|シリーズ" +) +_META_RE = re.compile(rf"^\s*(?:[-*]\s*)?\*\*(?:{_KEYS})\s*[:：]\*\*\s+\S") +_TOC_HEAD_RE = re.compile( + r"^\s{0,3}(?:#{1,6}\s*)?(?:table of contents|contents|toc|目录|目次|もくじ)\s*$", + re.I, +) +_LINK_RE = re.compile(r"(? str: + return ( + html.unescape(s) + .replace("\r\n", "\n") + .replace("\r", "\n") + .replace("\ufeff", "") + .replace("\u00a0", " ") + .replace("\u200b", "") + ) + + +def _is_internal(href: str) -> bool: + href = html.unescape(href).strip().lower() + return ( + href.startswith("#") + or href.endswith(".html") + or href.endswith(".xhtml") + or ".html#" in href + or ".xhtml#" in href + ) + + +def _is_toc_line(s: str) -> bool: + s = s.strip() + if not s: + return False + s = re.sub(r"^\s*(?:[-*+]|\d+\.)\s+", "", s) + m = re.fullmatch(r"\[([^\]]+)\]\(([^)]+)\)", s) + return bool((m and _is_internal(m.group(2))) or _DOTTED_TOC_RE.match(s)) + + +def _strip_head(text: str) -> str: + lines = _n(text).split("\n") + i = 0 + while i < len(lines) and not lines[i].strip(): + i += 1 + start = i + while i < len(lines) and _META_RE.match(lines[i].strip()): + i += 1 + if i - start >= 2: + while i < len(lines) and not lines[i].strip(): + i += 1 + else: + i = start + toc0, had_head = i, False + if i < len(lines) and _TOC_HEAD_RE.match(lines[i].strip()): + had_head = True + i += 1 + while i < len(lines) and not lines[i].strip(): + i += 1 + toc = 0 + while i < len(lines) and i - toc0 < 120: + s = lines[i].strip() + if not s: + if toc and i + 1 < len(lines) and _is_toc_line(lines[i + 1]): + i += 1 + continue + break + if not _is_toc_line(s): + break + toc += 1 + i += 1 + if toc >= 2 and (had_head or toc >= 3): + while i < len(lines) and not lines[i].strip(): + i += 1 + return "\n".join(lines[i:]).strip() + return "\n".join(lines[toc0:]).strip() + + +def _strip_links(text: str) -> str: + def repl(m: re.Match[str]) -> str: + label = html.unescape(m.group(1)).strip() + href = html.unescape(m.group(2)).strip().lower() + if not _is_internal(href): + return m.group(0) + if _FOOTNOTE_HREF_RE.search(href) or ( + href.startswith("#") and _FOOTNOTE_LABEL_RE.fullmatch(label) + ): + return "" + return label + + return _LINK_RE.sub(repl, _n(text)) + + +def _img_alt(m: re.Match[str]) -> str: + alt = re.sub(r"\s+", " ", html.unescape(m.group(1)).strip()) + if not alt or _GENERIC_ALT_RE.fullmatch(alt) or _FILENAME_ALT_RE.fullmatch(alt): + return "" + return alt + + +def _sanitize(text: str) -> str: + out, prev_blank, prev = [], True, "" + for raw in _n(text).split("\n"): + line = _IMG_RE.sub(_img_alt, raw) + line = _EMPTY_IMG_LINK_RE.sub("", line).rstrip() + s = line.strip() + if not s: + if not prev_blank: + out.append("") + prev_blank = True + continue + if _SEP_RE.match(s) or _NOISE_RE.match(s): + continue + norm = re.sub(r"^\s{0,3}#{1,6}\s*", "", s).strip("*_ ").casefold() + if norm and norm == prev and len(norm) <= 120: + continue + out.append(line) + prev_blank = False + prev = norm + return "\n".join(out).strip() + + +class EpubParser(BaseParser): + """Parse EPUB files via MarkItDown.""" + + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + from .markitdown_parser import MarkitdownParser + + result = await MarkitdownParser().parse(file_content, file_name) + text = _sanitize(_strip_links(_strip_head(result.text))) + return ParseResult(text=text, media=result.media) diff --git a/astrbot/core/knowledge_base/parsers/util.py b/astrbot/core/knowledge_base/parsers/util.py index 7a4463202..52cffa5cd 100644 --- a/astrbot/core/knowledge_base/parsers/util.py +++ b/astrbot/core/knowledge_base/parsers/util.py @@ -6,6 +6,10 @@ async def select_parser(ext: str) -> BaseParser: from .markitdown_parser import MarkitdownParser return MarkitdownParser() + if ext == ".epub": + from .epub_parser import EpubParser + + return EpubParser() if ext == ".pdf": from .pdf_parser import PDFParser diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py index 04e0cdd76..b5d8099bc 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py @@ -5,8 +5,13 @@ import base64 from collections.abc import AsyncGenerator from dataclasses import replace -from astrbot.core import logger -from astrbot.core.agent.message import Message +from astrbot.core import db_helper, logger +from astrbot.core.agent.message import ( + CheckpointData, + CheckpointMessageSegment, + Message, + dump_messages_with_checkpoints, +) from astrbot.core.agent.response import AgentStats from astrbot.core.astr_agent_run_util import AgentRunner, run_agent, run_live_agent from astrbot.core.astr_main_agent import ( @@ -67,6 +72,10 @@ class InternalAgentSubStage(Stage): self.max_step = 30 self.show_tool_use: bool = settings.get("show_tool_use_status", True) self.show_tool_call_result: bool = settings.get("show_tool_call_result", False) + self.buffer_intermediate_messages: bool = settings.get( + "buffer_intermediate_messages", + False, + ) self.show_reasoning = settings.get("display_reasoning_text", False) self.sanitize_context_by_modalities: bool = settings.get( "sanitize_context_by_modalities", @@ -286,6 +295,7 @@ class InternalAgentSubStage(Stage): self.show_tool_use, self.show_tool_call_result, show_reasoning=self.show_reasoning, + buffer_intermediate_messages=self.buffer_intermediate_messages, ), ), ) @@ -312,6 +322,7 @@ class InternalAgentSubStage(Stage): self.show_tool_use, self.show_tool_call_result, show_reasoning=self.show_reasoning, + buffer_intermediate_messages=self.buffer_intermediate_messages, ), ), ) @@ -342,6 +353,7 @@ class InternalAgentSubStage(Stage): self.show_tool_call_result, stream_to_general, show_reasoning=self.show_reasoning, + buffer_intermediate_messages=self.buffer_intermediate_messages, ): yield None final_resp = agent_runner.get_final_llm_resp() @@ -433,7 +445,8 @@ class InternalAgentSubStage(Stage): ): logger.debug("LLM 响应为空,不保存记录｡") return - message_to_save = [] + + messages_to_save: list[Message] = [] skipped_initial_system = False for message in all_messages: if message.role == "system" and (not skipped_initial_system): @@ -441,7 +454,25 @@ class InternalAgentSubStage(Stage): continue if message.role in ["assistant", "user"] and message._no_save: continue - message_to_save.append(message.model_dump()) + messages_to_save.append(message) + + checkpoint_id = event.get_extra("llm_checkpoint_id") + message_to_save = dump_messages_with_checkpoints(messages_to_save) + if isinstance(checkpoint_id, str) and checkpoint_id: + message_to_save.append( + CheckpointMessageSegment( + content=CheckpointData(id=checkpoint_id), + ).model_dump() + ) + + # if user_aborted: + # message_to_save.append( + # Message( + # role="assistant", + # content="[User aborted this request. Partial output before abort was preserved.]", + # ).model_dump() + # ) + token_usage = None if runner_stats: token_usage = llm_response.usage.total if llm_response.usage else None diff --git a/astrbot/core/platform/sources/kook/kook_adapter.py b/astrbot/core/platform/sources/kook/kook_adapter.py index 3fdeb5467..1632df587 100644 --- a/astrbot/core/platform/sources/kook/kook_adapter.py +++ b/astrbot/core/platform/sources/kook/kook_adapter.py @@ -13,12 +13,13 @@ from astrbot.api.platform import ( PlatformMetadata, register_platform_adapter, ) -from astrbot.core.message.components import File, Record, Video +from astrbot.core.message.components import BaseMessageComponent, File, Record, Video from astrbot.core.platform.astr_message_event import MessageSesion from .kook_client import KookClient from .kook_config import KookConfig from .kook_event import KookEvent +from .kook_roles_record import KookRolesRecord from .kook_types import ( ContainerModule, FileModule, @@ -27,14 +28,18 @@ from .kook_types import ( KmarkdownElement, KookCardMessageContainer, KookChannelType, + KookMarkdownMentionRolePart, + KookMentionTagName, KookMessageEventData, KookMessageType, KookModuleType, + KookRoleExtraType, PlainTextElement, SectionModule, ) -KOOK_AT_SELECTOR_REGEX = re.compile(r"\(met\)([^()]+)\(met\)") +KOOK_AT_SELECTOR_REGEX = re.compile(r"\((met|rol)\)([^()]+)\(\1\)") +AT_MENTION_PREFIX_REGEX = re.compile(r"^@[^\s]+(\s*-\s*[^\s]+)?\s*") @register_platform_adapter( @@ -56,6 +61,7 @@ class KookPlatformAdapter(Platform): self._reconnect_task = None self.running = False self._main_task = None + self._roles_cache = KookRolesRecord("", self.client.http_client) async def send_by_session( self, @@ -91,7 +97,7 @@ class KookPlatformAdapter(Platform): event_type = event.type if event_type in (KookMessageType.KMARKDOWN, KookMessageType.CARD): if self._should_ignore_event_by_bot_nickname(event.author_id): - logger.debug("[KOOK] 收到来自机器人自身的消息, 忽略此消息") + logger.debug("[KOOK] 判断此消息为来自机器人自身的消息, 忽略此消息") return try: abm = await self.convert_message(event) @@ -99,8 +105,18 @@ class KookPlatformAdapter(Platform): except Exception as e: logger.error(f"[KOOK] 消息处理异常: {e}") elif event_type == KookMessageType.SYSTEM: - logger.debug(f'[KOOK] 消息为系统通知, 通知类型为: "{event.extra.type}"') - logger.debug(f"[KOOK] 原始消息数据: {event.to_json()}") + match event.extra.type: + case KookRoleExtraType(): + # 此时 target_id 就是频道id(guild_id) + guild_id = event.target_id + logger.info( + f'[KOOK] 收到频道"{guild_id}"的角色更新通知, 类型为"{event.extra.type.value}", 刷新角色id缓存' + ) + self._roles_cache.clear_guild_roles_cache(int(guild_id)) + case _: + logger.debug( + f'[KOOK] 判断此消息为"{event.extra.type}"类型的系统通知, 因未实现此消息的处理流程而忽略此消息, 原始消息数据: {event.to_json()}' + ) async def run(self): """主运行循环""" @@ -131,6 +147,8 @@ class KookPlatformAdapter(Platform): logger.info("[KOOK] 尝试连接KOOK服务器...") # 尝试连接 + await self.client.get_bot_info() + self._roles_cache.set_bot_id(self.client.bot_id) success = await self.client.connect() if success: @@ -200,49 +218,86 @@ class KookPlatformAdapter(Platform): logger.info("[KOOK] 资源清理完成") - def _parse_kmarkdown_text_message( + async def _convert_text_message_to_component( self, - data: KookMessageEventData, - self_id: str, - ) -> tuple[list, str]: - kmarkdown = data.extra.kmarkdown - content = data.content or "" - if kmarkdown is None: - logger.error( - f'[KOOK] 无法转换"{KookMessageType.KMARKDOWN.name}"消息, 消息中找不到kmarkdown字段', - ) - logger.error(f"[KOOK] 原始消息内容: {data.to_json()}") - return [], "" + content: str, + raw_content: str, + mention_role_part: list[KookMarkdownMentionRolePart] | None = None, + guild_id: str | None = None, + mention_name_map: dict[str, str] | None = None, + ) -> tuple[list[BaseMessageComponent], str]: + # kook平台有一个角色(role)的概念,他表示拥有某一类权限的许多用户 + # 且角色本身也有一个自己的id,与正常用户id不同 + # 而在频道中是可以`@`角色的,而想要知道bot是否属于某个角色 + # 需要通过 `/user/view` 接口获取当前bot账号的某个频道下所属角色的id + # 为了解决 https://github.com/AstrBotDevs/AstrBot/issues/7539 + # 在确定机器人需要响应某个`(rol)xxx(rol)`时,需要将角色id替换装当前的bot id + # 包装成`At`机器人自己,而`At`的name就保留角色名称 + # 如果没有查询到角色id或者bot不属于某类角色, 则不处理此`(rol)xxx(rol)` + # 暂时想不到能在不修改原有消息内容的情况下处理这个角色mention的方案 - raw_content = kmarkdown.raw_content or content - if not isinstance(content, str): - content = str(content) - if not isinstance(raw_content, str): - raw_content = str(raw_content) - - # TODO 后面的pydantic类型替换,以后再来探索吧 :( - mention_name_map: dict[str, str] = {} - mention_part = kmarkdown.mention_part - if isinstance(mention_part, list): - for item in mention_part: - if not isinstance(item, dict): - continue - mention_id = item.get("id") - if mention_id is None: - continue - mention_name_map[str(mention_id)] = str(item.get("username", "")) - - components = [] + message_str = raw_content + bot_id = self.client.bot_id + bot_nickname = self.client.bot_nickname + bot_username = self.client.bot_username + components: list[BaseMessageComponent] = [] + if mention_name_map is None: + mention_name_map = {} cursor = 0 + + role_mention_counter = -1 + for match in KOOK_AT_SELECTOR_REGEX.finditer(content): if match.start() > cursor: - plain_text = content[cursor : match.start()] + plain_text = content[cursor : match.start()].strip(" ") if plain_text: components.append(Plain(text=plain_text)) - mention_target = match.group(1).strip() - if mention_target == "all": + tag_name = match.group(1) + mention_target = match.group(2).strip() + if tag_name == KookMentionTagName.MENTION and mention_target == "all": components.append(AtAll()) + elif tag_name == KookMentionTagName.ROLE: + role_mention_counter += 1 + role_id = 0 + role_mention_name = mention_target + if mention_role_part is not None: + if len(mention_role_part) > role_mention_counter: + role_mention_name = mention_role_part[role_mention_counter].name + role_id = mention_role_part[role_mention_counter].role_id + if ( + bot_nickname == role_mention_name + or bot_username == role_mention_name + ): + components.append( + At( + qq=bot_id, + name=role_mention_name, # 保留角色名称 + ) + ) + continue + if not mention_target.isdigit() and role_id == 0: + continue + + role_id = role_id or int(mention_target) + if not guild_id: + continue + + if not guild_id.isdigit(): + continue + + if not await self._roles_cache.has_role_in_channel( + role_id, int(guild_id) + ): + continue + + components.append( + At( + qq=bot_id, + name=role_mention_name, # 保留角色名称 + ) + ) + elif mention_target: components.append( At( @@ -253,11 +308,11 @@ class KookPlatformAdapter(Platform): cursor = match.end() if cursor < len(content): - tail_text = content[cursor:] + tail_text = content[cursor:].strip(" ") if tail_text: components.append(Plain(text=tail_text)) - message_str = raw_content + message_str = raw_content.strip() if components: for comp in components: if isinstance(comp, Plain): @@ -265,9 +320,8 @@ class KookPlatformAdapter(Platform): continue break if isinstance(comp, At): - if str(comp.qq) == str(self_id): - message_str = re.sub( - r"^@[^\s]+(\s*-\s*[^\s]+)?\s*", + if str(comp.qq) == str(self.client.bot_id): + message_str = AT_MENTION_PREFIX_REGEX.sub( "", message_str, count=1, @@ -281,10 +335,44 @@ class KookPlatformAdapter(Platform): return components, message_str - def _parse_card_message(self, data: KookMessageEventData) -> tuple[list, str]: + async def _parse_kmarkdown_message( + self, data: KookMessageEventData + ) -> tuple[list[BaseMessageComponent], str]: + kmarkdown = data.extra.kmarkdown + guild_id = data.extra.guild_id + mention_role_part = None + if kmarkdown: + mention_role_part = kmarkdown.mention_role_part + # 无法处理可能会收到的道具消息content,只能保留原样 + content = str(data.content) or "" + if kmarkdown is None: + logger.error( + f'[KOOK] 无法转换"{KookMessageType.KMARKDOWN.name}"消息, 消息中找不到kmarkdown字段' + ) + logger.error(f"[KOOK] 原始消息内容: {data.to_json()}") + return [], "" + + raw_content = kmarkdown.raw_content or content + + mention_name_map: dict[str, str] = {} + mention_part = kmarkdown.mention_part + for item in mention_part: + mention_id = item.id + if mention_id is None: + continue + mention_name_map[str(mention_id)] = str(item.username) + + return await self._convert_text_message_to_component( + content, raw_content, mention_role_part, guild_id, mention_name_map + ) + + async def _parse_card_message( + self, data: KookMessageEventData + ) -> tuple[list[BaseMessageComponent], str]: content = data.content if not isinstance(content, str): content = str(content) + guild_id = data.extra.guild_id card_list = KookCardMessageContainer.from_dict(json.loads(content)) @@ -315,18 +403,13 @@ class KookPlatformAdapter(Platform): logger.debug(f"[KOOK] 跳过或未处理模块: {module.type}") text = "".join(text_parts) - message = [] + message: list[BaseMessageComponent] = [] if text: - for search in KOOK_AT_SELECTOR_REGEX.finditer(text): - search_text = search.group(1).strip() - if search_text == "all": - message.append(AtAll()) - continue - message.append(At(qq=search_text)) - text = text.replace(f"(met){search_text}(met)", "") - - message.append(Plain(text=text)) + component_parts, text = await self._convert_text_message_to_component( + text, text, guild_id=guild_id + ) + message.extend(component_parts) for img_url in images: message.append(Image(file=img_url)) @@ -399,12 +482,10 @@ class KookPlatformAdapter(Platform): abm.message_id = data.msg_id or "unknown" if data.type == KookMessageType.KMARKDOWN: - message, message_str = self._parse_kmarkdown_text_message(data, abm.self_id) - abm.message = message - abm.message_str = message_str + abm.message, abm.message_str = await self._parse_kmarkdown_message(data) elif data.type == KookMessageType.CARD: try: - abm.message, abm.message_str = self._parse_card_message(data) + abm.message, abm.message_str = await self._parse_card_message(data) except Exception as exp: logger.error(f"[KOOK] 卡片消息解析失败: {exp}") logger.error(f"[KOOK] 原始消息内容: {data.to_json()}") diff --git a/astrbot/core/platform/sources/kook/kook_client.py b/astrbot/core/platform/sources/kook/kook_client.py index da6f3e3b6..8a2c0ee50 100644 --- a/astrbot/core/platform/sources/kook/kook_client.py +++ b/astrbot/core/platform/sources/kook/kook_client.py @@ -2,6 +2,7 @@ import asyncio import base64 import random import time +import traceback import zlib import aiohttp @@ -57,12 +58,18 @@ class KookClient: @property def bot_nickname(self): + """机器人昵称""" return self._bot_nickname @property def bot_username(self): + """机器人名称""" return self._bot_username + @property + def http_client(self): + return self._http_client + async def get_bot_info(self) -> None: """获取机器人账号信息""" url = KookApiPaths.USER_ME @@ -153,7 +160,6 @@ class KookClient: sn=self.last_sn, session_id=self.session_id, ) - await self.get_bot_info() if not gateway_url: return False @@ -217,8 +223,8 @@ class KookClient: except websockets.exceptions.ConnectionClosed: logger.warning("[KOOK] WebSocket连接已关闭") break - except Exception as e: - logger.error(f"[KOOK] 消息处理异常: {e}") + except Exception: + logger.error(f"[KOOK] 消息处理异常: {traceback.format_exc()}") break except Exception as e: @@ -238,11 +244,15 @@ class KookClient: await self.event_callback(data) case KookMessageSignal.HELLO: - assert isinstance(data, KookHelloEventData) + assert isinstance(data, KookHelloEventData), ( + f"期望 data 为 {KookHelloEventData.__name__}, 实际为 {type(data).__name__}，" + ) await self._handle_hello(data) case KookMessageSignal.RESUME_ACK: - assert isinstance(data, KookResumeAckEventData) + assert isinstance(data, KookResumeAckEventData), ( + f"期望 data 为 {KookResumeAckEventData.__name__}, 实际为 {type(data).__name__}，" + ) await self._handle_resume_ack(data) case KookMessageSignal.PONG: diff --git a/astrbot/core/platform/sources/kook/kook_roles_record.py b/astrbot/core/platform/sources/kook/kook_roles_record.py new file mode 100644 index 000000000..fca9660ca --- /dev/null +++ b/astrbot/core/platform/sources/kook/kook_roles_record.py @@ -0,0 +1,164 @@ +import asyncio +import time +from collections import OrderedDict +from dataclasses import dataclass + +import aiohttp +import pydantic + +from astrbot import logger + +from .kook_types import KookApiPaths, KookUserViewResponse + +USER_VIEW_REQUEST_TIMEOUT = aiohttp.ClientTimeout(total=3) +ROLES_CACHE_MAX_SIZE = 2000 +MAX_RETRY_TIMES = 3 +RETRY_INTERVAL_SECOND = 1 * 60 + + +@dataclass +class RolesCache: + value: set[int] | None = None + failed_count: int = 0 + latest_update_time: float = 0 + + def update(self, roles: set[int] | None) -> None: + if roles is not None: + self.failed_count = 0 + self.value = roles + self.latest_update_time = time.time() + + def add_failed(self): + self.failed_count += 1 + + def reset(self, without_value=False): + if not without_value: + self.value = None + self.failed_count = 0 + self.latest_update_time = 0 + + +class KookRolesRecord: + """自动和缓存获取机器人所需响应的消息频道的role信息""" + + def __init__(self, bot_id: str, http_client: aiohttp.ClientSession): + # self._locks: dict[int, asyncio.Lock] = defaultdict(asyncio.Lock) + self._lock = asyncio.Lock() + self._bot_id = bot_id + self._http_client = http_client + # TODO 这个些配置后续加到适配器配置项里 + self._cache_max_size = ROLES_CACHE_MAX_SIZE + self._max_retry_times = MAX_RETRY_TIMES + self._retry_interval = RETRY_INTERVAL_SECOND + self._roles_cache: OrderedDict[int, RolesCache] = OrderedDict() + self._pending_tasks: dict[int, asyncio.Future] = {} + + def set_bot_id(self, bot_id: str): + self._bot_id = bot_id + + def clear_guild_roles_cache(self, guild_id: int): + self._roles_cache.pop(guild_id, None) + self._pending_tasks.pop(guild_id, None) + + async def _fetch_roles_by_guild_id(self, guild_id: int) -> set[int] | None: + # 由于需要判断bot账号是属于某个角色(role)才会回复消息, + # 而后续来自同一个频道的消息,在第一次查这个role的时候, + # 会一直阻塞消息接收直到请求完成或者报错, + # 所以,这里特意调低了timeout时间,避免阻塞太久 + url = KookApiPaths.USER_VIEW + try: + async with self._http_client.get( + url, + params={ + "guild_id": guild_id, + "user_id": self._bot_id, + }, + # TODO 这个超时时间后续加到适配器配置项里 + timeout=USER_VIEW_REQUEST_TIMEOUT, + ) as resp: + if resp.status != 200: + logger.error( + f'[KOOK] 获取机器人在频道"{guild_id}"的角色id信息失败，状态码: {resp.status} , {await resp.text()}' + ) + return + try: + resp_content = KookUserViewResponse.from_dict(await resp.json()) + except pydantic.ValidationError as e: + logger.error( + f'[KOOK] 获取机器人在频道"{guild_id}"的角色id信息失败, 响应数据格式错误: \n{e}' + ) + logger.error(f"[KOOK] 响应内容: {await resp.text()}") + return + + if not resp_content.success(): + logger.error( + f'[KOOK] 获取机器人在频道"{guild_id}"的角色id信息失败: {resp_content.model_dump_json()}' + ) + return + + logger.info(f'[KOOK] 获取机器人在频道"{guild_id}"的角色id成功') + return set(resp_content.data.roles) + + except Exception as e: + logger.error( + f'[KOOK] 获取机器人在频道"{guild_id}"的角色id信息时请求异常: {e}' + ) + return + + async def has_role_in_channel(self, role_id: int, guild_id: int) -> bool: + if (cache := self._roles_cache.get(guild_id)) is not None: + self._roles_cache.move_to_end(guild_id) + roles = cache.value + if roles is not None: + return role_id in roles + + new_future: asyncio.Future[set[int] | None] = asyncio.Future() + actual_future: asyncio.Future[set[int] | None] = self._pending_tasks.setdefault( + guild_id, new_future + ) + + if actual_future is not new_future: + roles = await actual_future + if roles is None: + return False + return role_id in roles + + try: + if (cache := self._roles_cache.get(guild_id)) is not None: + if ( + cache.failed_count > self._max_retry_times + and time.time() - cache.latest_update_time < self._retry_interval + ): + new_future.set_result(None) + return False + + # 简单的容量控制 (LRU) + if len(self._roles_cache) + 1 > self._cache_max_size: + self._roles_cache.popitem(last=False) + + roles_set = await self._fetch_roles_by_guild_id(guild_id) + + cache = self._roles_cache.get(guild_id) + if cache is not None: + cache.update(roles_set) + self._roles_cache.move_to_end(guild_id) + else: + cache = RolesCache(roles_set, latest_update_time=time.time()) + self._roles_cache[guild_id] = cache + + result = False + if roles_set is None: + cache.add_failed() + else: + result = role_id in roles_set + + new_future.set_result(roles_set) + return result + except Exception as e: + new_future.set_result(None) + logger.error( + f'[KOOK] 获取机器人在频道"{guild_id}"的角色id信息时发生异常: {e}' + ) + return False + finally: + self._pending_tasks.pop(guild_id, None) diff --git a/astrbot/core/platform/sources/kook/kook_types.py b/astrbot/core/platform/sources/kook/kook_types.py index 5faada4bb..93e32da7c 100644 --- a/astrbot/core/platform/sources/kook/kook_types.py +++ b/astrbot/core/platform/sources/kook/kook_types.py @@ -2,7 +2,7 @@ import json from enum import Enum, IntEnum from typing import Annotated, Any, Literal -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator class KookApiPaths: @@ -13,6 +13,7 @@ class KookApiPaths: # 初始化相关 USER_ME = f"{BASE_URL}{API_VERSION_PATH}/user/me" + USER_VIEW = f"{BASE_URL}{API_VERSION_PATH}/user/view" GATEWAY_INDEX = f"{BASE_URL}{API_VERSION_PATH}/gateway/index" # 消息相关 @@ -23,6 +24,14 @@ class KookApiPaths: DIRECT_MESSAGE_CREATE = f"{BASE_URL}{API_VERSION_PATH}/direct-message/create" +class KookMentionTagName(str, Enum): + """用来匹配 `(tagName)value(tagName)` 格式里的tagName , 例如: `(met)all(met)` + 定义参见KMarkdown语法文档: https://developer.kookapp.cn/doc/kmarkdown""" + + MENTION = "met" + ROLE = "rol" + + class KookMessageType(IntEnum): """定义参见kook事件结构文档: https://developer.kookapp.cn/doc/event/event-introduction""" @@ -56,6 +65,14 @@ class KookModuleType(str, Enum): CARD = "card" +class KookRoleExtraType(str, Enum): + """定义参见kook事件结构文档: https://developer.kookapp.cn/doc/event/event-introduction""" + + ADDED_ROLE = "added_role" + DELETED_ROLE = "deleted_role" + UPDATED_ROLE = "updated_role" + + ThemeType = Literal[ "primary", "success", @@ -74,7 +91,9 @@ SectionMode = Literal["left", "right"] CountdownMode = Literal["day", "hour", "second"] -class KookBaseDataClass(BaseModel): +class KookBaseReceiveDataClass(BaseModel): + """接收数据基类,`to_dict`/`to_json`默认保证尽量json原样输出""" + model_config = ConfigDict( extra="allow", arbitrary_types_allowed=True, @@ -91,11 +110,51 @@ class KookBaseDataClass(BaseModel): def to_dict( self, - mode: Literal["json", "python"] = "python", + mode: Literal["json", "python"] | str = "json", + by_alias=True, + exclude_none=False, + exclude_unset=True, + ) -> dict: + """默认配置预期场景为尽量原样输出,若需要使用此数据类发送json数据, + 请`exclude_none=True, exclude_unset=False`""" + return self.model_dump( + by_alias=by_alias, + exclude_none=exclude_none, + mode=mode, + exclude_unset=exclude_unset, + ) + + def to_json( + self, + indent: int | None = None, + ensure_ascii=False, + by_alias=True, + exclude_none=False, + exclude_unset=True, + ) -> str: + """默认配置预期场景为尽量原样输出,若需要使用此数据类发送json数据, + 请`exclude_none=True, exclude_unset=False`""" + return self.model_dump_json( + indent=indent, + ensure_ascii=ensure_ascii, + by_alias=by_alias, + exclude_none=exclude_none, + exclude_unset=exclude_unset, + ) + + +class KookBaseSendDataClass(KookBaseReceiveDataClass): + """发送数据基类,`to_dict`/`to_json`保证默认输出内容格式包含接口格式所需最简格式内容""" + + def to_dict( + self, + mode: Literal["json", "python"] | str = "json", by_alias=True, exclude_none=True, exclude_unset=False, ) -> dict: + """默认配置预期场景为发送数据,若需要使用此数据类接收数据并尽量原样json输出, + 请`exclude_none=False, exclude_unset=True`""" return self.model_dump( by_alias=by_alias, exclude_none=exclude_none, @@ -111,6 +170,8 @@ class KookBaseDataClass(BaseModel): exclude_none=True, exclude_unset=False, ) -> str: + """默认配置预期场景为发送数据,若需要使用此数据类接收数据并尽量原样json输出, + 请`exclude_none=False, exclude_unset=True`""" return self.model_dump_json( indent=indent, ensure_ascii=ensure_ascii, @@ -120,7 +181,7 @@ class KookBaseDataClass(BaseModel): ) -class KookCardModelBase(KookBaseDataClass): +class KookCardModelBase(KookBaseSendDataClass): """卡片模块基类""" type: str @@ -256,10 +317,28 @@ AnyModule = Annotated[ ] -class KookCardMessage(KookBaseDataClass): +class KookCardMessage(KookBaseSendDataClass): """卡片定义文档详见 : https://developer.kookapp.cn/doc/cardmessage - 此类型不能直接to_json后发送,因为kook要求卡片容器json顶层必须是**列表** - 若要发送卡片消息,请使用KookCardMessageContainer + 适用于发送单个卡片消息 + 将此消息类型放入`Json`的data字段进行卡片消息发送,适配器会自动添加顶层的列表 + 若要发送多个卡片消息，推荐使用KookCardMessageContainer进行卡片消息组装 + + 使用方法： + ```python + chain = [] + chain.append( + Json( + data=KookCardMessage( + theme="info", + size="lg", + modules=[ + HeaderModule(text=PlainTextElement(content="test1")), + ], + ).to_dict() + ) + ) + yield event.chain_result(chain) + ``` """ model_config = ConfigDict(arbitrary_types_allowed=True) @@ -276,14 +355,69 @@ class KookCardMessage(KookBaseDataClass): class KookCardMessageContainer(list[KookCardMessage]): - """卡片消息容器(列表),此类型可以直接to_json后发送出去""" + """卡片消息容器(列表),可放入多个卡片消息(KookCardMessage) + + 使用方法: + ```python + chain = [] + chain.append( + Json( + data=KookCardMessageContainer( + [ + KookCardMessage( + theme="info", + size="lg", + modules=[ + HeaderModule(text=PlainTextElement(content="test1")), + ], + ) + ] + ).to_dict() + ) + ) + yield event.chain_result(chain) + ``` + """ def append(self, object: KookCardMessage) -> None: return super().append(object) - def to_json(self, indent: int | None = None, ensure_ascii: bool = True) -> str: + def to_dict( + self, + by_alias=True, + exclude_none=True, + exclude_unset=False, + ) -> list[dict]: + """默认配置预期场景为发送数据,若需要使用此数据类接收数据并尽量原样json输出, + 请`exclude_none=False, exclude_unset=True`""" + return [ + i.to_dict( + by_alias=by_alias, + exclude_none=exclude_none, + exclude_unset=exclude_unset, + ) + for i in self + ] + + def to_json( + self, + indent: int | None = None, + ensure_ascii: bool = True, + by_alias=True, + exclude_none=True, + exclude_unset=False, + ) -> str: + """默认配置预期场景为发送数据,若需要使用此数据类接收数据并尽量原样json输出, + 请`exclude_none=False, exclude_unset=True`""" return json.dumps( - [i.to_dict() for i in self], + [ + i.to_dict( + by_alias=by_alias, + exclude_none=exclude_none, + exclude_unset=exclude_unset, + ) + for i in self + ], indent=indent, ensure_ascii=ensure_ascii, ) @@ -302,8 +436,7 @@ class OrderMessage(BaseModel): class KookMessageSignal(IntEnum): """KOOK WebSocket 信令类型 - ws文档: https://developer.kookapp.cn/doc/websocket - """ + ws文档: https://developer.kookapp.cn/doc/websocket""" MESSAGE = 0 """server->client 消息(s包含聊天和通知消息)""" @@ -327,7 +460,7 @@ class KookChannelType(str, Enum): BROADCAST = "BROADCAST" -class KookAuthor(KookBaseDataClass): +class KookAuthor(KookBaseReceiveDataClass): id: str username: str identify_num: str @@ -340,25 +473,110 @@ class KookAuthor(KookBaseDataClass): roles: list[int] = Field(default_factory=list) -class KookKMarkdown(KookBaseDataClass): +class KookMarkdownMentionPart(KookBaseReceiveDataClass): + """ + 文档参考: https://developer.kookapp.cn/doc/event/message + """ + + id: str + username: str + full_name: str + avatar: str + + +class KookMarkdownMentionRolePart(KookBaseReceiveDataClass): + """ + 文档参考: https://developer.kookapp.cn/doc/event/message + """ + + role_id: int + name: str + color: int + color_type: int + color_map: list[Any] + position: int | None = None + hoist: int | None = None + mentionable: int | None = None + permissions: int | None = None + + +class KookKMarkdown(KookBaseReceiveDataClass): raw_content: str - mention_part: list[Any] = Field(default_factory=list) - mention_role_part: list[Any] = Field(default_factory=list) + mention_part: list[KookMarkdownMentionPart] = Field(default_factory=list) + mention_role_part: list[KookMarkdownMentionRolePart] = Field(default_factory=list) -class KookExtra(KookBaseDataClass): - type: int | str +class KookRole(KookBaseReceiveDataClass): + """服务器角色对象数据结构""" + + role_id: int = Field(alias="role_id") + name: str | None = None + color: int | None = None + position: int | None = None + hoist: int | None = 0 # 是否在成员列表中单独展示 + mentionable: int | None = 0 # 是否允许所有人提到该角色 + permissions: int | None = None + + +class KookRoleEventBody(KookBaseReceiveDataClass): + """ + 服务器角色相关事件 (added_role, updated_role, deleted_role) 的 Body 部分 + 文档参考: https://developer.kookapp.cn/doc/event/guild-role + """ + + role_id: int | None = None # 在 deleted_role 中通常只给 ID + name: str | None = None + color: int | None = None + position: int | None = None + hoist: int | None = None + mentionable: int | None = None + permissions: int | None = None + # 有些事件会将完整的 role 对象包裹在 body 里 + # 如果是 added_role 且需要处理更完整的结构，可以扩展 + + +class KookExtra(KookBaseReceiveDataClass): + """事件结构定义 + 文档参考 : https://developer.kookapp.cn/doc/event/event-introduction""" + + type: KookRoleExtraType | str | int + """当 type 非系统消息(255)时, type为int + + 当 type 为系统消息(255)时, type为str + """ + code: str | None = None - body: dict[str, Any] | None = None + body: KookRole | dict[str, Any] | None = None author: KookAuthor | None = None kmarkdown: KookKMarkdown | None = None last_msg_content: str | None = None mention: list[str] = Field(default_factory=list) mention_all: bool = False mention_here: bool = False + guild_id: str | None = None + guild_type: int | None = None + channel_name: str | None = None + visible_only: str | None = None + mention_no_at: list | None = None + mention_roles: list[int] | None = None + nav_channels: list | None = None + emoji: list | None = None + preview_content: str | None = None + channel_type: int | None = None + send_msg_device: int | None = None + + @field_validator("type", mode="before") + @classmethod + def parse_type(cls, value): + """优先尝试匹配枚举，失败则保留原值""" + if isinstance(value, str): + if value in {e.value for e in KookRoleExtraType}: + return KookRoleExtraType(value) + + return value -class KookMessageEventData(KookBaseDataClass): +class KookMessageEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.MESSAGE] = Field( KookMessageSignal.MESSAGE, exclude=True, @@ -369,7 +587,7 @@ class KookMessageEventData(KookBaseDataClass): type: KookMessageType target_id: str author_id: str - content: str | dict[str, Any] + content: str | dict[str, Any] # 道具消息时这里是dict msg_id: str msg_timestamp: int nonce: str @@ -377,7 +595,7 @@ class KookMessageEventData(KookBaseDataClass): extra: KookExtra -class KookHelloEventData(KookBaseDataClass): +class KookHelloEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.HELLO] = Field( KookMessageSignal.HELLO, exclude=True, @@ -388,7 +606,7 @@ class KookHelloEventData(KookBaseDataClass): session_id: str -class KookPingEventData(KookBaseDataClass): +class KookPingEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.PING] = Field( KookMessageSignal.PING, exclude=True, @@ -396,7 +614,7 @@ class KookPingEventData(KookBaseDataClass): """only for type hint""" -class KookPongEventData(KookBaseDataClass): +class KookPongEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.PONG] = Field( KookMessageSignal.PONG, exclude=True, @@ -404,7 +622,7 @@ class KookPongEventData(KookBaseDataClass): """only for type hint""" -class KookResumeEventData(KookBaseDataClass): +class KookResumeEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.RESUME] = Field( KookMessageSignal.RESUME, exclude=True, @@ -412,7 +630,7 @@ class KookResumeEventData(KookBaseDataClass): """only for type hint""" -class KookReconnectEventData(KookBaseDataClass): +class KookReconnectEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.RECONNECT] = Field( KookMessageSignal.RECONNECT, exclude=True, @@ -423,7 +641,7 @@ class KookReconnectEventData(KookBaseDataClass): err: str -class KookResumeAckEventData(KookBaseDataClass): +class KookResumeAckEventData(KookBaseReceiveDataClass): signal: Literal[KookMessageSignal.RESUME_ACK] = Field( KookMessageSignal.RESUME_ACK, exclude=True, @@ -433,7 +651,7 @@ class KookResumeAckEventData(KookBaseDataClass): session_id: str -class KookWebsocketEvent(KookBaseDataClass): +class KookWebsocketEvent(KookBaseReceiveDataClass): """KOOK WebSocket 原始推送结构""" signal: KookMessageSignal = Field( @@ -470,22 +688,22 @@ class KookWebsocketEvent(KookBaseDataClass): return data -class KookUserTag(KookBaseDataClass): +class KookUserTag(KookBaseReceiveDataClass): color: str bg_color: str text: str -class KookApiResponseBase(KookBaseDataClass): +class KookApiResponseBase(KookBaseReceiveDataClass): code: int message: str - data: Any + data: dict # 就算请求失败了也是空dict def success(self) -> bool: return self.code == 0 -class KookUserMeData(KookBaseDataClass): +class KookUserMeData(KookBaseReceiveDataClass): """USER_ME 接口返回的 'data' 字段主体""" id: str @@ -514,7 +732,47 @@ class KookUserMeResponse(KookApiResponseBase): data: KookUserMeData -class KookGatewayIndexData(KookBaseDataClass): +class KookUserMeViewData(KookBaseReceiveDataClass): + """USER_ME 接口返回的 'data' 字段主体""" + + class KookTagInfo(KookBaseReceiveDataClass): + color: str + bg_color: str + text: str + + id: str + username: str + identify_num: str + online: bool + os: str + status: int + avatar: str + vip_avatar: str + banner: str + nickname: str + roles: list[int] + is_vip: bool + vip_amp: bool + bot: bool + kpm_vip: str | None = None + wealth_level: int + bot_status: int + tag_info: KookTagInfo + mobile_verified: bool + is_sys: bool + client_id: str + verified: bool + joined_at: int + active_time: int + + +class KookUserViewResponse(KookApiResponseBase): + """USER_VIEW 完整响应结构""" + + data: KookUserMeViewData + + +class KookGatewayIndexData(KookBaseReceiveDataClass): url: str diff --git a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py index 781100efe..3ab2c32ee 100644 --- a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py +++ b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py @@ -44,6 +44,22 @@ def _patch_qq_botpy_formdata() -> None: _patch_qq_botpy_formdata() +# Retry decorator for QQ Official API transient errors (HTTP 500/504) +_qqofficial_retry = retry( + retry=retry_if_exception_type( + ( + botpy.errors.ServerError, + botpy.errors.SequenceNumberError, + OSError, + asyncio.TimeoutError, + ) + ), + stop=stop_after_attempt(5), + wait=wait_exponential(multiplier=2, min=2, max=30), + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, +) + class QQOfficialMessageEvent(AstrMessageEvent): MARKDOWN_NOT_ALLOWED_ERROR = "不允许发送原生 markdown" @@ -465,6 +481,8 @@ class QQOfficialMessageEvent(AstrMessageEvent): file_info=result["file_info"], ttl=result.get("ttl", 0), ) + except (botpy.errors.ServerError, botpy.errors.SequenceNumberError): + logger.error(f"上传媒体文件失败，共尝试5次后放弃: {file_source}") except Exception as e: logger.error(f"上传请求错误: {e}") return None diff --git a/astrbot/core/platform/sources/webchat/webchat_adapter.py b/astrbot/core/platform/sources/webchat/webchat_adapter.py index 0267f1b3f..ba229c515 100644 --- a/astrbot/core/platform/sources/webchat/webchat_adapter.py +++ b/astrbot/core/platform/sources/webchat/webchat_adapter.py @@ -230,6 +230,11 @@ class WebChatAdapter(Platform): payload.get("enable_streaming", True), ) message_event.set_extra("action_type", payload.get("action_type")) + message_event.set_extra("llm_checkpoint_id", payload.get("llm_checkpoint_id")) + message_event.set_extra( + "thread_selected_text", payload.get("thread_selected_text") + ) + self.commit_event(message_event) async def terminate(self) -> None: diff --git a/astrbot/core/platform_message_history_mgr.py b/astrbot/core/platform_message_history_mgr.py index b3ac69ac3..951f4b288 100644 --- a/astrbot/core/platform_message_history_mgr.py +++ b/astrbot/core/platform_message_history_mgr.py @@ -13,6 +13,7 @@ class PlatformMessageHistoryManager: content: dict, # TODO: parse from message chain sender_id: str | None = None, sender_name: str | None = None, + llm_checkpoint_id: str | None = None, ) -> PlatformMessageHistory: """Insert a new platform message history record.""" return await self.db.insert_platform_message_history( @@ -21,6 +22,7 @@ class PlatformMessageHistoryManager: content=content, sender_id=sender_id, sender_name=sender_name, + llm_checkpoint_id=llm_checkpoint_id, ) async def get( @@ -52,3 +54,20 @@ class PlatformMessageHistoryManager: user_id=user_id, offset_sec=offset_sec, ) + + async def update( + self, + message_id: int, + content: dict | None = None, + llm_checkpoint_id: str | None = None, + ) -> None: + """Update a platform message history record.""" + await self.db.update_platform_message_history( + message_id=message_id, + content=content, + llm_checkpoint_id=llm_checkpoint_id, + ) + + async def delete_by_id(self, message_id: int) -> None: + """Delete a platform message history record by ID.""" + await self.db.delete_platform_message_history_by_id(message_id) diff --git a/astrbot/core/provider/entities.py b/astrbot/core/provider/entities.py index 2e095cc3d..b610adf11 100644 --- a/astrbot/core/provider/entities.py +++ b/astrbot/core/provider/entities.py @@ -21,6 +21,7 @@ from astrbot.core.agent.message import ( ContentPart, ToolCall, ToolCallMessageSegment, + is_checkpoint_message, ) from astrbot.core.agent.tool import ToolSet from astrbot.core.db.po import Conversation @@ -151,6 +152,8 @@ class ProviderRequest: result_parts = [] for ctx in self.contexts: + if is_checkpoint_message(ctx): + continue role = ctx.get("role", "unknown") content = ctx.get("content", "") diff --git a/astrbot/core/provider/modalities.py b/astrbot/core/provider/modalities.py new file mode 100644 index 000000000..66ac74e9b --- /dev/null +++ b/astrbot/core/provider/modalities.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import copy +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +from astrbot import logger +from astrbot.core.agent.message import Message + + +@dataclass(slots=True) +class ContextSanitizeStats: + fixed_image_blocks: int = 0 + fixed_audio_blocks: int = 0 + fixed_tool_messages: int = 0 + removed_tool_calls: int = 0 + + @property + def changed(self) -> bool: + return bool( + self.fixed_image_blocks + or self.fixed_audio_blocks + or self.fixed_tool_messages + or self.removed_tool_calls + ) + + +def _message_to_dict(message: dict[str, Any] | Message) -> dict[str, Any] | None: + if isinstance(message, Message): + return dict(message.model_dump()) + if isinstance(message, dict): + return dict(copy.deepcopy(message)) + return None + + +def sanitize_contexts_by_modalities( + contexts: Sequence[dict[str, Any] | Message], + modalities: list[str] | None, +) -> tuple[list[dict[str, Any]], ContextSanitizeStats]: + if not contexts: + return [], ContextSanitizeStats() + if not modalities or not isinstance(modalities, list): + copied_contexts = [] + for msg in contexts: + copied_msg = _message_to_dict(msg) + if copied_msg: + copied_contexts.append(copied_msg) + return copied_contexts, ContextSanitizeStats() + + supports_image = "image" in modalities + supports_audio = "audio" in modalities + supports_tool_use = "tool_use" in modalities + if supports_image and supports_audio and supports_tool_use: + copied_contexts = [] + for msg in contexts: + copied_msg = _message_to_dict(msg) + if copied_msg: + copied_contexts.append(copied_msg) + return copied_contexts, ContextSanitizeStats() + + sanitized_contexts: list[dict[str, Any]] = [] + stats = ContextSanitizeStats() + + for raw_msg in contexts: + msg = _message_to_dict(raw_msg) + if not msg: + continue + role = msg.get("role") + if not role: + continue + + if not supports_tool_use: + if role == "tool": + stats.fixed_tool_messages += 1 + fixed_msg: dict[str, Any] = { + "role": "user", + "content": _tool_result_placeholder(msg.get("content")), + } + msg = fixed_msg + if role == "assistant" and "tool_calls" in msg: + stats.removed_tool_calls += 1 + msg.pop("tool_calls", None) + msg.pop("tool_call_id", None) + + if not supports_image or not supports_audio: + content = msg.get("content") + if isinstance(content, list): + filtered_parts: list[Any] = [] + removed_any_multimodal = False + for part in content: + if isinstance(part, dict): + part_type = str(part.get("type", "")).lower() + if not supports_image and part_type in {"image_url", "image"}: + removed_any_multimodal = True + stats.fixed_image_blocks += 1 + filtered_parts.append({"type": "text", "text": "[Image]"}) + continue + if not supports_audio and part_type in { + "audio_url", + "input_audio", + }: + removed_any_multimodal = True + stats.fixed_audio_blocks += 1 + filtered_parts.append({"type": "text", "text": "[Audio]"}) + continue + filtered_parts.append(part) + if removed_any_multimodal: + msg["content"] = filtered_parts + + if role == "assistant": + content = msg.get("content") + has_tool_calls = bool(msg.get("tool_calls")) + if not has_tool_calls: + if not content: + continue + if isinstance(content, str) and not content.strip(): + continue + + sanitized_contexts.append(msg) + + return sanitized_contexts, stats + + +def _tool_result_placeholder(content: Any) -> str: + if isinstance(content, str): + content_text = content.strip() + elif isinstance(content, list): + text_parts: list[str] = [] + for part in content: + if isinstance(part, dict): + part_type = str(part.get("type", "")).lower() + if part_type == "text": + text_parts.append(str(part.get("text", ""))) + elif part_type in {"image_url", "image"}: + text_parts.append("[Image]") + elif part_type in {"audio_url", "input_audio"}: + text_parts.append("[Audio]") + content_text = "\n".join(part for part in text_parts if part).strip() + else: + content_text = "" + if not content_text: + return "[Tool result]" + return f"[Tool result]\n{content_text}" + + +def log_context_sanitize_stats(stats: ContextSanitizeStats) -> None: + if not stats.changed: + return + logger.debug( + "context modality fix applied: " + "fixed_image_blocks=%s, fixed_audio_blocks=%s, " + "fixed_tool_messages=%s, removed_tool_calls=%s", + stats.fixed_image_blocks, + stats.fixed_audio_blocks, + stats.fixed_tool_messages, + stats.removed_tool_calls, + ) diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 230660155..261709d67 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -4,10 +4,7 @@ import os from collections.abc import AsyncGenerator from typing import Literal, TypeAlias, Union -import aiofiles -import anyio - -from astrbot.core.agent.message import ContentPart, Message +from astrbot.core.agent.message import ContentPart, Message, is_checkpoint_message from astrbot.core.agent.tool import ToolSet from astrbot.core.provider.entities import ( LLMResponse, @@ -191,6 +188,8 @@ class Provider(AbstractProvider): return [] dicts: list[dict] = [] for message in messages: + if is_checkpoint_message(message): + continue if isinstance(message, Message): dicts.append(message.model_dump()) else: diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index 64e6b3d4c..7a5f8bd59 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -20,6 +20,7 @@ from astrbot.core.provider.func_tool_manager import ToolSet from astrbot.core.provider.register import register_provider_adapter from astrbot.core.utils.io import download_image_by_url from astrbot.core.utils.network_utils import ( + create_proxy_client, is_connection_error, log_connection_failure, ) @@ -117,15 +118,13 @@ class ProviderAnthropic(Provider): http_client=self._create_http_client(provider_config), ) - def _create_http_client(self, provider_config: dict) -> httpx.AsyncClient | None: - """创建带代理的 HTTP 客户端""" - proxy = provider_config.get("proxy", "") - if proxy: - logger.info(f"[Anthropic] 使用代理: {proxy}") - return httpx.AsyncClient(proxy=proxy, headers=self.custom_headers) - if self.custom_headers: - return httpx.AsyncClient(headers=self.custom_headers) - return None + def _create_http_client(self, provider_config: dict) -> httpx.AsyncClient: + """创建带代理的 HTTP 客户端，使用系统 SSL 证书""" + return create_proxy_client( + "Anthropic", + provider_config.get("proxy", ""), + headers=self.custom_headers, + ) def _apply_thinking_config(self, payloads: dict) -> None: thinking_type = self.thinking_config.get("type", "") diff --git a/astrbot/core/provider/sources/minimax_token_plan_source.py b/astrbot/core/provider/sources/minimax_token_plan_source.py index 5a578424f..d226707fd 100644 --- a/astrbot/core/provider/sources/minimax_token_plan_source.py +++ b/astrbot/core/provider/sources/minimax_token_plan_source.py @@ -1,11 +1,15 @@ +from astrbot import logger from astrbot.core.provider.sources.anthropic_source import ProviderAnthropic from ..register import register_provider_adapter MINIMAX_TOKEN_PLAN_MODELS = [ "MiniMax-M2.7", + "MiniMax-M2.7-highspeed", "MiniMax-M2.5", + "MiniMax-M2.5-highspeed", "MiniMax-M2.1", + "MiniMax-M2.1-highspeed", "MiniMax-M2", ] @@ -43,9 +47,13 @@ class ProviderMiniMaxTokenPlan(ProviderAnthropic): configured_model = provider_config.get("model", "MiniMax-M2.7") if configured_model not in MINIMAX_TOKEN_PLAN_MODELS: - raise ValueError( - f"Unsupported model: {configured_model!r}. " - f"Supported models: {', '.join(MINIMAX_TOKEN_PLAN_MODELS)}" + logger.warning( + f"Configured model {configured_model!r} is not in the known " + f"Token Plan model list " + f"({', '.join(MINIMAX_TOKEN_PLAN_MODELS)}). " + f"The model may still work if your plan supports it. " + f"If you encounter errors, please check your plan's " + f"model availability." ) self.set_model(configured_model) diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index 9aa65e15d..fdd6e378f 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -439,7 +439,7 @@ class ProviderOpenAIOfficial(Provider): image_fallback_used, ) - def _create_http_client(self, provider_config: dict) -> httpx.AsyncClient | None: + def _create_http_client(self, provider_config: dict) -> httpx.AsyncClient: """创建带代理的 HTTP 客户端""" proxy = provider_config.get("proxy", "") return create_proxy_client("OpenAI", proxy) @@ -860,29 +860,26 @@ class ProviderOpenAIOfficial(Provider): # 工具集未提供 # Should be unreachable raise Exception("工具集未提供") - for tool in tools.list_tools(): - # Safely access function attribute (ChatCompletionMessageCustomToolCall doesn't have it) - tool_call_function = getattr(tool_call, "function", None) - if not tool_call_function: - continue - if ( - tool_call.type == "function" - and tool_call_function.name == tool.name - ): - # workaround for #1454 - arguments = tool_call_function.arguments - if isinstance(arguments, str): - args = json.loads(arguments) - else: - args = arguments - args_ls.append(args) - func_name_ls.append(tool_call_function.name) - tool_call_ids.append(tool_call.id) + + if tool_call.type == "function": + # workaround for #1454 + if isinstance(tool_call.function.arguments, str): + try: + args = json.loads(tool_call.function.arguments) + except json.JSONDecodeError as e: + logger.error(f"解析参数失败: {e}") + args = {} + else: + args = tool_call.function.arguments + args_ls.append(args) + func_name_ls.append(tool_call.function.name) + tool_call_ids.append(tool_call.id) # gemini-2.5 / gemini-3 series extra_content handling extra_content = getattr(tool_call, "extra_content", None) if extra_content is not None: tool_call_extra_content_dict[tool_call.id] = extra_content + llm_response.role = "tool" llm_response.tools_call_args = args_ls llm_response.tools_call_name = func_name_ls diff --git a/astrbot/core/star/config.py b/astrbot/core/star/config.py index e8f350b03..4893b27d2 100644 --- a/astrbot/core/star/config.py +++ b/astrbot/core/star/config.py @@ -1,4 +1,4 @@ -"""此功能已过时,参考 https://astrbot.app/dev/plugin.html#%E6%B3%A8%E5%86%8C%E6%8F%92%E4%BB%B6%E9%85%8D%E7%BD%AE-beta""" +"""此功能已过时，参考 https://docs.astrbot.app/dev/plugin.html#%E6%B3%A8%E5%86%8C%E6%8F%92%E4%BB%B6%E9%85%8D%E7%BD%AE-beta""" import json import os diff --git a/astrbot/core/star/filter/regex.py b/astrbot/core/star/filter/regex.py index 4c2a033ef..8e85c4d30 100644 --- a/astrbot/core/star/filter/regex.py +++ b/astrbot/core/star/filter/regex.py @@ -10,9 +10,9 @@ from . import HandlerFilter class RegexFilter(HandlerFilter): """正则表达式过滤器""" - def __init__(self, regex: str) -> None: - self.regex_str = regex + def __init__(self, regex: str | re.Pattern) -> None: self.regex = re.compile(regex) + self.regex_str = self.regex.pattern def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool: return bool(self.regex.search(event.get_message_str().strip())) diff --git a/astrbot/core/star/register/star_handler.py b/astrbot/core/star/register/star_handler.py index 76116d0bd..dfab95025 100644 --- a/astrbot/core/star/register/star_handler.py +++ b/astrbot/core/star/register/star_handler.py @@ -292,7 +292,7 @@ def register_platform_adapter_type( return decorator -def register_regex(regex: str, **kwargs): +def register_regex(regex: str | re.Pattern, **kwargs): """注册一个 Regex""" def decorator(awaitable): diff --git a/astrbot/core/tools/computer_tools/fs.py b/astrbot/core/tools/computer_tools/fs.py index 3a7db6e24..5e02e69b7 100644 --- a/astrbot/core/tools/computer_tools/fs.py +++ b/astrbot/core/tools/computer_tools/fs.py @@ -171,7 +171,7 @@ def _decode_escaped_text(value: str) -> str: @dataclass class FileReadTool(FunctionTool): name: str = "astrbot_file_read_tool" - description: str = "read file content." + description: str = "read file content. Supports text, image, and PDF (text extraction), docx and epub files." parameters: dict = field( default_factory=lambda: { "type": "object", diff --git a/astrbot/core/tools/message_tools.py b/astrbot/core/tools/message_tools.py index 3789677e0..b4c82a394 100644 --- a/astrbot/core/tools/message_tools.py +++ b/astrbot/core/tools/message_tools.py @@ -79,7 +79,21 @@ class SendMessageToUserTool(FunctionTool[AstrAgentContext]): context: ContextWrapper[AstrAgentContext], path: str, ) -> tuple[str, bool]: - if os.path.exists(path): + # if the path is relative, check if the file exists in user's local workspace + if not os.path.isabs(path): + unified_msg_origin = context.context.event.unified_msg_origin + if unified_msg_origin: + from astrbot.core.tools.computer_tools.util import workspace_root + + try: + ws_path = workspace_root(unified_msg_origin) + ws_candidate = (ws_path / path).resolve() + if ws_candidate.is_file() and ws_candidate.is_relative_to(ws_path): + return str(ws_candidate), False + except Exception: + pass + # check if the file exists in local environment (only allow absolute paths to prevent traversal) + elif os.path.isfile(path): return path, False try: diff --git a/astrbot/core/tools/web_search_tools.py b/astrbot/core/tools/web_search_tools.py index bc9d4d7f3..c5827c25a 100644 --- a/astrbot/core/tools/web_search_tools.py +++ b/astrbot/core/tools/web_search_tools.py @@ -201,6 +201,10 @@ async def _bocha_search( header = { "Authorization": f"Bearer {bocha_key}", "Content-Type": "application/json", + # Explicitly disable brotli encoding to avoid aiohttp >= 3.13.3 brotli + # decompression incompatibility (TypeError: process() takes exactly 1 argument). + # See: https://github.com/aio-libs/aiohttp/issues/11898 + "Accept-Encoding": "gzip, deflate", } async with ( aiohttp.ClientSession(trust_env=True) as session, diff --git a/astrbot/core/utils/network_utils.py b/astrbot/core/utils/network_utils.py index 03a870418..ca4fee365 100644 --- a/astrbot/core/utils/network_utils.py +++ b/astrbot/core/utils/network_utils.py @@ -1,9 +1,13 @@ """Network error handling utilities for providers.""" +import ssl + import httpx from astrbot import logger +_SYSTEM_SSL_CTX = ssl.create_default_context() + def is_connection_error(exc: BaseException) -> bool: """Check if an exception is a connection/network related error. @@ -86,21 +90,30 @@ def log_connection_failure( def create_proxy_client( provider_label: str, proxy: str | None = None, -) -> httpx.AsyncClient | None: + headers: dict[str, str] | None = None, + verify: ssl.SSLContext | str | bool | None = None, +) -> httpx.AsyncClient: """Create an httpx AsyncClient with proxy configuration if provided. + Uses the system SSL certificate store instead of certifi, which avoids + SSL verification failures for endpoints whose CA chain is not in certifi + but is trusted by the operating system. + Note: The caller is responsible for closing the client when done. Consider using the client as a context manager or calling aclose() explicitly. Args: provider_label: The provider name for log prefix (e.g., "OpenAI", "Gemini") proxy: The proxy address (e.g., "http://127.0.0.1:7890"), or None/empty + headers: Optional custom headers to include in every request + verify: Optional override for TLS verification. Defaults to the shared + system SSL context when not provided. Returns: - An httpx.AsyncClient configured with the proxy, or None if no proxy - + An httpx.AsyncClient created with the shared system SSL context; the proxy is applied only if one is provided. """ + resolved_verify = _SYSTEM_SSL_CTX if verify is None else verify if proxy: logger.info(f"[{provider_label}] 使用代理: {proxy}") - return httpx.AsyncClient(proxy=proxy) - return None + return httpx.AsyncClient(proxy=proxy, verify=resolved_verify, headers=headers) + return httpx.AsyncClient(verify=resolved_verify, headers=headers) diff --git a/astrbot/core/utils/t2i/renderer.py b/astrbot/core/utils/t2i/renderer.py index 45093cfe1..39d56e5f2 100644 --- a/astrbot/core/utils/t2i/renderer.py +++ b/astrbot/core/utils/t2i/renderer.py @@ -28,7 +28,7 @@ class HtmlRenderer: @return: 图片 URL 或者文件路径,取决于 return_url 参数｡ - @example: 参见 https://astrbot.app 插件开发部分｡ + @example: 参见 https://docs.astrbot.app 插件开发部分。 """ return await self.network_strategy.render_custom_template( tmpl_str, diff --git a/astrbot/dashboard/routes/chat.py b/astrbot/dashboard/routes/chat.py index c2e615498..8a1c859fc 100644 --- a/astrbot/dashboard/routes/chat.py +++ b/astrbot/dashboard/routes/chat.py @@ -4,13 +4,14 @@ import os import re import uuid from contextlib import asynccontextmanager -from pathlib import Path -from typing import Any +from copy import deepcopy +from typing import cast import anyio from quart import g, make_response, request, send_file from astrbot.core import logger, sp +from astrbot.core.agent.message import get_checkpoint_id, is_checkpoint_message from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase from astrbot.core.platform.message_type import MessageType @@ -80,6 +81,12 @@ class ChatRoute(Route): "POST", self.update_session_display_name, ), + "/chat/message/edit": ("POST", self.update_message), + "/chat/message/regenerate": ("POST", self.regenerate_message), + "/chat/thread/create": ("POST", self.create_thread), + "/chat/thread/get": ("GET", self.get_thread), + "/chat/thread/send": ("POST", self.send_thread_message), + "/chat/thread/delete": ("POST", self.delete_thread), "/chat/get_file": ("GET", self.get_file), "/chat/get_attachment": ("GET", self.get_attachment), "/chat/post_file": ("POST", self.post_file), @@ -260,6 +267,238 @@ class ChatRoute(Route): used_refs.append(payload) return {"used": used_refs} if used_refs else {} + def _sanitize_message_content(self, content: dict) -> dict: + """Normalize editable WebChat message content before persisting.""" + if not isinstance(content, dict): + raise ValueError("Missing key: content") + + normalized = deepcopy(content) + message_type = normalized.get("type") + if message_type not in {"user", "bot"}: + raise ValueError("Invalid key: content.type") + + message_parts = normalized.get("message") + if not isinstance(message_parts, list): + raise ValueError("Missing key: content.message") + normalized["message"] = strip_message_parts_path_fields(message_parts) + return normalized + + def _extract_platform_message_text(self, content: dict | None) -> str: + if not isinstance(content, dict): + return "" + message_parts = content.get("message") + if not isinstance(message_parts, list): + return "" + texts: list[str] = [] + for part in message_parts: + if isinstance(part, dict) and part.get("type") == "plain": + text = part.get("text") + if isinstance(text, str): + texts.append(text) + return "".join(texts) + + def _build_webchat_unified_msg_origin(self, session) -> str: + message_type = ( + MessageType.GROUP_MESSAGE.value + if session.is_group + else MessageType.FRIEND_MESSAGE.value + ) + return ( + f"{session.platform_id}:{message_type}:" + f"{session.platform_id}!{session.creator}!{session.session_id}" + ) + + def _build_thread_unified_msg_origin(self, creator: str, thread_id: str) -> str: + return ( + f"webchat:{MessageType.FRIEND_MESSAGE.value}:webchat!{creator}!{thread_id}" + ) + + def _serialize_thread(self, thread) -> dict: + return { + "thread_id": thread.thread_id, + "parent_session_id": thread.parent_session_id, + "parent_message_id": thread.parent_message_id, + "base_checkpoint_id": thread.base_checkpoint_id, + "selected_text": thread.selected_text, + "created_at": to_utc_isoformat(thread.created_at), + "updated_at": to_utc_isoformat(thread.updated_at), + } + + async def _delete_threads_by_ids(self, thread_ids: list[str], creator: str) -> None: + for thread_id in thread_ids: + unified_msg_origin = self._build_thread_unified_msg_origin( + creator, thread_id + ) + active_event_registry.request_agent_stop_all(unified_msg_origin) + await self.conv_mgr.delete_conversations_by_user_id(unified_msg_origin) + await self.platform_history_mgr.delete( + platform_id="webchat_thread", + user_id=thread_id, + offset_sec=99999999, + ) + webchat_queue_mgr.remove_queues(thread_id) + self.running_convs.pop(thread_id, None) + + async def _load_current_conversation_history(self, session) -> tuple[str, list]: + unified_msg_origin = self._build_webchat_unified_msg_origin(session) + conversation_id = await self.conv_mgr.get_curr_conversation_id( + unified_msg_origin + ) + if not conversation_id: + return "", [] + + conversation = await self.conv_mgr.get_conversation( + unified_msg_origin=unified_msg_origin, + conversation_id=conversation_id, + ) + if not conversation: + return "", [] + + try: + history = json.loads(conversation.history or "[]") + except json.JSONDecodeError: + return "", [] + return conversation_id, history if isinstance(history, list) else [] + + def _find_checkpoint_index( + self, history: list[dict], checkpoint_id: str + ) -> int | None: + for index, message in enumerate(history): + if get_checkpoint_id(message) == checkpoint_id: + return index + return None + + def _find_turn_range( + self, history: list[dict], checkpoint_id: str + ) -> tuple[int, int] | None: + checkpoint_index = self._find_checkpoint_index(history, checkpoint_id) + if checkpoint_index is None: + return None + + start = 0 + for index in range(checkpoint_index - 1, -1, -1): + if is_checkpoint_message(history[index]): + start = index + 1 + break + return start, checkpoint_index + + def _is_latest_checkpoint(self, history: list[dict], checkpoint_id: str) -> bool: + for message in reversed(history): + current_checkpoint_id = get_checkpoint_id(message) + if current_checkpoint_id: + return current_checkpoint_id == checkpoint_id + return False + + def _replace_user_conversation_content(self, original_content, edited_text: str): + if isinstance(original_content, str): + return edited_text + if not isinstance(original_content, list): + return edited_text + + result: list[dict] = [] + inserted_text = False + for part in original_content: + if not isinstance(part, dict): + result.append(part) + continue + if part.get("type") != "text": + result.append(part) + continue + text = part.get("text") + if isinstance(text, str) and text.startswith(""): + result.append(part) + continue + if not inserted_text and edited_text: + result.append({"type": "text", "text": edited_text}) + inserted_text = True + + if not inserted_text and edited_text: + result.insert(0, {"type": "text", "text": edited_text}) + return result + + def _replace_assistant_conversation_content( + self, + original_content, + edited_text: str, + reasoning: str, + ): + if isinstance(original_content, str): + return edited_text + if not isinstance(original_content, list): + return [{"type": "text", "text": edited_text}] if edited_text else [] + + result: list[dict] = [] + inserted_text = False + inserted_think = False + for part in original_content: + if not isinstance(part, dict): + result.append(part) + continue + if part.get("type") == "text": + if not inserted_text and edited_text: + result.append({"type": "text", "text": edited_text}) + inserted_text = True + continue + if part.get("type") == "think": + if not inserted_think and reasoning: + result.append({"type": "think", "think": reasoning}) + inserted_think = True + continue + result.append(part) + + if reasoning and not inserted_think: + result.insert(0, {"type": "think", "think": reasoning}) + if edited_text and not inserted_text: + result.append({"type": "text", "text": edited_text}) + return result + + def _find_turn_user_index( + self, history: list[dict], start: int, end: int + ) -> int | None: + for index in range(start, end): + message = history[index] + if isinstance(message, dict) and message.get("role") == "user": + return index + return None + + def _find_turn_final_assistant_index( + self, history: list[dict], start: int, end: int + ) -> int | None: + for index in range(end - 1, start - 1, -1): + message = history[index] + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + if message.get("tool_calls") and not message.get("content"): + continue + return index + return None + + async def _get_sorted_platform_history(self, session) -> list: + history_list = await self.platform_history_mgr.get( + platform_id=session.platform_id, + user_id=session.session_id, + page=1, + page_size=100000, + ) + history_list.sort(key=lambda item: (item.created_at, item.id)) + return history_list + + async def _delete_platform_history_after( + self, session, message_id: int + ) -> list[int]: + history_list = await self._get_sorted_platform_history(session) + should_delete = False + deleted_ids: list[int] = [] + for item in history_list: + if should_delete: + if item.id is not None: + deleted_ids.append(item.id) + await self.platform_history_mgr.delete_by_id(item.id) + continue + if item.id == message_id: + should_delete = True + return deleted_ids + async def _save_bot_message( self, webchat_conv_id: str, @@ -268,6 +507,8 @@ class ChatRoute(Route): reasoning: str, agent_stats: dict, refs: dict, + llm_checkpoint_id: str | None = None, + platform_history_id: str = "webchat", ): """保存 bot 消息到历史记录,返回保存的记录""" bot_message_parts = [] @@ -282,11 +523,12 @@ class ChatRoute(Route): if refs: new_his["refs"] = refs record = await self.platform_history_mgr.insert( - platform_id="webchat", + platform_id=platform_history_id, user_id=webchat_conv_id, content=new_his, sender_id="bot", sender_name="bot", + llm_checkpoint_id=llm_checkpoint_id, ) return record @@ -307,6 +549,9 @@ class ChatRoute(Route): selected_provider = post_data.get("selected_provider") selected_model = post_data.get("selected_model") enable_streaming = post_data.get("enable_streaming", True) + platform_history_id = post_data.get("_platform_history_id") or "webchat" + thread_selected_text = post_data.get("_thread_selected_text") + if not session_id: return Response().error("session_id is empty").to_json() webchat_conv_id = session_id @@ -318,10 +563,13 @@ class ChatRoute(Route): .to_json() ) message_id = str(uuid.uuid4()) + llm_checkpoint_id = post_data.get("_llm_checkpoint_id") or str(uuid.uuid4()) + skip_user_history = bool(post_data.get("_skip_user_history")) back_queue = webchat_queue_mgr.get_or_create_back_queue( message_id, webchat_conv_id, ) + saved_user_record = None async def stream(): client_disconnected = False @@ -338,6 +586,19 @@ class ChatRoute(Route): "session_id": webchat_conv_id, } yield f"data: {json.dumps(session_info, ensure_ascii=False)}\n\n" + if saved_user_record and not client_disconnected: + user_saved_info = { + "type": "user_message_saved", + "data": { + "id": saved_user_record.id, + "created_at": to_utc_isoformat( + saved_user_record.created_at + ), + "llm_checkpoint_id": llm_checkpoint_id, + }, + } + yield f"data: {json.dumps(user_saved_info, ensure_ascii=False)}\n\n" + async with track_conversation(self.running_convs, webchat_conv_id): while True: result, should_break = await _poll_webchat_stream_result( @@ -462,6 +723,8 @@ class ChatRoute(Route): accumulated_reasoning, agent_stats, refs, + llm_checkpoint_id, + platform_history_id, ) if saved_record and (not client_disconnected): saved_info = { @@ -471,6 +734,7 @@ class ChatRoute(Route): "created_at": to_utc_isoformat( saved_record.created_at, ), + "llm_checkpoint_id": llm_checkpoint_id, }, } try: @@ -498,25 +762,34 @@ class ChatRoute(Route): "selected_model": selected_model, "enable_streaming": enable_streaming, "message_id": message_id, + "llm_checkpoint_id": llm_checkpoint_id, + "thread_selected_text": thread_selected_text, }, ), ) message_parts_for_storage = strip_message_parts_path_fields(message_parts) - await self.platform_history_mgr.insert( - platform_id="webchat", - user_id=webchat_conv_id, - content={"type": "user", "message": message_parts_for_storage}, - sender_id=username, - sender_name=username, - ) - response = await make_response( - stream(), - { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Transfer-Encoding": "chunked", - "Connection": "keep-alive", - }, + + if not skip_user_history: + saved_user_record = await self.platform_history_mgr.insert( + platform_id=platform_history_id, + user_id=webchat_conv_id, + content={"type": "user", "message": message_parts_for_storage}, + sender_id=username, + sender_name=username, + llm_checkpoint_id=llm_checkpoint_id, + ) + + response = cast( + QuartResponse, + await make_response( + stream(), + { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Transfer-Encoding": "chunked", + "Connection": "keep-alive", + }, + ), ) response.timeout = None return response @@ -568,6 +841,10 @@ class ChatRoute(Route): user_id=session_id, offset_sec=99999999, ) + thread_ids = await self.db.delete_webchat_threads_by_parent_session(session_id) + await self._delete_threads_by_ids(thread_ids, username) + + # 删除与会话关联的配置路由 try: router = self.umop_config_router assert router is not None @@ -739,8 +1016,14 @@ class ChatRoute(Route): page_size=1000, ) history_res = [history.model_dump() for history in history_ls] - response_data: dict[str, Any] = { + threads = await self.db.get_webchat_threads_by_parent_session( + parent_session_id=session_id, + creator=username, + ) + + response_data = { "history": history_res, + "threads": [self._serialize_thread(thread) for thread in threads], "is_running": self.running_convs.get(session_id, False), } if project_info: @@ -751,6 +1034,396 @@ class ChatRoute(Route): } return Response().ok(data=response_data).to_json() + async def create_thread(self): + """Create or reuse a side thread from a selected assistant message.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + session_id = post_data.get("session_id") + parent_message_id = post_data.get("parent_message_id") + selected_text = str(post_data.get("selected_text") or "").strip() + if not session_id: + return Response().error("Missing key: session_id").__dict__ + if parent_message_id is None: + return Response().error("Missing key: parent_message_id").__dict__ + if not selected_text: + return Response().error("Missing key: selected_text").__dict__ + + try: + parent_message_id = int(parent_message_id) + except (TypeError, ValueError): + return Response().error("Invalid key: parent_message_id").__dict__ + + username = g.get("username", "guest") + session = await self.db.get_platform_session_by_id(session_id) + if not session: + return Response().error(f"Session {session_id} not found").__dict__ + if session.creator != username: + return Response().error("Permission denied").__dict__ + + parent_record = await self.db.get_platform_message_history_by_id( + parent_message_id + ) + if ( + not parent_record + or parent_record.platform_id != session.platform_id + or parent_record.user_id != session_id + ): + return Response().error("Parent message not found").__dict__ + if not isinstance(parent_record.content, dict): + return Response().error("Invalid parent message content").__dict__ + if parent_record.content.get("type") != "bot": + return Response().error("Only bot messages can create threads").__dict__ + + checkpoint_id = parent_record.llm_checkpoint_id + if not checkpoint_id: + return ( + Response().error("Parent message is not linked to LLM history").__dict__ + ) + + existing = await self.db.get_webchat_thread_by_parent_message_and_text( + parent_session_id=session_id, + parent_message_id=parent_message_id, + selected_text=selected_text, + creator=username, + ) + if existing: + return Response().ok(data=self._serialize_thread(existing)).__dict__ + + conversation_id, history = await self._load_current_conversation_history( + session + ) + turn_range = self._find_turn_range(history, checkpoint_id) + if not conversation_id or not turn_range: + return Response().error("Linked checkpoint not found").__dict__ + + _start, end = turn_range + base_history = history[: end + 1] + thread = await self.db.create_webchat_thread( + creator=username, + parent_session_id=session_id, + parent_message_id=parent_message_id, + base_checkpoint_id=checkpoint_id, + selected_text=selected_text, + ) + await self.conv_mgr.new_conversation( + unified_msg_origin=self._build_thread_unified_msg_origin( + username, + thread.thread_id, + ), + platform_id="webchat", + content=base_history, + ) + return Response().ok(data=self._serialize_thread(thread)).__dict__ + + async def get_thread(self): + """Get a side thread and its message history.""" + thread_id = request.args.get("thread_id") + if not thread_id: + return Response().error("Missing key: thread_id").__dict__ + + username = g.get("username", "guest") + thread = await self.db.get_webchat_thread_by_id(thread_id) + if not thread: + return Response().error(f"Thread {thread_id} not found").__dict__ + if thread.creator != username: + return Response().error("Permission denied").__dict__ + + history_ls = await self.platform_history_mgr.get( + platform_id="webchat_thread", + user_id=thread_id, + page=1, + page_size=1000, + ) + return ( + Response() + .ok( + data={ + "thread": self._serialize_thread(thread), + "history": [history.model_dump() for history in history_ls], + "is_running": self.running_convs.get(thread_id, False), + } + ) + .__dict__ + ) + + async def send_thread_message(self): + """Send a message inside a WebChat side thread.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + thread_id = post_data.get("thread_id") + if not thread_id: + return Response().error("Missing key: thread_id").__dict__ + + username = g.get("username", "guest") + thread = await self.db.get_webchat_thread_by_id(thread_id) + if not thread: + return Response().error(f"Thread {thread_id} not found").__dict__ + if thread.creator != username: + return Response().error("Permission denied").__dict__ + + return await self.chat( + { + "session_id": thread.thread_id, + "message": post_data.get("message", []), + "enable_streaming": post_data.get("enable_streaming", True), + "selected_provider": post_data.get("selected_provider"), + "selected_model": post_data.get("selected_model"), + "_platform_history_id": "webchat_thread", + "_thread_selected_text": thread.selected_text, + } + ) + + async def delete_thread(self): + """Delete a WebChat side thread and its isolated history.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + thread_id = post_data.get("thread_id") + if not thread_id: + return Response().error("Missing key: thread_id").__dict__ + + username = g.get("username", "guest") + thread = await self.db.get_webchat_thread_by_id(thread_id) + if not thread: + return Response().error(f"Thread {thread_id} not found").__dict__ + if thread.creator != username: + return Response().error("Permission denied").__dict__ + + await self.db.delete_webchat_thread(thread_id) + await self._delete_threads_by_ids([thread_id], username) + return Response().ok(data={"thread_id": thread_id}).__dict__ + + async def update_message(self): + """Update a persisted WebChat message and its linked LLM turn.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + session_id = post_data.get("session_id") + message_id = post_data.get("message_id") + content = post_data.get("content") + if not session_id: + return Response().error("Missing key: session_id").__dict__ + if message_id is None: + return Response().error("Missing key: message_id").__dict__ + + try: + message_id = int(message_id) + content = self._sanitize_message_content(content) + except (TypeError, ValueError) as exc: + return Response().error(str(exc)).__dict__ + + username = g.get("username", "guest") + session = await self.db.get_platform_session_by_id(session_id) + if not session: + return Response().error(f"Session {session_id} not found").__dict__ + if session.creator != username: + return Response().error("Permission denied").__dict__ + + record = await self.db.get_platform_message_history_by_id(message_id) + if not record: + return Response().error(f"Message {message_id} not found").__dict__ + if record.platform_id != session.platform_id or record.user_id != session_id: + return Response().error("Message does not belong to the session").__dict__ + if not isinstance(record.content, dict): + return Response().error("Invalid message content").__dict__ + if record.content.get("type") != content.get("type"): + return Response().error("Message type cannot be changed").__dict__ + if content.get("type") != "user": + return Response().error("Only user messages can be edited").__dict__ + + platform_history = await self._get_sorted_platform_history(session) + latest_user_record = next( + ( + item + for item in reversed(platform_history) + if isinstance(item.content, dict) and item.content.get("type") == "user" + ), + None, + ) + if not latest_user_record or latest_user_record.id != message_id: + return ( + Response().error("Only the latest user message can be edited").__dict__ + ) + + checkpoint_id = record.llm_checkpoint_id + if not checkpoint_id: + return ( + Response() + .error("This message is not linked to LLM history and cannot be edited") + .__dict__ + ) + + conversation_id, history = await self._load_current_conversation_history( + session + ) + turn_range = self._find_turn_range(history, checkpoint_id) + if not conversation_id or not turn_range: + return Response().error("Linked checkpoint not found").__dict__ + if not self._is_latest_checkpoint(history, checkpoint_id): + return Response().error("Only the latest turn can be edited").__dict__ + + start, _end = turn_range + + target_index = self._find_turn_user_index(history, start, _end) + if target_index is None: + return Response().error("Linked user message not found").__dict__ + + new_checkpoint_id = str(uuid.uuid4()) + truncated_history = history[:start] + await self.platform_history_mgr.update( + message_id=message_id, + content=content, + llm_checkpoint_id=new_checkpoint_id, + ) + deleted_message_ids = await self._delete_platform_history_after( + session, message_id + ) + thread_ids = await self.db.delete_webchat_threads_by_parent_message_ids( + session_id, + deleted_message_ids, + ) + await self._delete_threads_by_ids(thread_ids, username) + await self.conv_mgr.update_conversation( + unified_msg_origin=self._build_webchat_unified_msg_origin(session), + conversation_id=conversation_id, + history=truncated_history, + ) + await self.db.update_platform_session(session_id=session_id) + updated = await self.db.get_platform_message_history_by_id(message_id) + return ( + Response() + .ok( + data={ + "message": updated.model_dump() if updated else None, + "needs_regenerate": True, + "truncated_after_message": True, + } + ) + .__dict__ + ) + + async def regenerate_message(self): + """Regenerate the latest bot message linked to an LLM checkpoint.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + session_id = post_data.get("session_id") + message_id = post_data.get("message_id") + if not session_id: + return Response().error("Missing key: session_id").__dict__ + if message_id is None: + return Response().error("Missing key: message_id").__dict__ + + try: + message_id = int(message_id) + except (TypeError, ValueError): + return Response().error("Invalid key: message_id").__dict__ + + username = g.get("username", "guest") + session = await self.db.get_platform_session_by_id(session_id) + if not session: + return Response().error(f"Session {session_id} not found").__dict__ + if session.creator != username: + return Response().error("Permission denied").__dict__ + + target_record = await self.db.get_platform_message_history_by_id(message_id) + if not target_record: + return Response().error(f"Message {message_id} not found").__dict__ + if ( + target_record.platform_id != session.platform_id + or target_record.user_id != session_id + ): + return Response().error("Message does not belong to the session").__dict__ + if not isinstance(target_record.content, dict): + return Response().error("Invalid message content").__dict__ + if target_record.content.get("type") != "bot": + return Response().error("Only bot messages can be regenerated").__dict__ + + checkpoint_id = target_record.llm_checkpoint_id + if not checkpoint_id: + return Response().error("Message is not linked to LLM history").__dict__ + + conversation_id, history = await self._load_current_conversation_history( + session + ) + turn_range = self._find_turn_range(history, checkpoint_id) + if not conversation_id or not turn_range: + return Response().error("Linked checkpoint not found").__dict__ + if not self._is_latest_checkpoint(history, checkpoint_id): + return ( + Response().error("Regenerating older turns requires branching").__dict__ + ) + + start, end = turn_range + user_index = self._find_turn_user_index(history, start, end) + if user_index is None: + return Response().error("Linked user message not found").__dict__ + + platform_history = await self._get_sorted_platform_history(session) + source_user_record = next( + ( + item + for item in reversed(platform_history) + if item.llm_checkpoint_id == checkpoint_id + and isinstance(item.content, dict) + and item.content.get("type") == "user" + ), + None, + ) + if not source_user_record: + return Response().error("Linked user display message not found").__dict__ + + old_bot_record_ids = [ + item.id + for item in platform_history + if item.id is not None + and item.llm_checkpoint_id == checkpoint_id + and isinstance(item.content, dict) + and item.content.get("type") == "bot" + ] + if not old_bot_record_ids: + return Response().error("Linked bot display message not found").__dict__ + + new_checkpoint_id = str(uuid.uuid4()) + # The WebChat send path adds the current user message from the prompt. + # Remove the whole old turn here to avoid duplicating that user message. + new_history = history[:start] + history[end + 1 :] + await self.conv_mgr.update_conversation( + unified_msg_origin=self._build_webchat_unified_msg_origin(session), + conversation_id=conversation_id, + history=new_history, + ) + thread_ids = await self.db.delete_webchat_threads_by_parent_message_ids( + session_id, + old_bot_record_ids, + ) + await self._delete_threads_by_ids(thread_ids, username) + for old_bot_record_id in old_bot_record_ids: + await self.platform_history_mgr.delete_by_id(old_bot_record_id) + await self.platform_history_mgr.update( + message_id=source_user_record.id, + llm_checkpoint_id=new_checkpoint_id, + ) + + return await self.chat( + { + "session_id": session_id, + "message": source_user_record.content.get("message", []), + "enable_streaming": post_data.get("enable_streaming", True), + "selected_provider": post_data.get("selected_provider"), + "selected_model": post_data.get("selected_model"), + "_skip_user_history": True, + "_llm_checkpoint_id": new_checkpoint_id, + } + ) + async def update_session_display_name(self): """Update a Platform session's display name.""" post_data = await request.json diff --git a/astrbot/dashboard/routes/live_chat.py b/astrbot/dashboard/routes/live_chat.py index eef4dba74..f9dafc489 100644 --- a/astrbot/dashboard/routes/live_chat.py +++ b/astrbot/dashboard/routes/live_chat.py @@ -323,8 +323,9 @@ class LiveChatRoute(Route): text: str, media_parts: list, reasoning: str, - agent_stats: dict[str, Any], - refs: dict[str, Any], + agent_stats: dict, + refs: dict, + llm_checkpoint_id: str | None = None, ): """保存 bot 消息到历史记录｡""" bot_message_parts = [] @@ -348,6 +349,7 @@ class LiveChatRoute(Route): content=new_his, sender_id="bot", sender_name="bot", + llm_checkpoint_id=llm_checkpoint_id, ) async def _send_chat_payload(self, session: LiveChatSession, payload: dict) -> None: @@ -531,6 +533,7 @@ class LiveChatRoute(Route): session.is_processing = True session.should_interrupt = False back_queue = webchat_queue_mgr.get_or_create_back_queue(message_id, session_id) + llm_checkpoint_id = str(uuid.uuid4()) try: chat_queue = webchat_queue_mgr.get_or_create_queue(session_id) @@ -548,19 +551,31 @@ class LiveChatRoute(Route): "show_reasoning": show_reasoning, "enable_streaming": enable_streaming, "message_id": message_id, + "llm_checkpoint_id": llm_checkpoint_id, }, ), ) message_parts_for_storage = strip_message_parts_path_fields(message_parts) - mgr = self.platform_history_mgr - assert mgr is not None - await mgr.insert( + saved_user_record = await self.platform_history_mgr.insert( platform_id="webchat", user_id=session_id, content={"type": "user", "message": message_parts_for_storage}, sender_id=session.username, sender_name=session.username, + llm_checkpoint_id=llm_checkpoint_id, + ) + await self._send_chat_payload( + session, + { + "ct": "chat", + "type": "user_message_saved", + "data": { + "id": saved_user_record.id, + "created_at": to_utc_isoformat(saved_user_record.created_at), + "llm_checkpoint_id": llm_checkpoint_id, + }, + }, ) accumulated_parts = [] @@ -702,6 +717,7 @@ class LiveChatRoute(Route): accumulated_reasoning, agent_stats, refs, + llm_checkpoint_id, ) if saved_record: await self._send_chat_payload( @@ -714,6 +730,7 @@ class LiveChatRoute(Route): "created_at": to_utc_isoformat( saved_record.created_at, ), + "llm_checkpoint_id": llm_checkpoint_id, }, }, ) diff --git a/changelogs/v4.23.2.md b/changelogs/v4.23.2.md new file mode 100644 index 000000000..d6e91acea --- /dev/null +++ b/changelogs/v4.23.2.md @@ -0,0 +1,69 @@ +- [更新日志(简体中文)](#chinese) +- [Changelog(English)](#english) + + + +## What's Changed + +### 新增 + +- 知识库稀疏检索阶段新增 SQLite FTS5 支持，大幅优化万到十万级别分块时造成的召回时的显著卡顿。（[#7648](https://github.com/AstrBotDevs/AstrBot/pull/7648)） +- KOOK 平台新增角色提及支持，包含角色记录、事件解析、API 类型与相关测试。（[#7626](https://github.com/AstrBotDevs/AstrBot/pull/7626)） +- 新增 MiniMax Token Plan Provider。（[#7609](https://github.com/AstrBotDevs/AstrBot/pull/7609)） +- 新增 `on_agent_begin`、`on_using_llm_tool`、`on_llm_tool_respond`、`on_agent_done` 插件事件钩子，并更新插件开发文档。（[#7540](https://github.com/AstrBotDevs/AstrBot/pull/7540)） +- 为 t2i 模板启用 Shiki 代码高亮，并新增 VitePress 风格模板。（[#7501](https://github.com/AstrBotDevs/AstrBot/pull/7501)） + +### 优化 + +- 优化 `EmptyModelOutputError` 处理，不再严格强制校验模型工具输出的名字是否符合工具集，采用引导的方式让模型输出正确的工具调用。（[#7375](https://github.com/AstrBotDevs/AstrBot/pull/7375)） +- 优化 `SendMessageToUserTool` 工具在本地 workspace root 内解析相对文件路径的问题。（[#7668](https://github.com/AstrBotDevs/AstrBot/pull/7668)） + +### 修复 + +- 修复 Bocha 搜索返回 Brotli 压缩内容时出现 `Can not decode content-encoding: br` 的问题。（[#7655](https://github.com/AstrBotDevs/AstrBot/pull/7655)） +- 修复微信个人号主动 Cron 发送时 `context_token` 未持久化的问题。（[#7595](https://github.com/AstrBotDevs/AstrBot/pull/7595)） +- 修复 Anthropic Provider 默认 `max_tokens` 偏小导致工具调用输出文件操作时造成大量截断的问题。（[#7593](https://github.com/AstrBotDevs/AstrBot/pull/7593)） +- 修复 Cron `last_run_at` 在 Dashboard 中未按本地时区显示的问题。（[#7625](https://github.com/AstrBotDevs/AstrBot/pull/7625)） +- 修复 Cron 工具调度失败时静默继续的问题，并在失败时返回明确错误给模型。（[#7513](https://github.com/AstrBotDevs/AstrBot/pull/7513)） +- 修复 Telegram 媒体组异常被静默吞掉的问题。（[#7537](https://github.com/AstrBotDevs/AstrBot/pull/7537)） +- 修复会话限流计数为 0 时可能触发 `IndexError` 的问题。（[#7635](https://github.com/AstrBotDevs/AstrBot/pull/7635)） + +- 修复知识库上传错误处理与错误提示，提升失败时的可诊断性。（[#7534](https://github.com/AstrBotDevs/AstrBot/pull/7534), [#7536](https://github.com/AstrBotDevs/AstrBot/pull/7536)） +- 修复 Dashboard 聊天附件 401 问题。（[#7569](https://github.com/AstrBotDevs/AstrBot/pull/7569)） +- 修复 Dashboard 数字输入框在未编辑时失焦会重置为 0 的问题。（[#7560](https://github.com/AstrBotDevs/AstrBot/pull/7560)） +- 修复 Dashboard 列表项内代码块暗色模式传递问题，并移除 `ThemeAwareMarkdownCodeBlock` 中注入的默认值。（[#7667](https://github.com/AstrBotDevs/AstrBot/pull/7667), [commit](https://github.com/AstrBotDevs/AstrBot/commit/fd2ca702d)） +- 修复 Updater 请求不支持 SOCKS 代理的问题，并补充相关测试。（[#7615](https://github.com/AstrBotDevs/AstrBot/pull/7615)） +- 回滚 SCSS import warning 修复以避免相关兼容问题。（[#7616](https://github.com/AstrBotDevs/AstrBot/pull/7616)） + + + +## What's Changed (EN) + +### New Features + +- Added SQLite FTS5 support to the knowledge-base sparse retrieval stage, greatly reducing recall latency for knowledge bases with tens of thousands to hundreds of thousands of chunks. ([#7648](https://github.com/AstrBotDevs/AstrBot/pull/7648)) +- Added KOOK role mention support, including role records, event parsing, API types, and tests. ([#7626](https://github.com/AstrBotDevs/AstrBot/pull/7626)) +- Added the MiniMax Token Plan Provider. ([#7609](https://github.com/AstrBotDevs/AstrBot/pull/7609)) +- Added plugin event hooks for `on_agent_begin`, `on_using_llm_tool`, `on_llm_tool_respond`, and `on_agent_done`, with updated plugin development docs. ([#7540](https://github.com/AstrBotDevs/AstrBot/pull/7540)) +- Enabled Shiki highlighting for t2i templates and added a VitePress-style template. ([#7501](https://github.com/AstrBotDevs/AstrBot/pull/7501)) + +### Improvements + +- Improved `EmptyModelOutputError` handling by no longer strictly enforcing whether model tool output names match the tool set, and instead guiding the model toward valid tool calls. ([#7375](https://github.com/AstrBotDevs/AstrBot/pull/7375)) +- Improved relative file path resolution within the local workspace root for `SendMessageToUserTool`. ([#7668](https://github.com/AstrBotDevs/AstrBot/pull/7668)) + +### Bug Fixes + +- Fixed Bocha search failures caused by Brotli-compressed responses reporting `Can not decode content-encoding: br`. ([#7655](https://github.com/AstrBotDevs/AstrBot/pull/7655)) +- Fixed `context_token` persistence for proactive Cron sends in the Weixin OC adapter. ([#7595](https://github.com/AstrBotDevs/AstrBot/pull/7595)) +- Increased the Anthropic Provider default `max_tokens` value to avoid heavy truncation during tool-call output for file operations. ([#7593](https://github.com/AstrBotDevs/AstrBot/pull/7593)) +- Displayed Cron `last_run_at` in the local timezone in Dashboard. ([#7625](https://github.com/AstrBotDevs/AstrBot/pull/7625)) +- Returned explicit errors to the model when Cron tool scheduling fails instead of continuing silently. ([#7513](https://github.com/AstrBotDevs/AstrBot/pull/7513)) +- Prevented Telegram media group exceptions from being silently swallowed. ([#7537](https://github.com/AstrBotDevs/AstrBot/pull/7537)) +- Fixed an `IndexError` when session `rate_limit_count` is 0. ([#7635](https://github.com/AstrBotDevs/AstrBot/pull/7635)) +- Improved knowledge-base upload error handling and error messages for better diagnostics. ([#7534](https://github.com/AstrBotDevs/AstrBot/pull/7534), [#7536](https://github.com/AstrBotDevs/AstrBot/pull/7536)) +- Fixed Dashboard chat attachment 401 errors. ([#7569](https://github.com/AstrBotDevs/AstrBot/pull/7569)) +- Fixed numeric inputs resetting to 0 on blur when the value was not edited. ([#7560](https://github.com/AstrBotDevs/AstrBot/pull/7560)) +- Fixed dark-mode propagation for code blocks inside Dashboard list items and removed the injected default value from `ThemeAwareMarkdownCodeBlock`. ([#7667](https://github.com/AstrBotDevs/AstrBot/pull/7667), [commit](https://github.com/AstrBotDevs/AstrBot/commit/fd2ca702d)) +- Added SOCKS proxy support for updater requests and related tests. ([#7615](https://github.com/AstrBotDevs/AstrBot/pull/7615)) +- Reverted the SCSS import warning fix to avoid related compatibility issues. ([#7616](https://github.com/AstrBotDevs/AstrBot/pull/7616)) diff --git a/dashboard/pnpm-lock.yaml b/dashboard/pnpm-lock.yaml index ae65710eb..c60569271 100644 --- a/dashboard/pnpm-lock.yaml +++ b/dashboard/pnpm-lock.yaml @@ -5,8 +5,8 @@ settings: excludeLinksFromLockfile: false overrides: - immutable: 4.3.8 - lodash-es: 4.17.23 + immutable: "4.3.8" + lodash-es: "4.17.23" importers: diff --git a/dashboard/src/components/chat/ChatMessageList.vue b/dashboard/src/components/chat/ChatMessageList.vue new file mode 100644 index 000000000..a0e70d65f --- /dev/null +++ b/dashboard/src/components/chat/ChatMessageList.vue @@ -0,0 +1,1391 @@ + + + + + diff --git a/dashboard/src/components/chat/ProviderConfigDialog.vue b/dashboard/src/components/chat/ProviderConfigDialog.vue index 85bad2d93..dfa919904 100644 --- a/dashboard/src/components/chat/ProviderConfigDialog.vue +++ b/dashboard/src/components/chat/ProviderConfigDialog.vue @@ -141,7 +141,6 @@ @delete-provider="deleteProvider" @add-model-provider="addModelProvider" /> - >>>>>>> origin ======= - >>>>>>> origin
+ + + + + + {{ tm("actions.retry") }} + + + + + + + + + + {{ provider.id }} + + + {{ provider.model }} + + + mdi-eye-outline + + + mdi-music-note-outline + + + mdi-wrench + + + mdi-brain + + + + + +
+ {{ tm("actions.noAvailableModels") }} +
+ + + + + + + + + diff --git a/dashboard/src/components/chat/ThreadPanel.vue b/dashboard/src/components/chat/ThreadPanel.vue new file mode 100644 index 000000000..cd0431309 --- /dev/null +++ b/dashboard/src/components/chat/ThreadPanel.vue @@ -0,0 +1,543 @@ + + + + + diff --git a/dashboard/src/components/chat/ThreadedMarkdownMessagePart.vue b/dashboard/src/components/chat/ThreadedMarkdownMessagePart.vue new file mode 100644 index 000000000..96e25087c --- /dev/null +++ b/dashboard/src/components/chat/ThreadedMarkdownMessagePart.vue @@ -0,0 +1,89 @@ + + + diff --git a/dashboard/src/components/chat/message_list_comps/ActionRef.vue b/dashboard/src/components/chat/message_list_comps/ActionRef.vue index 0fc681710..fabcf7596 100644 --- a/dashboard/src/components/chat/message_list_comps/ActionRef.vue +++ b/dashboard/src/components/chat/message_list_comps/ActionRef.vue @@ -22,7 +22,7 @@ +{{ refs.used.length - 3 }} - + {{ tm("refs.sources") }}
@@ -64,10 +64,13 @@ export default { .refs-container { display: flex; align-items: center; - margin-left: 8px; - padding: 4px 8px; - border-radius: 12px; + min-height: 24px; + padding: 0 6px; + border-radius: 8px; + color: inherit; cursor: pointer; + font-size: 12px; + line-height: 24px; transition: background-color; } @@ -79,6 +82,7 @@ export default { display: flex; align-items: center; position: relative; + min-height: 24px; } .ref-avatar { @@ -86,7 +90,6 @@ export default { height: 20px; border-radius: 50%; opacity: 0.9; - background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); display: flex; align-items: center; justify-content: center; @@ -118,4 +121,11 @@ export default { opacity: 0.7; font-weight: 500; } + +.refs-label { + margin-left: 6px; + color: inherit; + font-size: 12px; + line-height: 24px; +} diff --git a/dashboard/src/components/chat/message_list_comps/RefsSidebar.vue b/dashboard/src/components/chat/message_list_comps/RefsSidebar.vue index 118c797d5..90f12555e 100644 --- a/dashboard/src/components/chat/message_list_comps/RefsSidebar.vue +++ b/dashboard/src/components/chat/message_list_comps/RefsSidebar.vue @@ -133,14 +133,16 @@ export default { display: flex; justify-content: space-between; align-items: center; - padding: 16px 20px; + padding: 14px 16px 8px; flex-shrink: 0; } .sidebar-title { - font-size: 18px; + font-size: 16px; font-weight: 600; - color: var(--v-theme-primaryText); + color: rgb(var(--v-theme-on-surface)); + line-height: 1.4; + margin: 0; } .refs-list { @@ -176,7 +178,6 @@ export default { display: flex; align-items: center; justify-content: center; - background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); } .ref-item-favicon { diff --git a/dashboard/src/components/chat/message_list_comps/ThreadNode.vue b/dashboard/src/components/chat/message_list_comps/ThreadNode.vue new file mode 100644 index 000000000..6e1aad62b --- /dev/null +++ b/dashboard/src/components/chat/message_list_comps/ThreadNode.vue @@ -0,0 +1,63 @@ + + + + + diff --git a/dashboard/src/components/shared/ThemeAwareMarkdownCodeBlock.vue b/dashboard/src/components/shared/ThemeAwareMarkdownCodeBlock.vue index f90127218..902cc5111 100644 --- a/dashboard/src/components/shared/ThemeAwareMarkdownCodeBlock.vue +++ b/dashboard/src/components/shared/ThemeAwareMarkdownCodeBlock.vue @@ -13,7 +13,7 @@ diff --git a/dashboard/src/i18n/locales/en-US/features/alkaid/knowledge-base.json b/dashboard/src/i18n/locales/en-US/features/alkaid/knowledge-base.json index f4d8a33fd..98d5aea37 100644 --- a/dashboard/src/i18n/locales/en-US/features/alkaid/knowledge-base.json +++ b/dashboard/src/i18n/locales/en-US/features/alkaid/knowledge-base.json @@ -70,7 +70,7 @@ }, "upload": { "title": "Upload Files to Knowledge Base", - "subtitle": "Supports txt, pdf, word, excel and other formats", + "subtitle": "Supports txt, pdf, epub, word, excel and other formats", "dropzone": "Drag and drop files here or click to upload", "chunkSettings": { "title": "Chunk Settings", @@ -152,4 +152,4 @@ "preRequisite": "Hint: Please go to the plugin market to install astrbot_plugin_url_2_knowledge_base and follow the instructions in the plugin documentation to complete the playwright installation before using this feature.", "allChunksUploaded": "All chunks uploaded successfully" } -} \ No newline at end of file +} diff --git a/dashboard/src/i18n/locales/en-US/features/chat.json b/dashboard/src/i18n/locales/en-US/features/chat.json index 4b87cc8ab..ccbe35aae 100644 --- a/dashboard/src/i18n/locales/en-US/features/chat.json +++ b/dashboard/src/i18n/locales/en-US/features/chat.json @@ -39,6 +39,9 @@ "actions": { "copy": "Copy", "regenerate": "Regenerate", + "retry": "Retry", + "retryWithModel": "Use another model", + "noAvailableModels": "No available models", "like": "Like", "dislike": "Dislike", "share": "Share", @@ -60,6 +63,15 @@ "done": "Done", "running": "Running" }, + "thread": { + "title": "Thread", + "askInThread": "Ask in thread", + "count": "{count} threads", + "createFailed": "Failed to create thread", + "delete": "Delete thread", + "confirmDelete": "Delete this thread? This action cannot be undone.", + "placeholder": "Ask about this excerpt..." + }, "conversation": { "newConversation": "New Conversation", "noHistory": "No conversation history", diff --git a/dashboard/src/i18n/locales/en-US/features/config-metadata.json b/dashboard/src/i18n/locales/en-US/features/config-metadata.json index 105e731b4..4f8a1f91f 100644 --- a/dashboard/src/i18n/locales/en-US/features/config-metadata.json +++ b/dashboard/src/i18n/locales/en-US/features/config-metadata.json @@ -292,6 +292,10 @@ "description": "Output Tool Call Results", "hint": "Only takes effect when \"Output Function Call Status\" is enabled, and shows at most 70 characters." }, + "buffer_intermediate_messages": { + "description": "Merge Agent Intermediate Messages", + "hint": "When enabled, intermediate text generated during multi-step tool calls in non-streaming mode will be buffered and sent as a single merged reply after the Agent finishes." + }, "sanitize_context_by_modalities": { "description": "Sanitize History by Modalities", "hint": "When enabled, sanitizes contexts before each LLM request by removing image blocks and tool-call structures that the current provider's modalities do not support (this changes what the model sees)." diff --git a/dashboard/src/i18n/locales/en-US/features/extension.json b/dashboard/src/i18n/locales/en-US/features/extension.json index 233f28dd5..0a4a7fe48 100644 --- a/dashboard/src/i18n/locales/en-US/features/extension.json +++ b/dashboard/src/i18n/locales/en-US/features/extension.json @@ -144,7 +144,8 @@ "updated": "Last Updated", "updateStatus": "Update Status", "ascending": "Ascending", - "descending": "Descending" + "descending": "Descending", + "pinUpdatesOnTop": "Pin Updates on Top" }, "tags": { "danger": "Danger" diff --git a/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json b/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json index 90d3e6158..0ef8fa3d4 100644 --- a/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json +++ b/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json @@ -46,7 +46,7 @@ "title": "Upload Document", "selectFile": "Select File", "dropzone": "Drop files here or click to select", - "supportedFormats": "Supported formats: ", + "supportedFormats": "Supported formats: .txt, .md, .pdf, .docx, .epub, .xls, .xlsx", "maxSize": "Max file size: 128MB", "chunkSettings": "Chunk Settings", "batchSettings": "Batch Settings", diff --git a/dashboard/src/i18n/locales/ru-RU/features/config-metadata.json b/dashboard/src/i18n/locales/ru-RU/features/config-metadata.json index 3c59e5eef..6be4a7026 100644 --- a/dashboard/src/i18n/locales/ru-RU/features/config-metadata.json +++ b/dashboard/src/i18n/locales/ru-RU/features/config-metadata.json @@ -292,6 +292,10 @@ "description": "Выводить результаты работы инструментов", "hint": "Работает только при включенном статусе вызова функций. Показывает макс. 70 символов." }, + "buffer_intermediate_messages": { + "description": "Объединять промежуточные сообщения Agent", + "hint": "Если включено, промежуточный текст, созданный во время многошаговых вызовов инструментов в непотоковом режиме, будет буферизован и отправлен одним объединенным ответом после завершения Agent." + }, "sanitize_context_by_modalities": { "description": "Очистка истории по модальностям", "hint": "Если включено, очищает контекст перед запросом, удаляя блоки (например, изображения), которые не поддерживаются выбранным провайдером." diff --git a/dashboard/src/i18n/locales/ru-RU/features/extension.json b/dashboard/src/i18n/locales/ru-RU/features/extension.json index d473de830..c085885db 100644 --- a/dashboard/src/i18n/locales/ru-RU/features/extension.json +++ b/dashboard/src/i18n/locales/ru-RU/features/extension.json @@ -143,7 +143,8 @@ "updated": "Дате обновления", "updateStatus": "Статусу обновления", "ascending": "По возрастанию", - "descending": "По убыванию" + "descending": "По убыванию", + "pinUpdatesOnTop": "Обновления сверху" }, "tags": { "danger": "Опасно" diff --git a/dashboard/src/i18n/locales/ru-RU/features/knowledge-base/detail.json b/dashboard/src/i18n/locales/ru-RU/features/knowledge-base/detail.json index 2720bb2a5..303b0e1c8 100644 --- a/dashboard/src/i18n/locales/ru-RU/features/knowledge-base/detail.json +++ b/dashboard/src/i18n/locales/ru-RU/features/knowledge-base/detail.json @@ -46,7 +46,7 @@ "title": "Добавление контента", "selectFile": "Файл", "dropzone": "Нажмите или перетащите файл сюда", - "supportedFormats": "Форматы: ", + "supportedFormats": "Форматы: .txt, .md, .pdf, .docx, .epub, .xls, .xlsx", "maxSize": "Максимум: 128MB", "chunkSettings": "Фрагментация", "batchSettings": "Пакетная обработка", @@ -115,4 +115,4 @@ "saveFailed": "Ошибка сохранения", "tips": "Внимание! Изменение этих параметров повлияет на будущую выдачу базы знаний." } -} \ No newline at end of file +} diff --git a/dashboard/src/i18n/locales/zh-CN/features/alkaid/knowledge-base.json b/dashboard/src/i18n/locales/zh-CN/features/alkaid/knowledge-base.json index 2467b1efc..d4c7ea839 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/alkaid/knowledge-base.json +++ b/dashboard/src/i18n/locales/zh-CN/features/alkaid/knowledge-base.json @@ -70,7 +70,7 @@ }, "upload": { "title": "上传文件到知识库", - "subtitle": "支持 txt、pdf、word、excel 等多种格式", + "subtitle": "支持 txt、pdf、epub、word、excel 等多种格式", "dropzone": "拖放文件到这里或点击上传", "chunkSettings": { "title": "分片设置", @@ -152,4 +152,4 @@ "preRequisite": "提示：请先前往插件市场安装 astrbot_plugin_url_2_knowledge_base 并根据插件文档内的指示完成 playwright 安装后才可使用本功能", "allChunksUploaded": "所有分片上传成功" } -} \ No newline at end of file +} diff --git a/dashboard/src/i18n/locales/zh-CN/features/chat.json b/dashboard/src/i18n/locales/zh-CN/features/chat.json index 1c3eea444..cc837015c 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/chat.json +++ b/dashboard/src/i18n/locales/zh-CN/features/chat.json @@ -39,6 +39,9 @@ "actions": { "copy": "复制", "regenerate": "重新生成", + "retry": "重试", + "retryWithModel": "使用其他模型", + "noAvailableModels": "暂无可用模型", "like": "点赞", "dislike": "踩", "share": "分享", @@ -60,6 +63,15 @@ "done": "已完成", "running": "运行中" }, + "thread": { + "title": "分支", + "askInThread": "在新分支中询问", + "count": "{count} 个分支", + "createFailed": "创建分支失败", + "delete": "删除分支", + "confirmDelete": "确定要删除这个分支吗？此操作无法撤销。", + "placeholder": "继续追问这段内容..." + }, "conversation": { "newConversation": "新的聊天", "noHistory": "暂无对话历史", diff --git a/dashboard/src/i18n/locales/zh-CN/features/config-metadata.json b/dashboard/src/i18n/locales/zh-CN/features/config-metadata.json index 3a721de75..e9a21bcc7 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/config-metadata.json +++ b/dashboard/src/i18n/locales/zh-CN/features/config-metadata.json @@ -290,6 +290,10 @@ "description": "输出函数调用返回结果", "hint": "仅在启用“输出函数调用状态”时生效，且最多展示 70 个字符。" }, + "buffer_intermediate_messages": { + "description": "合并 Agent 中间消息", + "hint": "开启后，非流式模式下多步工具调用过程中产生的中间文本将缓冲，待 Agent 完成后合并为一条回复发送。" + }, "sanitize_context_by_modalities": { "description": "按模型能力清理历史上下文", "hint": "开启后，在每次请求 LLM 前会按当前模型提供商中所选择的模型能力删除对话中不支持的图片/工具调用结构（会改变模型看到的历史）" @@ -1651,4 +1655,4 @@ "helpMiddle": "或", "helpSuffix": "。" } -} \ No newline at end of file +} diff --git a/dashboard/src/i18n/locales/zh-CN/features/extension.json b/dashboard/src/i18n/locales/zh-CN/features/extension.json index 04eaa8bfa..c06f38142 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/extension.json +++ b/dashboard/src/i18n/locales/zh-CN/features/extension.json @@ -147,7 +147,8 @@ "updated": "更新时间", "updateStatus": "更新状态", "ascending": "升序", - "descending": "降序" + "descending": "降序", + "pinUpdatesOnTop": "有更新置顶" }, "tags": { "danger": "危险" diff --git a/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json b/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json index deb61f933..7473c6c1e 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json +++ b/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json @@ -46,7 +46,7 @@ "title": "上传文档", "selectFile": "选择文件", "dropzone": "拖放文件到这里或点击选择", - "supportedFormats": "支持的格式: ", + "supportedFormats": "支持的格式: .txt, .md, .pdf, .docx, .epub, .xls, .xlsx", "maxSize": "最大文件大小: 128MB", "chunkSettings": "分块设置", "batchSettings": "批处理设置", diff --git a/dashboard/src/views/extension/extensionPreferenceStorage.mjs b/dashboard/src/views/extension/extensionPreferenceStorage.mjs new file mode 100644 index 000000000..31cd90bbe --- /dev/null +++ b/dashboard/src/views/extension/extensionPreferenceStorage.mjs @@ -0,0 +1,84 @@ +export const SHOW_RESERVED_PLUGINS_STORAGE_KEY = "showReservedPlugins"; +export const PLUGIN_LIST_VIEW_MODE_STORAGE_KEY = "pluginListViewMode"; +export const PIN_UPDATES_ON_TOP_STORAGE_KEY = "pinUpdatesOnTop"; + +/** + * Resolve the storage backend for reading preferences. + * Pass `null` to explicitly disable storage access in callers/tests. + */ +const getStorageForRead = (storageOverride) => { + if (storageOverride === null) { + return null; + } + if (storageOverride !== undefined) { + return typeof storageOverride?.getItem === "function" + ? storageOverride + : null; + } + if (typeof window === "undefined") { + return null; + } + try { + const localStorage = window.localStorage ?? null; + return typeof localStorage?.getItem === "function" ? localStorage : null; + } catch { + return null; + } +}; + +/** + * Resolve the storage backend for writing preferences. + * Pass `null` to explicitly disable storage access in callers/tests. + */ +const getStorageForWrite = (storageOverride) => { + if (storageOverride === null) { + return null; + } + if (storageOverride !== undefined) { + return typeof storageOverride?.setItem === "function" + ? storageOverride + : null; + } + if (typeof window === "undefined") { + return null; + } + try { + const localStorage = window.localStorage ?? null; + return typeof localStorage?.setItem === "function" ? localStorage : null; + } catch { + return null; + } +}; + +export const readBooleanPreference = (key, fallback, storage) => { + const targetStorage = getStorageForRead(storage); + if (!targetStorage) { + return fallback; + } + + try { + const saved = targetStorage.getItem(key); + if (saved === "true") { + return true; + } + if (saved === "false") { + return false; + } + return fallback; + } catch { + return fallback; + } +}; + +export const writeBooleanPreference = (key, value, storage) => { + const targetStorage = getStorageForWrite(storage); + if (!targetStorage) { + return; + } + + try { + targetStorage.setItem(key, String(value)); + } catch { + // Ignore restricted storage environments. + } +}; diff --git a/dashboard/src/views/knowledge-base/KBList.vue b/dashboard/src/views/knowledge-base/KBList.vue index f554472b4..010e5cdbb 100644 --- a/dashboard/src/views/knowledge-base/KBList.vue +++ b/dashboard/src/views/knowledge-base/KBList.vue @@ -13,7 +13,7 @@ variant="text" size="small" color="grey" - href="https://astrbot.app/use/knowledge-base.html" + href="https://docs.astrbot.app/use/knowledge-base.html" target="_blank" /> diff --git a/dashboard/tests/extensionPreferenceStorage.test.mjs b/dashboard/tests/extensionPreferenceStorage.test.mjs new file mode 100644 index 000000000..e864f5fae --- /dev/null +++ b/dashboard/tests/extensionPreferenceStorage.test.mjs @@ -0,0 +1,75 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; + +import { + PIN_UPDATES_ON_TOP_STORAGE_KEY, + readBooleanPreference, + writeBooleanPreference, +} from '../src/views/extension/extensionPreferenceStorage.mjs'; + +test("readBooleanPreference returns fallback when storage access throws", () => { + const storage = { + getItem() { + throw new Error("SecurityError"); + }, + }; + + assert.equal( + readBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, storage), + true, + ); +}); + +test("readBooleanPreference parses stored boolean strings", () => { + const storage = { + getItem(key) { + return key === PIN_UPDATES_ON_TOP_STORAGE_KEY ? "false" : null; + }, + }; + + assert.equal( + readBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, storage), + false, + ); +}); + +test("readBooleanPreference treats explicit null storage as unavailable", () => { + assert.equal( + readBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, null), + true, + ); +}); + +test("readBooleanPreference treats invalid storage overrides as unavailable", () => { + assert.equal( + readBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, {}), + true, + ); +}); + +test("writeBooleanPreference stores boolean strings and swallows storage errors", () => { + const writes = []; + const storage = { + setItem(key, value) { + writes.push([key, value]); + throw new Error("QuotaExceededError"); + }, + }; + + assert.doesNotThrow(() => + writeBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, storage), + ); + assert.deepEqual(writes, [[PIN_UPDATES_ON_TOP_STORAGE_KEY, "true"]]); +}); + +test("writeBooleanPreference ignores explicit null storage", () => { + assert.doesNotThrow(() => + writeBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, null), + ); +}); + +test("writeBooleanPreference ignores invalid storage overrides", () => { + assert.doesNotThrow(() => + writeBooleanPreference(PIN_UPDATES_ON_TOP_STORAGE_KEY, true, {}), + ); +}); diff --git a/tests/test_computer_fs_tools.py b/tests/test_computer_fs_tools.py index 7fe33fe21..a9f6fa16d 100644 --- a/tests/test_computer_fs_tools.py +++ b/tests/test_computer_fs_tools.py @@ -92,6 +92,95 @@ def _make_large_text() -> str: return "".join(f"line-{index:05d}-{'x' * 48}\n" for index in range(6000)) +def _make_epub_bytes(*, chapter_count: int = 1) -> bytes: + manifest_items = [ + '' + ] + spine_items = [''] + nav_links = [] + + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode="w") as archive: + archive.writestr( + "mimetype", + "application/epub+zip", + compress_type=zipfile.ZIP_STORED, + ) + archive.writestr( + "META-INF/container.xml", + """ + + + + + +""", + ) + + for index in range(1, chapter_count + 1): + manifest_items.append( + f'' + ) + spine_items.append(f'') + nav_links.append(f'
  • Chapter {index}
    • ') + archive.writestr( + f"OEBPS/chapter{index}.xhtml", + f""" + + + Chapter {index} + + +

    Chapter {index}

    +

    Paragraph {index}

    + + +""", + ) + + archive.writestr( + "OEBPS/nav.xhtml", + """ + + + Navigation + + + + + +""".format(links="".join(nav_links)), + ) + archive.writestr( + "OEBPS/content.opf", + """ + + + test-book + Test Book + en + + + {manifest} + + + {spine} + + +""".format( + manifest="".join(manifest_items), + spine="".join(spine_items), + ), + ) + + return buffer.getvalue() + + def test_detect_text_encoding_allows_utf8_probe_cut_mid_character(): sample = '{"results": ["中文内容"]}'.encode()[:-1] @@ -230,6 +319,36 @@ async def test_file_read_tool_reads_docx_via_parser_and_magic( assert result == "doc-line-1\ndoc-line-2\n" +def test_is_epub_bytes_rejects_plain_zip_archive(): + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode="w") as archive: + archive.writestr("README.txt", "hello") + + assert file_read_utils._is_epub_bytes(buffer.getvalue()) is False + + +@pytest.mark.asyncio +async def test_file_read_tool_reads_epub_via_parser_and_magic( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +): + workspace = _setup_local_fs_tools(monkeypatch, tmp_path) + epub_path = workspace / "novel.bin" + epub_path.write_bytes(_make_epub_bytes(chapter_count=2)) + + async def _fake_parse_epub(_file_bytes: bytes, _file_name: str) -> str: + return "# Chapter 1\n\nParagraph 1\n" + + monkeypatch.setattr(file_read_utils, "_parse_local_epub_text", _fake_parse_epub) + + result = await fs_tools.FileReadTool().call( + _make_context(), + path="novel.bin", + ) + + assert result == "# Chapter 1\n\nParagraph 1\n" + + @pytest.mark.asyncio async def test_file_read_tool_stores_long_converted_document_in_workspace( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_conversation_checkpoint.py b/tests/test_conversation_checkpoint.py new file mode 100644 index 000000000..96113b7d1 --- /dev/null +++ b/tests/test_conversation_checkpoint.py @@ -0,0 +1,110 @@ +import pytest + +from astrbot.core.agent.message import ( + CheckpointData, + CheckpointMessageSegment, + Message, + bind_checkpoint_messages, + dump_messages_with_checkpoints, + get_checkpoint_id, + strip_checkpoint_messages, +) +from astrbot.core.provider.provider import Provider +from astrbot.dashboard.routes.chat import ChatRoute + + +def test_checkpoint_message_segment_round_trip(): + message = CheckpointMessageSegment(content=CheckpointData(id="cp-1")) + + dumped = message.model_dump() + + assert dumped == {"role": "_checkpoint", "content": {"id": "cp-1"}} + assert get_checkpoint_id(dumped) == "cp-1" + assert Message.model_validate(dumped).content == CheckpointData(id="cp-1") + + +def test_checkpoint_requires_checkpoint_data(): + with pytest.raises(ValueError, match="checkpoint message content"): + Message(role="_checkpoint", content="cp-1") + + +def test_checkpoint_data_is_only_allowed_for_checkpoint_role(): + with pytest.raises(ValueError, match="CheckpointData is only allowed"): + Message(role="user", content=CheckpointData(id="cp-1")) + + +def test_strip_checkpoint_messages(): + history = [ + {"role": "user", "content": "hello"}, + {"role": "_checkpoint", "content": {"id": "cp-1"}}, + {"role": "assistant", "content": "world"}, + ] + + assert strip_checkpoint_messages(history) == [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + + +def test_bind_and_dump_checkpoint_messages_preserves_boundaries(): + history = [ + {"role": "user", "content": "old user"}, + {"role": "assistant", "content": "old bot"}, + {"role": "_checkpoint", "content": {"id": "cp-1"}}, + {"role": "user", "content": "next user"}, + ] + + messages = bind_checkpoint_messages(history) + + assert len(messages) == 3 + assert messages[1]._checkpoint_after == CheckpointData(id="cp-1") + assert dump_messages_with_checkpoints(messages) == [ + {"role": "user", "content": "old user"}, + {"role": "assistant", "content": "old bot"}, + {"role": "_checkpoint", "content": {"id": "cp-1"}}, + {"role": "user", "content": "next user"}, + ] + + +def test_dump_checkpoint_messages_drops_checkpoint_when_message_is_dropped(): + history = [ + {"role": "user", "content": "old user"}, + {"role": "assistant", "content": "old bot"}, + {"role": "_checkpoint", "content": {"id": "cp-1"}}, + {"role": "user", "content": "latest user"}, + ] + + messages = bind_checkpoint_messages(history) + + assert dump_messages_with_checkpoints(messages[2:]) == [ + {"role": "user", "content": "latest user"}, + ] + + +def test_provider_ensure_message_to_dicts_skips_checkpoints(): + messages = [ + Message(role="user", content="hello"), + CheckpointMessageSegment(content=CheckpointData(id="cp-1")), + {"role": "assistant", "content": "world"}, + {"role": "_checkpoint", "content": {"id": "cp-2"}}, + ] + + assert Provider._ensure_message_to_dicts(object(), messages) == [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + + +def test_chat_route_find_turn_range(): + route = ChatRoute.__new__(ChatRoute) + history = [ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + {"role": "_checkpoint", "content": {"id": "cp-1"}}, + {"role": "user", "content": "c"}, + {"role": "assistant", "content": "d"}, + {"role": "_checkpoint", "content": {"id": "cp-2"}}, + ] + + assert route._find_turn_range(history, "cp-2") == (3, 5) + assert route._find_turn_range(history, "missing") is None diff --git a/tests/test_epub_parser.py b/tests/test_epub_parser.py new file mode 100644 index 000000000..a53f89f36 --- /dev/null +++ b/tests/test_epub_parser.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import io +import zipfile + +import pytest + +from astrbot.core.knowledge_base.parsers.epub_parser import EpubParser +from astrbot.core.knowledge_base.parsers.util import select_parser + + +def _make_epub_bytes() -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode="w") as archive: + archive.writestr( + "mimetype", + "application/epub+zip", + compress_type=zipfile.ZIP_STORED, + ) + archive.writestr( + "META-INF/container.xml", + """ + + + + + +""", + ) + archive.writestr( + "OEBPS/nav.xhtml", + """ + + + Navigation + + + + + +""", + ) + archive.writestr( + "OEBPS/chapter2.xhtml", + """ + + + Chapter 2 + + +

    Second

    +

    Beta paragraph.

    + + +""", + ) + archive.writestr( + "OEBPS/chapter1.xhtml", + """ + + + Chapter 1 + + +

    First

    +

    Alpha paragraph.

    +
      +
    • Point A
      • +
    • Point B
      • +
    + + +""", + ) + archive.writestr( + "OEBPS/content.opf", + """ + + + test-book + Test Book + en + + + + + + + + + + + + +""", + ) + + return buffer.getvalue() + + +def _make_epub_bytes_with_generic_content() -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, mode="w") as archive: + archive.writestr( + "mimetype", + "application/epub+zip", + compress_type=zipfile.ZIP_STORED, + ) + archive.writestr( + "META-INF/container.xml", + """ + + + + + +""", + ) + archive.writestr( + "OEBPS/chapter1.xhtml", + """ + + + Chapter 1 + + +

    First

    + Lead text +

    Piura*5, continued.

    + +
    Inside div
    +
    Inside section
    + + + + + +
    Cell ACell B
    + + +""", + ) + archive.writestr( + "OEBPS/content.opf", + """ + + + test-book + Test Book + en + + + + + + + + +""", + ) + + return buffer.getvalue() + + +@pytest.mark.asyncio +async def test_select_parser_supports_epub(): + parser = await select_parser(".epub") + + assert isinstance(parser, EpubParser) + + +@pytest.mark.asyncio +async def test_epub_parser_reads_spine_order_as_text(): + result = await EpubParser().parse(_make_epub_bytes(), "book.epub") + + assert result.media == [] + assert "**Title:**" not in result.text + assert "[Chapter 2](chapter2.xhtml)" not in result.text + assert result.text.startswith("1. Chapter 2") + assert "2. Chapter 1" in result.text + assert "Beta paragraph." in result.text + assert "# First" in result.text + assert "* Point A" in result.text + assert result.text.index("1. Chapter 2") < result.text.index("## Second") + assert result.text.index("## Second") < result.text.index("# First") + + +@pytest.mark.asyncio +async def test_epub_parser_preserves_generic_container_text(): + result = await EpubParser().parse( + _make_epub_bytes_with_generic_content(), + "book.epub", + ) + + assert "**Title:**" not in result.text + assert "# First" in result.text + assert "Lead text" in result.text + assert r"Piura\*5, continued." in result.text + assert "filepos" not in result.text + assert r"[\*5]" not in result.text + assert "Image00000.jpg" not in result.text + assert "![](" not in result.text + assert "\n\n\n" not in result.text + assert "Inside div" in result.text + assert "Inside section" in result.text + assert "| Cell A | Cell B |" in result.text diff --git a/tests/test_kook/data/kook_api_response_user_me.json b/tests/test_kook/data/kook_api_response_user_me.json new file mode 100644 index 000000000..8fefa63a4 --- /dev/null +++ b/tests/test_kook/data/kook_api_response_user_me.json @@ -0,0 +1,36 @@ +{ + "code": 0, + "message": "操作成功", + "data": { + "id": "573092175", + "username": "bot_username", + "identify_num": "9561", + "online": false, + "os": "Websocket", + "status": 0, + "avatar": "https://example.com", + "vip_avatar": "https://example.com", + "banner": "", + "nickname": "bot_nickname", + "roles": [], + "is_vip": false, + "vip_amp": false, + "bot": true, + "kpm_vip": null, + "wealth_level": 0, + "bot_status": 0, + "tag_info": { + "color": "#0096FF", + "bg_color": "#0096FF33", + "text": "机器人" + }, + "mobile_verified": true, + "is_sys": false, + "client_id": "g3nsxNQhNMZFKatU", + "verified": false, + "mobile_prefix": "86", + "mobile": "****", + "invited_count": 0, + "intent": 255 + } +} \ No newline at end of file diff --git a/tests/test_kook/data/kook_api_response_user_view.json b/tests/test_kook/data/kook_api_response_user_view.json new file mode 100644 index 000000000..3dfcdb08e --- /dev/null +++ b/tests/test_kook/data/kook_api_response_user_view.json @@ -0,0 +1,36 @@ +{ + "code": 0, + "message": "操作成功", + "data": { + "id": "573092175", + "username": "bot_username", + "identify_num": "9561", + "online": false, + "os": "Websocket", + "status": 0, + "avatar": "https://img.kookapp.cn/assets/bot.png?x-oss-process=style/icon", + "vip_avatar": "https://img.kookapp.cn/assets/bot.png?x-oss-process=style/icon", + "banner": "", + "nickname": "bot_nickname", + "roles": [ + 13726212 + ], + "is_vip": false, + "vip_amp": false, + "bot": true, + "kpm_vip": null, + "wealth_level": 0, + "bot_status": 0, + "tag_info": { + "color": "#0096FF", + "bg_color": "#0096FF33", + "text": "机器人" + }, + "mobile_verified": true, + "is_sys": false, + "client_id": "g3nsxNQhNMZFKatU", + "verified": false, + "joined_at": 1772260532000, + "active_time": 1776418003694 + } +} \ No newline at end of file diff --git a/tests/test_kook/data/kook_ws_event_group_message.json b/tests/test_kook/data/kook_ws_event_group_message.json index dcab6e901..53bd50481 100644 --- a/tests/test_kook/data/kook_ws_event_group_message.json +++ b/tests/test_kook/data/kook_ws_event_group_message.json @@ -26,7 +26,7 @@ "banner": "", "nickname": "some_username", "roles": [ - 63724577 + 63423577 ], "is_vip": false, "vip_amp": false, diff --git a/tests/test_kook/data/kook_ws_event_group_message_with_mention.json b/tests/test_kook/data/kook_ws_event_group_message_with_mention.json new file mode 100644 index 000000000..50f598a9e --- /dev/null +++ b/tests/test_kook/data/kook_ws_event_group_message_with_mention.json @@ -0,0 +1,88 @@ +{ + "s": 0, + "d": { + "channel_type": "GROUP", + "type": 9, + "target_id": "2732467349811313213", + "author_id": "7324688132731983", + "content": "(rol)25555643(rol) /help (met)3351526782(met) (met)all(met) ", + "msg_id": "9b047d81-40fe-41af-ad39-916ae77e6b20", + "msg_timestamp": 1776405840600, + "nonce": "r2oQkO7kRpNSgmsv7TNl2zOA", + "from_type": 1, + "extra": { + "type": 9, + "code": "", + "author": { + "id": "3351526782", + "username": "some_username", + "identify_num": "4198", + "nickname": "some_username", + "bot": false, + "online": true, + "avatar": "https://example.com", + "vip_avatar": "https://example.com", + "status": 1, + "roles": [ + 63423577 + ], + "os": "Websocket", + "banner": "", + "is_vip": false, + "vip_amp": false, + "nameplate": [], + "wealth_level": 0, + "is_sys": false + }, + "kmarkdown": { + "raw_content": "@some_role /help @some_username @全体成员", + "mention_part": [ + { + "id": "3351526782", + "username": "some_username", + "full_name": "some_username#4198", + "avatar": "https://example.com", + "wealth_level": 0 + } + ], + "mention_role_part": [ + { + "role_id": 25555643, + "name": "some_role", + "color": 0, + "color_type": 1, + "color_map": [] + } + ], + "channel_part": [], + "spl": [] + }, + "last_msg_content": "some_username：@some_role /help @some_username @全体成员", + "mention": [ + "3351526782" + ], + "mention_all": true, + "mention_here": false, + "guild_id": "1239678456780469", + "guild_type": 0, + "channel_name": "聊天大厅", + "visible_only": "", + "mention_no_at": [], + "mention_roles": [ + 25555643 + ], + "nav_channels": [], + "emoji": [], + "preview_content": "", + "channel_type": 1, + "send_msg_device": 0 + } + }, + "extra": { + "verifyToken": "kW4FH_ASHio1hosd", + "encryptKey": "", + "callbackUrl": "", + "intent": 255 + }, + "sn": 1 +} \ No newline at end of file diff --git a/tests/test_kook/data/kook_ws_event_group_system_message_update_role.json b/tests/test_kook/data/kook_ws_event_group_system_message_update_role.json new file mode 100644 index 000000000..ae8ea6c17 --- /dev/null +++ b/tests/test_kook/data/kook_ws_event_group_system_message_update_role.json @@ -0,0 +1,38 @@ +{ + "s": 0, + "d": { + "channel_type": "GROUP", + "type": 255, + "target_id": "6059671921720469", + "author_id": "1", + "content": "[系统消息]", + "msg_id": "35de353d-aa69-40a8-b329-e4498ed8d30d", + "msg_timestamp": 1776498088031, + "nonce": "", + "from_type": 1, + "extra": { + "type": "updated_role", + "body": { + "role_id": 16752184, + "name": "some_username", + "color": 0, + "position": 5, + "hoist": 0, + "mentionable": 0, + "permissions": 5571747823, + "desc": "", + "color_type": 1, + "color_map": {}, + "type": 1, + "op_permissions": 0 + } + } + }, + "sn": 4, + "extra": { + "verifyToken": "kW4FH_ASHio1hosd", + "encryptKey": "", + "callbackUrl": "", + "intent": 255 + } +} \ No newline at end of file diff --git a/tests/test_kook/shared.py b/tests/test_kook/shared.py index f5ef18b8b..366b608db 100644 --- a/tests/test_kook/shared.py +++ b/tests/test_kook/shared.py @@ -1,5 +1,111 @@ -from pathlib import Path +import json +from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock +import aiohttp +from astrbot.api.platform import AstrBotMessage, MessageType +from astrbot.core.message.components import ( + File, + Record, +) + + +from pathlib import Path + CURRENT_DIR = Path(__file__).parent TEST_DATA_DIR = CURRENT_DIR / "data" + + +class KookEventDataPath: + GROUP_MESSAGE_WITH_MENTION = ( + TEST_DATA_DIR / "kook_ws_event_group_message_with_mention.json" + ) + GROUP_MESSAGE = TEST_DATA_DIR / "kook_ws_event_group_message.json" + HELLO = TEST_DATA_DIR / "kook_ws_event_hello.json" + MESSAGE_WITH_CARD_1 = TEST_DATA_DIR / "kook_ws_event_message_with_card_1.json" + MESSAGE_WITH_CARD_2 = TEST_DATA_DIR / "kook_ws_event_message_with_card_2.json" + PING = TEST_DATA_DIR / "kook_ws_event_ping.json" + PONG = TEST_DATA_DIR / "kook_ws_event_pong.json" + PRIVATE_MESSAGE = TEST_DATA_DIR / "kook_ws_event_private_message.json" + PRIVATE_SYSTEM_MESSAGE = TEST_DATA_DIR / "kook_ws_event_private_system_message.json" + RECONNECT_ERR = TEST_DATA_DIR / "kook_ws_event_reconnect_err.json" + RESUME_ACK = TEST_DATA_DIR / "kook_ws_event_resume_ack.json" + RESUME = TEST_DATA_DIR / "kook_ws_event_resume.json" + GROUP_SYSTEM_MESSAGE_UPDATE_ROLE = TEST_DATA_DIR / "kook_ws_event_group_system_message_update_role.json" + + +class KookApiDataPath: + USER_ME = TEST_DATA_DIR / "kook_api_response_user_me.json" + USER_VIEW = TEST_DATA_DIR / "kook_api_response_user_view.json" + + +def mock_kook_client(upload_asset_return: str, send_text_return: str): + client = MagicMock() + + client.upload_asset = AsyncMock(return_value=upload_asset_return) + client.send_text = AsyncMock(return_value=send_text_return) + return client + + +def mock_http_client( + http_method: str = "get", + return_value: str | dict | list | None = None, + status: int = 200, +): + """Mock aiohttp ClientSession""" + + if isinstance(return_value, (dict, list)): + response_text = json.dumps(return_value) + else: + response_text = return_value or "{}" + + mock_response = MagicMock() + mock_response.status = status + mock_response.text = AsyncMock(return_value=response_text) + mock_response.json = AsyncMock( + return_value=json.loads(response_text) if response_text else {} + ) + mock_response.read = AsyncMock(return_value=response_text.encode()) + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.__aexit__ = AsyncMock(return_value=None) + mock_session = MagicMock() + + async def mock_method(*args, **kwargs): + return mock_response + + setattr(mock_session, http_method.lower(), mock_method) + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + + mock_session.close = AsyncMock() + + return mock_session + + +def mock_file_message(input: str): + message = MagicMock(spec=File) + message.get_file = AsyncMock(return_value=input) + return message + + +def mock_record_message(input: str): + message = MagicMock(spec=Record) + message.text = input + message.convert_to_file_path = AsyncMock(return_value=input) + return message + + +def mock_astrbot_message(): + message = AstrBotMessage() + message.type = MessageType.OTHER_MESSAGE + message.group_id = "test" + message.session_id = "test" + message.message_id = "test" + return message + + +def mock_kook_roles_record(bot_id: str, http_client: aiohttp.ClientSession): + instance = AsyncMock() + instance.has_role_in_channel = AsyncMock(return_value=True) + return instance diff --git a/tests/test_kook/test_kook_client.py b/tests/test_kook/test_kook_client.py new file mode 100644 index 000000000..22556386b --- /dev/null +++ b/tests/test_kook/test_kook_client.py @@ -0,0 +1,156 @@ +import asyncio +from dataclasses import dataclass, field +import json +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from astrbot.core.message.components import ( + At, + AtAll, + BaseMessageComponent, + Plain, + Record, +) +from astrbot.core.platform.sources.kook.kook_client import KookClient +from astrbot.core.platform.sources.kook.kook_config import KookConfig +from astrbot.core.platform.sources.kook.kook_types import ( + KookMessageEventData, + KookWebsocketEvent, +) +from tests.test_kook.shared import ( + KookEventDataPath, + mock_http_client, + mock_kook_roles_record, +) + +TEST_BOT_ID = 1234567891 +TEST_BOT_USERNAME = "test_username" +TEST_BOT_NICKNAME = "test_nickname" + + +def mock_kook_client(config: KookConfig, event_callback): + class MockKookClient: + def __init__(self, config, callback): + self.bot_id = TEST_BOT_ID + self.bot_nickname = TEST_BOT_NICKNAME + self.bot_username = TEST_BOT_USERNAME + self.http_client = mock_http_client() + self.connect = AsyncMock() + self.close = AsyncMock() + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + return MockKookClient(config, event_callback) + + +def get_json_field(content: dict, json_field_path: list[str | int]) -> Any: + expend_value = content + for key in json_field_path: + expend_value = expend_value[key] + return expend_value + + +@dataclass +class JsonFieldPaths: + message_str: list[int | str] = field(default_factory=list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "expected_json_data_path, expected_message_str, expected_message_components", + [ + ( + KookEventDataPath.GROUP_MESSAGE_WITH_MENTION, + ["d", "extra", "kmarkdown", "raw_content"], + [ + # 这里默认机器人一定属于某个角色id + At(qq=TEST_BOT_ID, name="some_role"), + Plain(text="/help"), + At(qq=3351526782, name="some_username"), + AtAll(qq="all", name=""), + ], + ), + ( + KookEventDataPath.GROUP_MESSAGE, + ["d", "extra", "kmarkdown", "raw_content"], + [Plain(text="done!")], + ), + ( + KookEventDataPath.MESSAGE_WITH_CARD_1, + "[audio]", + [ + Plain(text="[audio]"), + Record( + file="https://img.kookapp.cn/attachments/2026-03/03/69a6841c3125d.wav", + url="", + text=None, + path=None, + ), + ], + ), + ( + KookEventDataPath.MESSAGE_WITH_CARD_2, + ["d", "extra", "kmarkdown", "raw_content"], + [ + Plain(text="(met)"), + Plain(text="all(met) #hello \\*\\*world\\*\\* [audio]\n😆"), + Record( + file="https://img.kookapp.cn/attachments/2026-03/03/69a6841c3125d.wav", + url="", + text=None, + path=None, + ), + ], + ), + ( + KookEventDataPath.PRIVATE_MESSAGE, + ["d", "extra", "kmarkdown", "raw_content"], + [Plain(text="/help")], + ), + ], +) +async def test_kook_event_warp_message( + expected_json_data_path: Path, + expected_message_str: list[int | str] | str, + expected_message_components: list[BaseMessageComponent], +): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + "astrbot.core.platform.sources.kook.kook_adapter.KookClient", mock_kook_client + ) + monkeypatch.setattr( + "astrbot.core.platform.sources.kook.kook_adapter.KookRolesRecord", + mock_kook_roles_record, + ) + + from astrbot.core.platform.sources.kook.kook_adapter import KookPlatformAdapter + + adapter = KookPlatformAdapter({}, {}, asyncio.Queue()) + + raw_event_str = expected_json_data_path.read_text(encoding="utf-8") + raw_event = json.loads(raw_event_str) + event = KookWebsocketEvent.from_json( + raw_event_str, + ) + assert isinstance(event.data, KookMessageEventData) + + astrbotMessage = await adapter.convert_message(event.data) + assert astrbotMessage.self_id == TEST_BOT_ID + assert astrbotMessage.sender.user_id == raw_event["d"]["author_id"] + assert ( + astrbotMessage.sender.nickname == raw_event["d"]["extra"]["author"]["username"] + ) + assert astrbotMessage.raw_message == raw_event["d"] + assert astrbotMessage.message_id == raw_event["d"]["msg_id"] + assert astrbotMessage.message == expected_message_components + if isinstance(expected_message_str, str): + assert astrbotMessage.message_str == expected_message_str + else: + assert get_json_field(raw_event, expected_message_str) diff --git a/tests/test_kook/test_kook_event.py b/tests/test_kook/test_kook_event.py index 5fe73a510..616d596c8 100644 --- a/tests/test_kook/test_kook_event.py +++ b/tests/test_kook/test_kook_event.py @@ -1,10 +1,8 @@ -from unittest.mock import AsyncMock, MagicMock +import json import pytest -from astrbot.api.platform import AstrBotMessage, MessageType, PlatformMetadata, Unknown -from astrbot.api.event import MessageChain +from astrbot.api.platform import PlatformMetadata, Unknown from astrbot.core.message.components import ( - File, Image, Plain, Video, @@ -12,44 +10,18 @@ from astrbot.core.message.components import ( AtAll, BaseMessageComponent, Json, - Record, Reply, ) from astrbot.core.platform.sources.kook.kook_event import KookEvent from astrbot.core.platform.sources.kook.kook_types import KookMessageType, OrderMessage - - -async def mock_kook_client(upload_asset_return: str, send_text_return: str): - # 1. Mock 掉整个 KookClient 类 - client = MagicMock() - - client.upload_asset = AsyncMock(return_value=upload_asset_return) - client.send_text = AsyncMock(return_value=send_text_return) - return client - - -def mock_file_message(input: str): - message = MagicMock(spec=File) - message.get_file = AsyncMock(return_value=input) - return message - - -def mock_record_message(input: str): - message = MagicMock(spec=Record) - message.text = input - message.convert_to_file_path = AsyncMock(return_value=input) - return message - - -def mock_astrbot_message(): - message = AstrBotMessage() - message.type = MessageType.OTHER_MESSAGE - message.group_id = "test" - message.session_id = "test" - message.message_id = "test" - return message +from tests.test_kook.shared import ( + mock_astrbot_message, + mock_file_message, + mock_kook_client, + mock_record_message, +) @pytest.mark.asyncio @@ -161,7 +133,7 @@ async def test_kook_event_warp_message( expected_output: OrderMessage, expected_error: type[BaseException] | None, ): - client = await mock_kook_client( + client = mock_kook_client( upload_asset_return, "", ) @@ -184,5 +156,20 @@ async def test_kook_event_warp_message( return result = await event._wrap_message(1, input_message) - assert result == expected_output - \ No newline at end of file + + expected_output_text: str | list | dict = expected_output.text + is_json_text = False + try: + expected_output_text = json.loads(expected_output_text) + is_json_text = True + except: + pass + + if is_json_text: + assert json.loads(result.text) == expected_output_text + else: + assert result.text == expected_output_text + + assert result.index == expected_output.index + assert result.type == expected_output.type + assert result.reply_id == expected_output.reply_id diff --git a/tests/test_kook/test_kook_types.py b/tests/test_kook/test_kook_types.py index 85c39622c..89568168a 100644 --- a/tests/test_kook/test_kook_types.py +++ b/tests/test_kook/test_kook_types.py @@ -15,16 +15,19 @@ from astrbot.core.platform.sources.kook.kook_types import ( ImageGroupModule, InviteModule, KmarkdownElement, + KookApiResponseBase, KookCardMessage, KookMessageSignal, KookModuleType, + KookUserMeResponse, + KookUserViewResponse, KookWebsocketEvent, ParagraphStructure, PlainTextElement, SectionModule, KookCardMessageContainer, ) -from tests.test_kook.shared import TEST_DATA_DIR +from tests.test_kook.shared import TEST_DATA_DIR, KookApiDataPath, KookEventDataPath def test_kook_card_message_container_append(): @@ -109,29 +112,31 @@ def test_all_kook_card_type(): ).to_json(indent=4, ensure_ascii=False) assert json_output == expect_json_data + @pytest.mark.parametrize( - "expected_json_data_filename", + "expected_json_data_path", [ - ("kook_ws_event_group_message.json"), - ("kook_ws_event_hello.json"), - ("kook_ws_event_message_with_card_1.json"), - ("kook_ws_event_message_with_card_2.json"), - ("kook_ws_event_ping.json"), - ("kook_ws_event_pong.json"), - ("kook_ws_event_private_message.json"), - ("kook_ws_event_private_system_message.json"), - ("kook_ws_event_reconnect_err.json"), - ("kook_ws_event_resume_ack.json"), - ("kook_ws_event_resume.json"), - + (KookEventDataPath.GROUP_MESSAGE_WITH_MENTION), + (KookEventDataPath.GROUP_MESSAGE), + (KookEventDataPath.HELLO), + (KookEventDataPath.MESSAGE_WITH_CARD_1), + (KookEventDataPath.MESSAGE_WITH_CARD_2), + (KookEventDataPath.PING), + (KookEventDataPath.PONG), + (KookEventDataPath.PRIVATE_MESSAGE), + (KookEventDataPath.PRIVATE_SYSTEM_MESSAGE), + (KookEventDataPath.RECONNECT_ERR), + (KookEventDataPath.RESUME_ACK), + (KookEventDataPath.RESUME), + (KookEventDataPath.GROUP_SYSTEM_MESSAGE_UPDATE_ROLE), ], ) -def test_websocket_event_type_parse(expected_json_data_filename:str): - expected_json_data_str =(TEST_DATA_DIR / expected_json_data_filename).read_text(encoding="utf-8") +def test_websocket_event_type_parse(expected_json_data_path: Path): + expected_json_data_str = (expected_json_data_path).read_text(encoding="utf-8") event = KookWebsocketEvent.from_json( expected_json_data_str, ) - event_dict = event.to_dict(mode="json",exclude_unset=True,exclude_none=False) + event_dict = event.to_dict() assert event_dict == json.loads(expected_json_data_str) @@ -141,8 +146,27 @@ def test_websocket_event_create(): data=None, sn=0, ) - assert ping_data.to_dict(mode="json")== { + assert ping_data.to_dict(exclude_none=True, exclude_unset=False) == { "s": KookMessageSignal.PING.value, "sn": 0, } - \ No newline at end of file + + +@pytest.mark.parametrize( + "expected_json_data_path, expected_dataclass", + [ + (KookApiDataPath.USER_ME, KookUserMeResponse), + (KookApiDataPath.USER_VIEW, KookUserViewResponse), + ], +) +def test_api_response_type_parse( + expected_json_data_path: Path, expected_dataclass: type[KookApiResponseBase] +): + expected_json_data_str = (expected_json_data_path).read_text(encoding="utf-8") + + response_body = expected_dataclass.from_json( + expected_json_data_str, + ) + + body_dict = response_body.to_dict() + assert body_dict == json.loads(expected_json_data_str) diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index a66f1199d..96152d27e 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -14,6 +14,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..") from astrbot.core.agent.agent import Agent from astrbot.core.agent.handoff import HandoffTool from astrbot.core.agent.hooks import BaseAgentRunHooks +from astrbot.core.agent.message import ImageURLPart, Message, TextPart from astrbot.core.agent.run_context import ContextWrapper from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner from astrbot.core.agent.tool import FunctionTool, ToolSet @@ -150,6 +151,25 @@ class MockErrProvider(MockProvider): ) +class CapturingProvider(MockProvider): + def __init__(self, modalities: list[str]): + super().__init__() + self.provider_config["modalities"] = modalities + self.received_contexts = [] + self.received_func_tools = [] + self.should_call_tools = False + + async def text_chat(self, **kwargs) -> LLMResponse: + self.call_count += 1 + self.received_contexts.append(kwargs.get("contexts")) + self.received_func_tools.append(kwargs.get("func_tool")) + return LLMResponse( + role="assistant", + completion_text="final", + usage=TokenUsage(input_other=10, output=5), + ) + + class MockEmptyOutputThenSuccessProvider(MockProvider): def __init__(self, failures_before_success: int = 1): super().__init__() @@ -609,6 +629,99 @@ async def test_tool_result_includes_all_calltoolresult_content( ] +@pytest.mark.asyncio +async def test_runner_replaces_runtime_image_context_before_provider_call( + runner, provider_request, mock_hooks +): + provider = CapturingProvider(modalities=["tool_use"]) + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=MockToolExecutor, + agent_hooks=mock_hooks, + streaming=False, + ) + + runner.run_context.messages.append( + Message( + role="user", + content=[ + TextPart(text="Review this image"), + ImageURLPart( + image_url=ImageURLPart.ImageURL( + url="data:image/png;base64,dGVzdA==" + ) + ), + ], + ) + ) + + async for _ in runner.step_until_done(1): + pass + + assert provider.received_contexts + sent_context = provider.received_contexts[0] + assert sent_context[-1]["content"] == [ + {"type": "text", "text": "Review this image"}, + {"type": "text", "text": "[Image]"}, + ] + assert len(runner.run_context.messages[-2].content) == 2 + + +@pytest.mark.asyncio +async def test_runner_builds_placeholder_for_unsupported_request_image( + runner, mock_hooks, tool_set +): + provider = CapturingProvider(modalities=["tool_use"]) + request = ProviderRequest( + prompt="Describe it", + image_urls=["/path/that/should/not/be/read.jpg"], + func_tool=tool_set, + contexts=[], + ) + + await runner.reset( + provider=provider, + request=request, + run_context=ContextWrapper(context=None), + tool_executor=MockToolExecutor, + agent_hooks=mock_hooks, + streaming=False, + ) + + async for _ in runner.step_until_done(1): + pass + + sent_context = provider.received_contexts[0] + assert sent_context[-1]["content"] == [ + {"type": "text", "text": "Describe it"}, + {"type": "text", "text": "[Image]"}, + ] + + +@pytest.mark.asyncio +async def test_runner_clears_tools_for_provider_without_tool_use( + runner, provider_request, mock_hooks, mock_tool_executor +): + provider = CapturingProvider(modalities=["text"]) + + await runner.reset( + provider=provider, + request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, + streaming=False, + ) + + async for _ in runner.step_until_done(1): + pass + + assert provider.received_func_tools == [None] + + @pytest.mark.asyncio async def test_same_tool_consecutive_results_include_escalating_guidance( runner, mock_tool_executor, mock_hooks diff --git a/tests/unit/test_astr_main_agent.py b/tests/unit/test_astr_main_agent.py index 4d51d3d68..dd0b2dbe5 100644 --- a/tests/unit/test_astr_main_agent.py +++ b/tests/unit/test_astr_main_agent.py @@ -8,7 +8,7 @@ from astrbot.core import astr_main_agent as ama from astrbot.core.agent.mcp_client import MCPTool from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.conversation_mgr import Conversation -from astrbot.core.message.components import File, Image, Plain, Reply +from astrbot.core.message.components import File, Image, Plain, Reply, Video from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.platform.platform_metadata import PlatformMetadata from astrbot.core.provider import Provider @@ -663,144 +663,6 @@ class TestDecorateLlmRequest: assert req.prompt == "Hello" -class TestModalitiesFix: - """Tests for _modalities_fix function.""" - - def test_modalities_fix_image_not_supported(self, mock_provider): - """Test modality fix when image is not supported.""" - module = ama - mock_provider.provider_config = {"modalities": ["text"]} - req = ProviderRequest(prompt="Hello", image_urls=["/path/to/image.jpg"]) - - module._modalities_fix(mock_provider, req) - - assert "[Image]" in req.prompt - assert req.image_urls == [] - - def test_modalities_fix_tool_not_supported(self, mock_provider): - """Test modality fix when tool is not supported.""" - module = ama - mock_provider.provider_config = {"modalities": ["text", "image"]} - req = ProviderRequest(prompt="Hello") - req.func_tool = ToolSet() - req.func_tool.add_tool( - FunctionTool( - name="dummy_tool", - description="dummy", - parameters={"type": "object", "properties": {}}, - ) - ) - - module._modalities_fix(mock_provider, req) - - assert req.func_tool is None - - def test_modalities_fix_all_supported(self, mock_provider): - """Test modality fix when all features are supported.""" - module = ama - mock_provider.provider_config = {"modalities": ["image", "tool_use"]} - tool_set = ToolSet() - tool_set.add_tool( - FunctionTool( - name="dummy_tool", - description="dummy", - parameters={"type": "object", "properties": {}}, - ) - ) - req = ProviderRequest( - prompt="Hello", - image_urls=["/path/to/image.jpg"], - func_tool=tool_set, - ) - - module._modalities_fix(mock_provider, req) - - assert req.prompt == "Hello" - assert len(req.image_urls) == 1 - assert req.func_tool is not None - - -class TestSanitizeContextByModalities: - """Tests for _sanitize_context_by_modalities function.""" - - def test_sanitize_no_op(self, mock_provider): - """Test sanitize when disabled or modalities support everything.""" - module = ama - config = module.MainAgentBuildConfig( - tool_call_timeout=60, sanitize_context_by_modalities=False - ) - mock_provider.provider_config = {"modalities": ["image", "tool_use"]} - req = ProviderRequest(contexts=[{"role": "user", "content": "Hello"}]) - - module._sanitize_context_by_modalities(config, mock_provider, req) - - assert len(req.contexts) == 1 - - def test_sanitize_removes_tool_messages(self, mock_provider): - """Test sanitize removes tool messages when tool_use not supported.""" - module = ama - config = module.MainAgentBuildConfig( - tool_call_timeout=60, sanitize_context_by_modalities=True - ) - mock_provider.provider_config = {"modalities": ["image"]} - req = ProviderRequest( - contexts=[ - {"role": "user", "content": "Hello"}, - {"role": "tool", "content": "Tool result"}, - ] - ) - - module._sanitize_context_by_modalities(config, mock_provider, req) - - assert len(req.contexts) == 1 - assert req.contexts[0]["role"] == "user" - - def test_sanitize_removes_tool_calls(self, mock_provider): - """Test sanitize removes tool_calls from assistant messages.""" - module = ama - config = module.MainAgentBuildConfig( - tool_call_timeout=60, sanitize_context_by_modalities=True - ) - mock_provider.provider_config = {"modalities": ["image"]} - req = ProviderRequest( - contexts=[ - { - "role": "assistant", - "content": "Response", - "tool_calls": [{"name": "tool"}], - } - ] - ) - - module._sanitize_context_by_modalities(config, mock_provider, req) - - assert "tool_calls" not in req.contexts[0] - - def test_sanitize_removes_image_blocks(self, mock_provider): - """Test sanitize removes image blocks when image not supported.""" - module = ama - config = module.MainAgentBuildConfig( - tool_call_timeout=60, sanitize_context_by_modalities=True - ) - mock_provider.provider_config = {"modalities": ["tool_use"]} - req = ProviderRequest( - contexts=[ - { - "role": "user", - "content": [ - {"type": "text", "text": "Hello"}, - {"type": "image_url", "url": "image.jpg"}, - ], - } - ] - ) - - module._sanitize_context_by_modalities(config, mock_provider, req) - - assert len(req.contexts[0]["content"]) == 1 - assert req.contexts[0]["content"][0]["type"] == "text" - - class TestPluginToolFix: """Tests for _plugin_tool_fix function.""" @@ -1018,6 +880,146 @@ class TestBuildMainAgent: assert result is not None + @pytest.mark.asyncio + async def test_build_main_agent_with_video_attachment( + self, mock_event, mock_context, mock_provider + ): + """Test building main agent with video attachments.""" + module = ama + mock_video = Video(file="file:///path/to/video.mp4") + mock_event.message_obj.message = [mock_video] + + mock_context.get_provider_by_id.return_value = None + mock_context.get_using_provider.return_value = mock_provider + mock_context.get_config.return_value = {} + + conv_mgr = mock_context.conversation_manager + _setup_conversation_for_build(conv_mgr) + + with ( + patch("astrbot.core.astr_main_agent.AgentRunner") as mock_runner_cls, + patch("astrbot.core.astr_main_agent.AstrAgentContext"), + ): + mock_runner = MagicMock() + mock_runner.reset = AsyncMock() + mock_runner_cls.return_value = mock_runner + + result = await module.build_main_agent( + event=mock_event, + plugin_context=mock_context, + config=module.MainAgentBuildConfig(tool_call_timeout=60), + ) + + assert result is not None + assert [ + part.text for part in result.provider_request.extra_user_content_parts + ] == ["[Video Attachment: name video.mp4, path path/to/video.mp4]"] + + @pytest.mark.asyncio + async def test_build_main_agent_with_quoted_video_attachment( + self, mock_event, mock_context, mock_provider + ): + """Test building main agent with quoted video attachments.""" + module = ama + mock_video = Video(file="file:///path/to/quoted-video.mp4") + mock_reply = Reply( + id="reply-1", + chain=[mock_video], + sender_nickname="", + message_str="quoted message", + ) + mock_event.message_obj.message = [Plain(text="Hello"), mock_reply] + + mock_context.get_provider_by_id.return_value = None + mock_context.get_using_provider.return_value = mock_provider + mock_context.get_config.return_value = {} + + conv_mgr = mock_context.conversation_manager + _setup_conversation_for_build(conv_mgr) + + with ( + patch("astrbot.core.astr_main_agent.AgentRunner") as mock_runner_cls, + patch("astrbot.core.astr_main_agent.AstrAgentContext"), + ): + mock_runner = MagicMock() + mock_runner.reset = AsyncMock() + mock_runner_cls.return_value = mock_runner + + result = await module.build_main_agent( + event=mock_event, + plugin_context=mock_context, + config=module.MainAgentBuildConfig(tool_call_timeout=60), + ) + + assert result is not None + assert ( + "[Video Attachment in quoted message: " + "name quoted-video.mp4, path path/to/quoted-video.mp4]" + ) in [part.text for part in result.provider_request.extra_user_content_parts] + + @pytest.mark.asyncio + async def test_build_main_agent_skips_video_attachment_when_conversion_fails( + self, mock_event, mock_context, mock_provider + ): + """Test video attachment failures do not abort request construction.""" + module = ama + mock_video = Video(file="file:///path/to/direct.mp4") + mock_quoted_video = Video(file="file:///path/to/quoted.mp4") + mock_reply = Reply( + id="reply-1", + chain=[mock_quoted_video], + sender_nickname="", + message_str="quoted message", + ) + mock_event.message_obj.message = [mock_video, mock_reply] + + mock_context.get_provider_by_id.return_value = None + mock_context.get_using_provider.return_value = mock_provider + mock_context.get_config.return_value = {} + + conv_mgr = mock_context.conversation_manager + _setup_conversation_for_build(conv_mgr) + + async def _raise_video_conversion_error(self): + if self.file.endswith("direct.mp4"): + raise RuntimeError("direct") + raise RuntimeError("quoted") + + with ( + patch("astrbot.core.astr_main_agent.AgentRunner") as mock_runner_cls, + patch("astrbot.core.astr_main_agent.AstrAgentContext"), + patch("astrbot.core.astr_main_agent.logger") as mock_logger, + patch.object( + Video, + "convert_to_file_path", + AsyncMock(side_effect=_raise_video_conversion_error), + ), + ): + mock_runner = MagicMock() + mock_runner.reset = AsyncMock() + mock_runner_cls.return_value = mock_runner + + result = await module.build_main_agent( + event=mock_event, + plugin_context=mock_context, + config=module.MainAgentBuildConfig(tool_call_timeout=60), + ) + + assert result is not None + assert not any( + "Video Attachment" in part.text + for part in result.provider_request.extra_user_content_parts + ) + assert mock_logger.error.call_count == 2 + assert ( + "Error processing video attachment" + in mock_logger.error.call_args_list[0][0][0] + ) + assert ( + "Error processing quoted video attachment" + in mock_logger.error.call_args_list[1][0][0] + ) + @pytest.mark.asyncio async def test_build_main_agent_no_prompt_no_images( self, mock_event, mock_context, mock_provider diff --git a/tests/unit/test_document_storage_fts.py b/tests/unit/test_document_storage_fts.py index 753c371ef..a7dd32c94 100644 --- a/tests/unit/test_document_storage_fts.py +++ b/tests/unit/test_document_storage_fts.py @@ -1,3 +1,5 @@ +import sqlite3 + import pytest from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage @@ -73,3 +75,29 @@ async def test_document_storage_fts_delete_skips_missing_fts_row(tmp_path): assert await storage.get_document_by_doc_id("legacy-chunk") is None await storage.close() + + +@pytest.mark.asyncio +async def test_document_storage_fts_recovers_from_legacy_non_fts_table(tmp_path): + db_path = tmp_path / "doc.db" + conn = sqlite3.connect(db_path) + conn.execute("CREATE TABLE documents_fts (rowid INTEGER PRIMARY KEY)") + conn.commit() + conn.close() + + storage = DocumentStorage(str(db_path)) + await storage.initialize() + + assert storage.fts5_available is True + + await storage.insert_document( + doc_id="legacy-fix", + text="legacy fts recovery text", + metadata={"kb_doc_id": "doc-1", "kb_id": "kb-1", "chunk_index": 0}, + ) + results = await storage.search_sparse(["legacy"], limit=10) + + assert results is not None + assert [result["doc_id"] for result in results] == ["legacy-fix"] + + await storage.close() diff --git a/tests/unit/test_mcp_client_schema.py b/tests/unit/test_mcp_client_schema.py new file mode 100644 index 000000000..0c3d9bc6a --- /dev/null +++ b/tests/unit/test_mcp_client_schema.py @@ -0,0 +1,117 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from astrbot.core.agent.mcp_client import MCPTool, _normalize_mcp_input_schema + + +class TestNormalizeMcpInputSchema: + def test_lifts_property_level_required_booleans_to_parent_required_array(self): + schema = { + "type": "object", + "properties": { + "stock_code": {"type": "string", "required": True}, + "market": {"type": "string", "required": False}, + }, + } + + normalized = _normalize_mcp_input_schema(schema) + + assert normalized["required"] == ["stock_code"] + assert "required" not in normalized["properties"]["stock_code"] + assert "required" not in normalized["properties"]["market"] + assert schema["properties"]["stock_code"]["required"] is True + + def test_preserves_existing_required_arrays_while_fixing_nested_objects(self): + schema = { + "type": "object", + "required": ["server"], + "properties": { + "server": { + "type": "object", + "required": ["transport"], + "properties": { + "transport": {"type": "string"}, + "stock_code": {"type": "string", "required": True}, + "market": {"type": "string", "required": False}, + }, + } + }, + } + + normalized = _normalize_mcp_input_schema(schema) + + assert normalized["required"] == ["server"] + assert normalized["properties"]["server"]["required"] == [ + "transport", + "stock_code", + ] + assert ( + "required" + not in normalized["properties"]["server"]["properties"]["stock_code"] + ) + assert ( + "required" not in normalized["properties"]["server"]["properties"]["market"] + ) + + def test_preserves_parent_required_flag_for_nested_object_properties(self): + schema = { + "type": "object", + "properties": { + "server": { + "type": "object", + "required": True, + "properties": { + "transport": {"type": "string", "required": True}, + }, + } + }, + } + + normalized = _normalize_mcp_input_schema(schema) + + assert normalized["required"] == ["server"] + assert normalized["properties"]["server"]["required"] == ["transport"] + assert ( + "required" + not in normalized["properties"]["server"]["properties"]["transport"] + ) + + def test_ignores_non_boolean_required_values_and_non_dict_properties(self): + schema = { + "type": "object", + "properties": { + "server": "invalid-property-schema", + "market": {"type": "string", "required": "yes"}, + "stock_code": {"type": "string", "required": True}, + }, + } + + normalized = _normalize_mcp_input_schema(schema) + + assert normalized["required"] == ["stock_code"] + assert normalized["properties"]["server"] == "invalid-property-schema" + assert normalized["properties"]["market"]["required"] == "yes" + assert "required" not in normalized["properties"]["stock_code"] + assert schema["properties"]["server"] == "invalid-property-schema" + assert schema["properties"]["market"]["required"] == "yes" + + +class TestMCPToolSchemaNormalization: + def test_mcp_tool_accepts_property_level_required_booleans(self): + mcp_tool = SimpleNamespace( + name="quote_lookup", + description="Lookup a quote", + inputSchema={ + "type": "object", + "properties": { + "stock_code": {"type": "string", "required": True}, + "market": {"type": "string", "required": False}, + }, + }, + ) + + tool = MCPTool(mcp_tool, MagicMock(), "gf-securities") + + assert tool.parameters["required"] == ["stock_code"] + assert "required" not in tool.parameters["properties"]["stock_code"] + assert "required" not in tool.parameters["properties"]["market"] diff --git a/tests/unit/test_network_utils.py b/tests/unit/test_network_utils.py new file mode 100644 index 000000000..ea3505e38 --- /dev/null +++ b/tests/unit/test_network_utils.py @@ -0,0 +1,51 @@ +import ssl + +import pytest + +from astrbot.core.utils import network_utils + + +def test_create_proxy_client_reuses_shared_ssl_context( + monkeypatch: pytest.MonkeyPatch, +): + captured_calls: list[dict] = [] + headers = {"X-Test-Header": "value"} + + class _FakeAsyncClient: + def __init__(self, **kwargs): + captured_calls.append(kwargs) + + monkeypatch.setattr(network_utils.httpx, "AsyncClient", _FakeAsyncClient) + + network_utils.create_proxy_client("OpenAI") + network_utils.create_proxy_client("OpenAI", proxy="http://127.0.0.1:7890") + network_utils.create_proxy_client("OpenAI", headers=headers) + network_utils.create_proxy_client("OpenAI", proxy="") + + assert len(captured_calls) == 4 + assert "proxy" not in captured_calls[0] + assert captured_calls[1]["proxy"] == "http://127.0.0.1:7890" + assert captured_calls[2]["headers"] is headers + assert "proxy" not in captured_calls[3] + assert isinstance(captured_calls[0]["verify"], ssl.SSLContext) + assert captured_calls[0]["verify"] is captured_calls[1]["verify"] + assert captured_calls[1]["verify"] is captured_calls[2]["verify"] + assert captured_calls[2]["verify"] is captured_calls[3]["verify"] + + +def test_create_proxy_client_allows_verify_override( + monkeypatch: pytest.MonkeyPatch, +): + captured_calls: list[dict] = [] + custom_verify = ssl.create_default_context() + + class _FakeAsyncClient: + def __init__(self, **kwargs): + captured_calls.append(kwargs) + + monkeypatch.setattr(network_utils.httpx, "AsyncClient", _FakeAsyncClient) + + network_utils.create_proxy_client("OpenAI", verify=custom_verify) + + assert len(captured_calls) == 1 + assert captured_calls[0]["verify"] is custom_verify diff --git a/video-fix.patch b/video-fix.patch new file mode 100644 index 000000000..77f879683 Binary files /dev/null and b/video-fix.patch differ