diff --git a/astrbot/cli/commands/cmd_init.py b/astrbot/cli/commands/cmd_init.py index 662b26e67..3c93c0f8a 100644 --- a/astrbot/cli/commands/cmd_init.py +++ b/astrbot/cli/commands/cmd_init.py @@ -112,7 +112,7 @@ async def initialize_astrbot( effective_admin_username = ( admin_username.strip() if admin_username - else str(DEFAULT_CONFIG["dashboard"]["username"]) + else str(DEFAULT_CONFIG["dashboard"]["username"]) # type: ignore[index] ) if admin_username: config = ensure_config_file() diff --git a/astrbot/cli/commands/cmd_run.py b/astrbot/cli/commands/cmd_run.py index 69a3aed62..7dec21678 100644 --- a/astrbot/cli/commands/cmd_run.py +++ b/astrbot/cli/commands/cmd_run.py @@ -332,12 +332,29 @@ def run( lock_file = astrbot_root / "astrbot.lock" lock = FileLock(lock_file, timeout=5) with lock.acquire(): - click.echo("AstrBot is running...") - if backend_only: - click.echo("Visit the dashboard at : https://dash.astrbot.men/") - click.echo("Backend Requests : localhost or based on https") + # Use TUI if in interactive TTY mode + if sys.stdin.isatty() and sys.stdout.isatty(): + from astrbot.cli.commands.cmd_run_tui import run_tui - asyncio.run(run_astrbot(astrbot_root)) + async def wrapped_startup() -> None: + await run_astrbot(astrbot_root) + + asyncio.run( + run_tui( + startup_coro=wrapped_startup, + astrbot_root=astrbot_root, + backend_only=backend_only, + host=os.environ.get("ASTRBOT_HOST"), + port=os.environ.get("ASTRBOT_PORT"), + ) + ) + else: + click.echo("AstrBot is running...") + if backend_only: + click.echo("Visit the dashboard at : https://dash.astrbot.men/") + click.echo("Backend Requests : localhost or based on https") + + asyncio.run(run_astrbot(astrbot_root)) except KeyboardInterrupt: click.echo("AstrBot has been shut down.") except Timeout: diff --git a/astrbot/cli/commands/cmd_run_tui.py b/astrbot/cli/commands/cmd_run_tui.py new file mode 100644 index 000000000..092a0999f --- /dev/null +++ b/astrbot/cli/commands/cmd_run_tui.py @@ -0,0 +1,307 @@ +"""AstrBot Run TUI - A beautiful textual interface for running AstrBot. + +This module provides a Textual-based TUI for `astrbot run` with: +- Animated ASCII logo +- Live log viewer +- Platform status indicators +- Only activates in interactive TTY environments +""" + +from __future__ import annotations + +import sys +import typing +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Any + +from textual.app import App, ComposeResult +from textual.binding import Binding +from textual.containers import Container, Horizontal, Vertical +from textual.reactive import reactive +from textual.widgets import Footer, Header, Log, Static + +if typing.TYPE_CHECKING: + from rich.console import Console + from rich.style import Style + from rich.text import Text +else: + Console: Any = None + Style: Any = None + Text: Any = None + + +# AstrBot ASCII Logo +ASTRBOT_LOGO = r""" + ___ _______.___________..______ .______ ______ .___________. + / \ / | || _ \ | _ \ / __ \ | | + / ^ \ | (----`---| |----`| |_) | | |_) | | | | | `---| |----` + / /_\ \ \ \ | | | / | _ < | | | | | | + / _____ \ .----) | | | | |\ \----.| |_) | | `--' | | | +/__/ \__\ |_______/ |__| | _| `._____||______/ \______/ |__| +""" + + +class AstrBotRunTUI(App): + """Textual TUI for AstrBot run command.""" + + CSS = """ + Screen { + background: $surface; + } + + #logo-container { + height: auto; + padding: 1 2; + background: $surface-darken-1; + border: solid $primary; + } + + #logo-text { + color: $primary; + text-style: bold; + font-family: "JetBrains Mono", "Fira Code", monospace; + } + + #main-container { + height: 1fr; + } + + #log-section { + border: solid $accent; + height: 70%; + margin: 1 2; + } + + #log-header { + background: $accent-darken-1; + padding: 1 2; + color: $text; + text-style: bold; + } + + Log { + background: $surface-darken-2; + color: $text; + border: solid $accent-darken-2; + } + + #status-section { + height: auto; + padding: 1 2; + background: $surface-darken-1; + border-top: solid $primary; + } + + .status-item { + padding: 0 2; + } + + .status-ok { + color: $success; + text-style: bold; + } + + .status-pending { + color: $warning; + } + + .status-label { + color: $text-muted; + } + + .hidden { + display: none; + } + """ + + BINDINGS: typing.ClassVar[list[Binding]] = [ + Binding("q", "quit", "Quit", show=True), + Binding("ctrl+c", "quit", "Quit", show=False), + Binding("l", "toggle_logs", "Toggle Logs", show=True), + ] + + log_visible = reactive(True) + + def __init__( + self, + startup_coro: Callable[[], Awaitable[Any]], + astrbot_root: Path, + backend_only: bool = False, + host: str | None = None, + port: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.startup_coro = startup_coro + self.astrbot_root = astrbot_root + self.backend_only = backend_only + self.host = host + self.port = port + self._animation_frame = 0 + self._startup_done = False + self._log_lines: list[str] = [] + self.console: Any = Console() if Console else None + + def compose(self) -> ComposeResult: + """Create child widgets.""" + yield Header() + + # Animated Logo + with Container(id="logo-container"): + yield Static(self._get_animated_logo(), id="logo-text") + + # Main content + with Vertical(id="main-container"): + # Log viewer + with Container( + id="log-section", classes="" if self.log_visible else "hidden" + ): + yield Static("📋 Live Logs", id="log-header") + yield Log(id="log-viewer") + + # Status bar + with Horizontal(id="status-section"): + yield Static("🌟 AstrBot", classes="status-item status-ok") + yield Static( + f"📁 {self.astrbot_root.name}", + classes="status-item", + id="root-status", + ) + if not self.backend_only: + dashboard_url = ( + f"http://{self.host or 'localhost'}:{self.port or '6185'}" + ) + yield Static( + f"🌐 Dashboard: [link]{dashboard_url}[/link]", + classes="status-item", + id="dashboard-status", + ) + yield Static( + "⚡ Running", classes="status-item status-ok", id="run-status" + ) + + yield Footer() + + def on_mount(self) -> None: + """Called when app is mounted.""" + self.title = "AstrBot" + self.sub_title = "AI Chatbot Framework" + + # Start the startup coroutine + self.set_timer(0.1, self._run_startup) + + # Animate logo + self.set_interval(0.5, self._animate_logo) + + # Get the log widget and configure it + log_widget = self.query_one("#log-viewer", Log) + log_widget.write_line("🚀 AstrBot TUI initialized") + log_widget.write_line(f"📁 Running from: {self.astrbot_root}") + if not self.backend_only: + log_widget.write_line( + f"🌐 Dashboard will be available at: {self.host or 'localhost'}:{self.port or '6185'}" + ) + log_widget.write_line("") + + def _get_animated_logo(self) -> str: + """Get the logo with optional animation effect.""" + lines = ASTRBOT_LOGO.strip().split("\n") + + if self.console and hasattr(self, "_animation_frame"): + # Create animated version with color cycling + frame = self._animation_frame % 4 + colors = ["#00D9FF", "#00FF87", "#FFD700", "#FF6B6B"] + color = colors[frame] + + text = Text() + for i, line in enumerate(lines): + style = Style(color=color, bold=True) if i == 0 else Style(color=color) + text.append(line + "\n", style=style) + return str(text) + + return ASTRBOT_LOGO + + def _animate_logo(self) -> None: + """Update the animated logo.""" + self._animation_frame = (self._animation_frame + 1) % 4 + logo_widget = self.query_one("#logo-text", Static) + logo_widget.update(self._get_animated_logo()) + + async def _run_startup(self) -> None: + """Run the AstrBot startup coroutine.""" + if self._startup_done: + return + self._startup_done = True + + try: + log_widget = self.query_one("#log-viewer", Log) + log_widget.write_line("⏳ Initializing AstrBot...") + + await self.startup_coro() + + log_widget.write_line("") + log_widget.write_line("✅ AstrBot started successfully!") + except Exception as e: + log_widget = self.query_one("#log-viewer", Log) + log_widget.write_line(f"❌ Error during startup: {e}") + log_widget.write_line("Check logs for details.") + + def action_toggle_logs(self) -> None: + """Toggle log visibility.""" + self.log_visible = not self.log_visible + log_section = self.query_one("#log-section", Container) + if self.log_visible: + log_section.remove_class("hidden") + else: + log_section.add_class("hidden") + + async def action_quit(self) -> None: + """Quit the application.""" + self.exit() + + def write_log(self, message: str) -> None: + """Write a message to the log viewer (can be called from outside).""" + log_widget = self.query_one("#log-viewer", Log) + log_widget.write_line(message) + + +def is_interactive_tty() -> bool: + """Check if we're running in an interactive TTY.""" + return sys.stdin.isatty() and sys.stdout.isatty() + + +async def run_tui( + startup_coro: Callable[[], Awaitable[Any]], + astrbot_root: Path, + backend_only: bool = False, + host: str | None = None, + port: str | None = None, +) -> None: + """Run the AstrBot TUI. + + Args: + startup_coro: Coroutine to run on startup + astrbot_root: AstrBot root directory + backend_only: Whether backend-only mode is enabled + host: Dashboard host + port: Dashboard port + """ + if not is_interactive_tty(): + # Not interactive, run without TUI + await startup_coro() + return + + app = AstrBotRunTUI( + startup_coro=startup_coro, + astrbot_root=astrbot_root, + backend_only=backend_only, + host=host, + port=port, + ) + + try: + await app.run_async() + except Exception: + # Fallback to non-TUI mode + await startup_coro() diff --git a/astrbot/cli/commands/tui_async.py b/astrbot/cli/commands/tui_async.py index 642d66a4a..a57d4dcb0 100644 --- a/astrbot/cli/commands/tui_async.py +++ b/astrbot/cli/commands/tui_async.py @@ -117,8 +117,8 @@ class TUIClient: # Create new session for TUI new_session_resp = await self._client.get( - "/api/chat/new_session", - params={"platform_id": "webchat"}, + "/api/tui/new_session", + params={"platform_id": "tui"}, headers=self._headers, ) if new_session_resp.status_code != 200: @@ -163,7 +163,7 @@ class TUIClient: try: resp = await self._client.get( - "/api/chat/get_session", + "/api/tui/get_session", params={"session_id": self.conversation_id}, headers=self._headers, ) @@ -317,9 +317,9 @@ class TUIClient: self.state.status = "Waiting for response..." try: - # Format umo for webchat + # Format umo for tui umo = ( - f"webchat:FriendMessage:webchat!{self.username}!{self.conversation_id}" + f"tui:FriendMessage:tui!{self.username}!{self.conversation_id}" ) # Reset parser for new stream @@ -328,7 +328,7 @@ class TUIClient: # Send message and stream response using proper SSE async with self._client.stream( "POST", - "/api/chat/chat", + "/api/tui/chat", headers=self._headers, json={ "umo": umo, diff --git a/astrbot/core/agent/runners/coze/coze_api_client.py b/astrbot/core/agent/runners/coze/coze_api_client.py index 41408ffb0..ce1de24d4 100644 --- a/astrbot/core/agent/runners/coze/coze_api_client.py +++ b/astrbot/core/agent/runners/coze/coze_api_client.py @@ -145,7 +145,7 @@ class CozeAPIClient: session = await self._ensure_session() url = f"{self.api_base}/v3/chat" - payload = { + payload: dict[str, Any] = { "bot_id": bot_id, "user_id": user_id, "stream": stream, diff --git a/astrbot/core/astrbot_config_mgr.py b/astrbot/core/astrbot_config_mgr.py index ecc6782fd..c42418319 100644 --- a/astrbot/core/astrbot_config_mgr.py +++ b/astrbot/core/astrbot_config_mgr.py @@ -13,7 +13,7 @@ from astrbot.core.utils.shared_preferences import SharedPreferences _VT = TypeVar("_VT") -class ConfInfo(TypedDict): +class ConfInfo(TypedDict, total=False): """Configuration information for a specific session or platform.""" id: str # UUID of the configuration or "default" diff --git a/astrbot/core/persona_mgr.py b/astrbot/core/persona_mgr.py index da865febd..01b6620af 100644 --- a/astrbot/core/persona_mgr.py +++ b/astrbot/core/persona_mgr.py @@ -433,7 +433,7 @@ class PersonaManager: user_turn = not user_turn try: - persona = Personality( + persona = Personality( # type: ignore[misc] **persona_cfg, _begin_dialogs_processed=bd_processed, _mood_imitation_dialogs_processed="", # deprecated diff --git a/astrbot/core/platform/manager.py b/astrbot/core/platform/manager.py index cfb5f0c35..9ff9177cb 100644 --- a/astrbot/core/platform/manager.py +++ b/astrbot/core/platform/manager.py @@ -10,6 +10,7 @@ from astrbot.core.utils.webhook_utils import ensure_platform_webhook_config from .platform import Platform, PlatformStatus from .register import platform_cls_map +from .sources.tui.tui_adapter import TUIAdapter from .sources.webchat.webchat_adapter import WebChatAdapter PLATFORM_ADAPTER_MODULES: dict[str, str] = { @@ -117,6 +118,11 @@ class PlatformManager: self.platform_insts.append(webchat_inst) self._start_platform_task("webchat", webchat_inst) + # TUI + tui_inst = TUIAdapter({}, self.settings, self.event_queue) + self.platform_insts.append(tui_inst) + self._start_platform_task("tui", tui_inst) + async def load_platform(self, platform_config: dict) -> None: """实例化一个平台""" # 动态导入 diff --git a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py index 864953400..9b108fa6e 100644 --- a/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py +++ b/astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py @@ -3,7 +3,7 @@ import base64 import os import random import uuid -from typing import cast +from typing import Any, cast import aiofiles import anyio @@ -75,7 +75,7 @@ class QQOfficialMessageEvent(AstrMessageEvent): # 先标记事件层“已执行发送操作”,避免异常路径遗漏 await super().send_streaming(generator, use_fallback) # QQ C2C 流式协议:开始/中间分片使用 state=1,结束分片使用 state=10 - stream_payload = {"state": 1, "id": None, "index": 0, "reset": False} + stream_payload: dict[str, Any] = {"state": 1, "id": None, "index": 0, "reset": False} last_edit_time = 0 # 上次发送分片的时间 throttle_interval = 1 # 分片间最短间隔 (秒) ret = None @@ -489,7 +489,7 @@ class QQOfficialMessageEvent(AstrMessageEvent): ) -> Media | None: """上传媒体文件""" # 构建基础payload - payload = {"file_type": file_type, "srv_send_msg": srv_send_msg} + payload: dict[str, Any] = {"file_type": file_type, "srv_send_msg": srv_send_msg} if file_name: payload["file_name"] = file_name @@ -572,7 +572,7 @@ class QQOfficialMessageEvent(AstrMessageEvent): logger.error(f"[QQOfficial] post_c2c_message: 响应不是 dict: {result}") return None - return message.Message(**result) + return message.Message(**cast(dict[str, Any], result)) # type: ignore[arg-type] @staticmethod async def _parse_to_qqofficial(message: MessageChain): diff --git a/astrbot/core/platform/sources/telegram/tg_event.py b/astrbot/core/platform/sources/telegram/tg_event.py index f338a488a..4f8dc08ac 100644 --- a/astrbot/core/platform/sources/telegram/tg_event.py +++ b/astrbot/core/platform/sources/telegram/tg_event.py @@ -307,7 +307,7 @@ class TelegramPlatformEvent(AstrMessageEvent): else: send_coro = client.send_photo media_kwarg = {"photo": image_path} - await send_coro(**media_kwarg, **cast(Any, payload)) + await send_coro(**cast(Any, media_kwarg), **cast(Any, payload)) elif isinstance(i, File): path = await i.get_file() name = i.name or os.path.basename(path) diff --git a/astrbot/core/platform/sources/tui/__init__.py b/astrbot/core/platform/sources/tui/__init__.py new file mode 100644 index 000000000..6ac1a858a --- /dev/null +++ b/astrbot/core/platform/sources/tui/__init__.py @@ -0,0 +1,5 @@ +from .tui_adapter import TUIAdapter +from .tui_event import TUIMessageEvent +from .tui_queue_mgr import TUIQueueMgr, tui_queue_mgr + +__all__ = ["TUIAdapter", "TUIMessageEvent", "TUIQueueMgr", "tui_queue_mgr"] diff --git a/astrbot/core/platform/sources/tui/tui_adapter.py b/astrbot/core/platform/sources/tui/tui_adapter.py new file mode 100644 index 000000000..8d4de15ca --- /dev/null +++ b/astrbot/core/platform/sources/tui/tui_adapter.py @@ -0,0 +1,239 @@ +import asyncio +import os +import time +from collections.abc import Callable, Coroutine +from typing import Any, cast + +from astrbot import logger +from astrbot.core import db_helper +from astrbot.core.db.po import PlatformMessageHistory +from astrbot.core.message.message_event_result import MessageChain +from astrbot.core.platform import ( + AstrBotMessage, + MessageMember, + MessageType, + Platform, + PlatformMetadata, +) +from astrbot.core.platform.astr_message_event import MessageSesion +from astrbot.core.platform.register import register_platform_adapter +from astrbot.core.platform.sources.webchat.message_parts_helper import ( + message_chain_to_storage_message_parts, + parse_webchat_message_parts, +) +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + +from .tui_event import TUIMessageEvent +from .tui_queue_mgr import TUIQueueMgr, tui_queue_mgr + + +def _extract_conversation_id(session_id: str) -> str: + """Extract raw TUI conversation id from event/session id.""" + if session_id.startswith("tui!"): + parts = session_id.split("!", 2) + if len(parts) == 3: + return parts[2] + return session_id + + +class QueueListener: + def __init__( + self, + tui_queue_mgr: TUIQueueMgr, + callback: Callable, + stop_event: asyncio.Event, + ) -> None: + self.tui_queue_mgr = tui_queue_mgr + self.callback = callback + self.stop_event = stop_event + + async def run(self) -> None: + """Register callback and keep adapter task alive.""" + self.tui_queue_mgr.set_listener(self.callback) + try: + await self.stop_event.wait() + finally: + await self.tui_queue_mgr.clear_listener() + + +@register_platform_adapter("tui", "tui") +class TUIAdapter(Platform): + def __init__( + self, + platform_config: dict, + platform_settings: dict, + event_queue: asyncio.Queue, + ) -> None: + super().__init__(platform_config, event_queue) + + self.settings = platform_settings + self.imgs_dir = os.path.join(get_astrbot_data_path(), "tui", "imgs") + self.attachments_dir = os.path.join(get_astrbot_data_path(), "attachments") + os.makedirs(self.imgs_dir, exist_ok=True) + os.makedirs(self.attachments_dir, exist_ok=True) + + self.metadata = PlatformMetadata( + name="tui", + description="tui", + id="tui", + support_proactive_message=True, + ) + self._shutdown_event = asyncio.Event() + self._tui_queue_mgr = tui_queue_mgr + + async def send_by_session( + self, + session: MessageSesion, + message_chain: MessageChain, + ) -> None: + conversation_id = _extract_conversation_id(session.session_id) + active_request_ids = self._tui_queue_mgr.list_back_request_ids(conversation_id) + stream_request_ids = [ + req_id for req_id in active_request_ids if not req_id.startswith("ws_sub_") + ] + target_request_ids = stream_request_ids or active_request_ids + + if not target_request_ids: + try: + await self._save_proactive_message(conversation_id, message_chain) + except Exception as e: + logger.error( + f"[TUIAdapter] Failed to save proactive message: {e}", + exc_info=True, + ) + await super().send_by_session(session, message_chain) + return + + for request_id in target_request_ids: + await TUIMessageEvent._send( + request_id, + message_chain, + session.session_id, + streaming=True, + emit_complete=True, + ) + + if not stream_request_ids: + try: + await self._save_proactive_message(conversation_id, message_chain) + except Exception as e: + logger.error( + f"[TUIAdapter] Failed to save proactive message: {e}", + exc_info=True, + ) + + await super().send_by_session(session, message_chain) + + async def _save_proactive_message( + self, + conversation_id: str, + message_chain: MessageChain, + ) -> None: + message_parts = await message_chain_to_storage_message_parts( + message_chain, + insert_attachment=db_helper.insert_attachment, + attachments_dir=self.attachments_dir, + ) + if not message_parts: + return + + await db_helper.insert_platform_message_history( + platform_id="tui", + user_id=conversation_id, + content={"type": "bot", "message": message_parts}, + sender_id="bot", + sender_name="bot", + ) + + async def _get_message_history( + self, message_id: int + ) -> PlatformMessageHistory | None: + return await db_helper.get_platform_message_history_by_id(message_id) + + async def _parse_message_parts( + self, + message_parts: list, + depth: int = 0, + max_depth: int = 1, + ) -> tuple[list, list[str]]: + """Parse message parts list, return message components and plain text lists.""" + + async def get_reply_parts( + message_id: Any, + ) -> tuple[list[dict], str | None, str | None] | None: + history = await self._get_message_history(message_id) + if not history or not history.content: + return None + + reply_parts = history.content.get("message", []) + if not isinstance(reply_parts, list): + return None + + return reply_parts, history.sender_id, history.sender_name + + components, text_parts, _ = await parse_webchat_message_parts( + message_parts, + strict=False, + include_empty_plain=True, + verify_media_path_exists=False, + reply_history_getter=get_reply_parts, + current_depth=depth, + max_reply_depth=max_depth, + cast_reply_id_to_str=False, + ) + return components, text_parts + + async def convert_message(self, data: tuple) -> AstrBotMessage: + username, cid, payload = data + + abm = AstrBotMessage() + abm.self_id = "tui" + abm.sender = MessageMember(username, username) + + abm.type = MessageType.FRIEND_MESSAGE + + abm.session_id = f"tui!{username}!{cid}" + + abm.message_id = payload.get("message_id") + + message_parts = payload.get("message", []) + abm.message, message_str_parts = await self._parse_message_parts(message_parts) + + logger.debug(f"TUIAdapter: {abm.message}") + + abm.timestamp = int(time.time()) + abm.message_str = "".join(message_str_parts) + abm.raw_message = data + return abm + + def run(self) -> Coroutine[Any, Any, None]: + async def callback(data: tuple) -> None: + abm = await self.convert_message(data) + await self.handle_msg(abm) + + bot = QueueListener(self._tui_queue_mgr, callback, self._shutdown_event) + return bot.run() + + def meta(self) -> PlatformMetadata: + return self.metadata + + async def handle_msg(self, message: AstrBotMessage) -> None: + message_event = TUIMessageEvent( + message_str=message.message_str, + message_obj=message, + platform_meta=self.meta(), + session_id=message.session_id, + ) + + _, _, payload = cast(tuple[Any, Any, dict[str, Any]], message.raw_message) + message_event.set_extra("selected_provider", payload.get("selected_provider")) + message_event.set_extra("selected_model", payload.get("selected_model")) + message_event.set_extra( + "enable_streaming", payload.get("enable_streaming", True) + ) + message_event.set_extra("action_type", payload.get("action_type")) + + self.commit_event(message_event) + + async def terminate(self) -> None: + self._shutdown_event.set() diff --git a/astrbot/core/platform/sources/tui/tui_event.py b/astrbot/core/platform/sources/tui/tui_event.py new file mode 100644 index 000000000..6ea7ec7eb --- /dev/null +++ b/astrbot/core/platform/sources/tui/tui_event.py @@ -0,0 +1,203 @@ +import base64 +import json +import os +import shutil +import uuid + +import aiofiles + +from astrbot.api import logger +from astrbot.api.event import AstrMessageEvent, MessageChain +from astrbot.api.message_components import File, Image, Json, Plain, Record +from astrbot.core.utils.astrbot_path import get_astrbot_data_path + +from .tui_queue_mgr import tui_queue_mgr + +attachments_dir = os.path.join(get_astrbot_data_path(), "attachments") + + +def _extract_conversation_id(session_id: str) -> str: + """Extract raw TUI conversation id from event/session id.""" + if session_id.startswith("tui!"): + parts = session_id.split("!", 2) + if len(parts) == 3: + return parts[2] + return session_id + + +class TUIMessageEvent(AstrMessageEvent): + def __init__(self, message_str, message_obj, platform_meta, session_id) -> None: + super().__init__(message_str, message_obj, platform_meta, session_id) + os.makedirs(attachments_dir, exist_ok=True) + + @staticmethod + async def _send( + message_id: str, + message: MessageChain | None, + session_id: str, + streaming: bool = False, + emit_complete: bool = False, + ) -> str | None: + request_id = str(message_id) + conversation_id = _extract_conversation_id(session_id) + tui_back_queue = tui_queue_mgr.get_or_create_back_queue( + request_id, + conversation_id, + ) + if not message: + await tui_back_queue.put( + { + "type": "end", + "data": "", + "streaming": False, + "message_id": message_id, + }, + ) + return + + data = "" + for comp in message.chain: + if isinstance(comp, Plain): + data = comp.text + await tui_back_queue.put( + { + "type": "plain", + "data": data, + "streaming": streaming, + "chain_type": message.type, + "message_id": message_id, + }, + ) + elif isinstance(comp, Json): + await tui_back_queue.put( + { + "type": "plain", + "data": json.dumps(comp.data, ensure_ascii=False), + "streaming": streaming, + "chain_type": message.type, + "message_id": message_id, + }, + ) + elif isinstance(comp, Image): + filename = f"{uuid.uuid4()!s}.jpg" + path = os.path.join(attachments_dir, filename) + image_base64 = await comp.convert_to_base64() + async with aiofiles.open(path, "wb") as f: + await f.write(base64.b64decode(image_base64)) + data = f"[IMAGE]{filename}" + await tui_back_queue.put( + { + "type": "image", + "data": data, + "streaming": streaming, + "message_id": message_id, + }, + ) + elif isinstance(comp, Record): + filename = f"{uuid.uuid4()!s}.wav" + path = os.path.join(attachments_dir, filename) + record_base64 = await comp.convert_to_base64() + async with aiofiles.open(path, "wb") as f: + await f.write(base64.b64decode(record_base64)) + data = f"[RECORD]{filename}" + await tui_back_queue.put( + { + "type": "record", + "data": data, + "streaming": streaming, + "message_id": message_id, + }, + ) + elif isinstance(comp, File): + file_path = await comp.get_file() + original_name = comp.name or os.path.basename(file_path) + ext = os.path.splitext(original_name)[1] or "" + filename = f"{uuid.uuid4()!s}{ext}" + dest_path = os.path.join(attachments_dir, filename) + shutil.copy2(file_path, dest_path) + data = f"[FILE]{filename}" + await tui_back_queue.put( + { + "type": "file", + "data": data, + "streaming": streaming, + "message_id": message_id, + }, + ) + else: + logger.debug(f"TUI ignores: {comp.type}") + + if emit_complete: + await tui_back_queue.put( + { + "type": "complete", + "data": data, + "streaming": streaming, + "chain_type": message.type, + "message_id": message_id, + }, + ) + + return data + + async def send(self, message: MessageChain | None) -> None: + message_id = self.message_obj.message_id + await TUIMessageEvent._send(message_id, message, session_id=self.session_id) + await super().send(MessageChain([])) + + async def send_streaming(self, generator, use_fallback: bool = False) -> None: + final_data = "" + reasoning_content = "" + message_id = self.message_obj.message_id + request_id = str(message_id) + conversation_id = _extract_conversation_id(self.session_id) + tui_back_queue = tui_queue_mgr.get_or_create_back_queue( + request_id, + conversation_id, + ) + async for chain in generator: + if chain.type == "audio_chunk": + audio_b64 = "" + text = None + + if chain.chain and isinstance(chain.chain[0], Plain): + audio_b64 = chain.chain[0].text + + if len(chain.chain) > 1 and isinstance(chain.chain[1], Json): + text = chain.chain[1].data.get("text") + + payload = { + "type": "audio_chunk", + "data": audio_b64, + "streaming": True, + "message_id": message_id, + } + if text: + payload["text"] = text + + await tui_back_queue.put(payload) + continue + + r = await TUIMessageEvent._send( + message_id=message_id, + message=chain, + session_id=self.session_id, + streaming=True, + ) + if not r: + continue + if chain.type == "reasoning": + reasoning_content += chain.get_plain_text() + else: + final_data += r + + await tui_back_queue.put( + { + "type": "complete", + "data": final_data, + "reasoning": reasoning_content, + "streaming": True, + "message_id": message_id, + }, + ) + await super().send_streaming(generator, use_fallback) diff --git a/astrbot/core/platform/sources/tui/tui_queue_mgr.py b/astrbot/core/platform/sources/tui/tui_queue_mgr.py new file mode 100644 index 000000000..ac770f820 --- /dev/null +++ b/astrbot/core/platform/sources/tui/tui_queue_mgr.py @@ -0,0 +1,164 @@ +import asyncio +from collections.abc import Awaitable, Callable + +from astrbot import logger + + +class TUIQueueMgr: + def __init__(self, queue_maxsize: int = 128, back_queue_maxsize: int = 512) -> None: + self.queues: dict[str, asyncio.Queue] = {} + """Conversation ID to asyncio.Queue mapping""" + self.back_queues: dict[str, asyncio.Queue] = {} + """Request ID to asyncio.Queue mapping for responses""" + self._conversation_back_requests: dict[str, set[str]] = {} + self._request_conversation: dict[str, str] = {} + self._queue_close_events: dict[str, asyncio.Event] = {} + self._listener_tasks: dict[str, asyncio.Task] = {} + self._listener_callback: Callable[[tuple], Awaitable[None]] | None = None + self.queue_maxsize = queue_maxsize + self.back_queue_maxsize = back_queue_maxsize + + def get_or_create_queue(self, conversation_id: str) -> asyncio.Queue: + """Get or create a queue for the given conversation ID""" + if conversation_id not in self.queues: + self.queues[conversation_id] = asyncio.Queue(maxsize=self.queue_maxsize) + self._queue_close_events[conversation_id] = asyncio.Event() + self._start_listener_if_needed(conversation_id) + return self.queues[conversation_id] + + def get_or_create_back_queue( + self, + request_id: str, + conversation_id: str | None = None, + ) -> asyncio.Queue: + """Get or create a back queue for the given request ID""" + if request_id not in self.back_queues: + self.back_queues[request_id] = asyncio.Queue( + maxsize=self.back_queue_maxsize + ) + if conversation_id: + self._request_conversation[request_id] = conversation_id + if conversation_id not in self._conversation_back_requests: + self._conversation_back_requests[conversation_id] = set() + self._conversation_back_requests[conversation_id].add(request_id) + return self.back_queues[request_id] + + def remove_back_queue(self, request_id: str) -> None: + """Remove back queue for the given request ID""" + self.back_queues.pop(request_id, None) + conversation_id = self._request_conversation.pop(request_id, None) + if conversation_id: + request_ids = self._conversation_back_requests.get(conversation_id) + if request_ids is not None: + request_ids.discard(request_id) + if not request_ids: + self._conversation_back_requests.pop(conversation_id, None) + + def remove_queues(self, conversation_id: str) -> None: + """Remove queues for the given conversation ID""" + for request_id in list( + self._conversation_back_requests.get(conversation_id, set()) + ): + self.remove_back_queue(request_id) + self._conversation_back_requests.pop(conversation_id, None) + self.remove_queue(conversation_id) + + def remove_queue(self, conversation_id: str) -> None: + """Remove input queue and listener for the given conversation ID""" + self.queues.pop(conversation_id, None) + + close_event = self._queue_close_events.pop(conversation_id, None) + if close_event is not None: + close_event.set() + + task = self._listener_tasks.pop(conversation_id, None) + if task is not None: + task.cancel() + + def list_back_request_ids(self, conversation_id: str) -> list[str]: + """List active back-queue request IDs for a conversation.""" + return list(self._conversation_back_requests.get(conversation_id, set())) + + def has_queue(self, conversation_id: str) -> bool: + """Check if a queue exists for the given conversation ID""" + return conversation_id in self.queues + + def set_listener( + self, + callback: Callable[[tuple], Awaitable[None]], + ) -> None: + self._listener_callback = callback + for conversation_id in list(self.queues.keys()): + self._start_listener_if_needed(conversation_id) + + async def clear_listener(self) -> None: + self._listener_callback = None + for close_event in list(self._queue_close_events.values()): + close_event.set() + self._queue_close_events.clear() + + listener_tasks = list(self._listener_tasks.values()) + for task in listener_tasks: + task.cancel() + if listener_tasks: + await asyncio.gather(*listener_tasks, return_exceptions=True) + self._listener_tasks.clear() + + def _start_listener_if_needed(self, conversation_id: str) -> None: + if self._listener_callback is None: + return + if conversation_id in self._listener_tasks: + task = self._listener_tasks[conversation_id] + if not task.done(): + return + queue = self.queues.get(conversation_id) + close_event = self._queue_close_events.get(conversation_id) + if queue is None or close_event is None: + return + task = asyncio.create_task( + self._listen_to_queue(conversation_id, queue, close_event), + name=f"tui_listener_{conversation_id}", + ) + self._listener_tasks[conversation_id] = task + task.add_done_callback( + lambda _: self._listener_tasks.pop(conversation_id, None) + ) + logger.debug(f"Started listener for TUI conversation: {conversation_id}") + + async def _listen_to_queue( + self, + conversation_id: str, + queue: asyncio.Queue, + close_event: asyncio.Event, + ) -> None: + while True: + get_task = asyncio.create_task(queue.get()) + close_task = asyncio.create_task(close_event.wait()) + try: + done, pending = await asyncio.wait( + {get_task, close_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + if close_task in done: + break + data = get_task.result() + if self._listener_callback is None: + continue + try: + await self._listener_callback(data) + except Exception as e: + logger.error( + f"Error processing message from TUI conversation {conversation_id}: {e}" + ) + except asyncio.CancelledError: + break + finally: + if not get_task.done(): + get_task.cancel() + if not close_task.done(): + close_task.cancel() + + +tui_queue_mgr = TUIQueueMgr() diff --git a/astrbot/core/provider/sources/mimo_api_common.py b/astrbot/core/provider/sources/mimo_api_common.py index d3bf75e66..56b15bab7 100644 --- a/astrbot/core/provider/sources/mimo_api_common.py +++ b/astrbot/core/provider/sources/mimo_api_common.py @@ -1,6 +1,7 @@ import base64 import uuid from pathlib import Path +from typing import Any, cast from urllib.parse import urlparse import httpx @@ -60,7 +61,7 @@ def create_http_client(timeout: int | None, proxy: str) -> httpx.AsyncClient: if proxy: logger.info("[MiMo API] Using proxy: %s", proxy) client_kwargs["proxy"] = proxy - return httpx.AsyncClient(**client_kwargs) + return httpx.AsyncClient(**cast(dict[str, Any], client_kwargs)) def build_api_url(api_base: str) -> str: diff --git a/astrbot/core/star/base.py b/astrbot/core/star/base.py index efd2527b1..8c3d743f1 100644 --- a/astrbot/core/star/base.py +++ b/astrbot/core/star/base.py @@ -1,7 +1,8 @@ from __future__ import annotations import logging -from typing import Any, Protocol +from asyncio import Queue +from typing import TYPE_CHECKING, Any, Protocol from astrbot.core import html_renderer from astrbot.core.utils.command_parser import CommandParserMixin @@ -9,6 +10,11 @@ from astrbot.core.utils.plugin_kv_store import PluginKVStoreMixin from .star import StarMetadata, star_map, star_registry +if TYPE_CHECKING: + from astrbot.core.provider.func_tool_manager import FunctionToolManager + from astrbot.core.provider.manager import ProviderManager + from astrbot.core.provider.provider import Provider + logger = logging.getLogger("astrbot") @@ -21,6 +27,18 @@ class Star(CommandParserMixin, PluginKVStoreMixin): class _ContextLike(Protocol): def get_config(self, umo: str | None = None) -> Any: ... + def get_using_provider(self, umo: str | None = None) -> "Provider | None": ... + + def get_llm_tool_manager(self) -> "FunctionToolManager": ... + + def get_event_queue(self) -> Queue[Any]: ... + + @property + def conversation_manager(self) -> Any: ... + + @property + def provider_manager(self) -> "ProviderManager": ... + def __init__(self, context: _ContextLike, config: dict | None = None) -> None: self.context = context diff --git a/astrbot/core/tools/send_message.py b/astrbot/core/tools/send_message.py index fc58f41fb..f605bf818 100644 --- a/astrbot/core/tools/send_message.py +++ b/astrbot/core/tools/send_message.py @@ -8,7 +8,7 @@ from __future__ import annotations import json import os import uuid -from typing import Any, TypedDict +from typing import Any, TypedDict, cast import anyio from pydantic import Field @@ -111,21 +111,23 @@ class SendMessageToUserTool(FunctionTool[AstrAgentContext]): if not isinstance(msg, dict): return f"error: messages[{idx}] should be an object." - if "type" not in msg: + msg_dict: dict[str, Any] = cast(dict[str, Any], msg) + + if "type" not in msg_dict: return f"error: messages[{idx}].type is required." - msg_type = str(msg["type"]).lower() + msg_type = str(msg_dict["type"]).lower() _file_from_sandbox = False try: if msg_type == "plain": - text = str(msg.get("text", "")).strip() + text = str(msg_dict.get("text", "")).strip() if not text: return f"error: messages[{idx}].text is required for plain component." components.append(Comp.Plain(text=text)) elif msg_type == "image": - path = msg.get("path") - url = msg.get("url") + path = msg_dict.get("path") + url = msg_dict.get("url") if path: ( local_path, @@ -137,8 +139,8 @@ class SendMessageToUserTool(FunctionTool[AstrAgentContext]): else: return f"error: messages[{idx}] must include path or url for image component." elif msg_type == "record": - path = msg.get("path") - url = msg.get("url") + path = msg_dict.get("path") + url = msg_dict.get("url") if path: ( local_path, @@ -150,8 +152,8 @@ class SendMessageToUserTool(FunctionTool[AstrAgentContext]): else: return f"error: messages[{idx}] must include path or url for record component." elif msg_type == "video": - path = msg.get("path") - url = msg.get("url") + path = msg_dict.get("path") + url = msg_dict.get("url") if path: ( local_path, @@ -163,10 +165,10 @@ class SendMessageToUserTool(FunctionTool[AstrAgentContext]): else: return f"error: messages[{idx}] must include path or url for video component." elif msg_type == "file": - path = msg.get("path") - url = msg.get("url") + path = msg_dict.get("path") + url = msg_dict.get("url") name = ( - msg.get("text") + msg_dict.get("text") or (os.path.basename(path) if path else "") or (os.path.basename(url) if url else "") or "file" @@ -182,7 +184,7 @@ class SendMessageToUserTool(FunctionTool[AstrAgentContext]): else: return f"error: messages[{idx}] must include path or url for file component." elif msg_type == "mention_user": - mention_user_id = msg.get("mention_user_id") + mention_user_id = msg_dict.get("mention_user_id") if not mention_user_id: return f"error: messages[{idx}].mention_user_id is required for mention_user component." components.append( diff --git a/astrbot/core/utils/shared_preferences.py b/astrbot/core/utils/shared_preferences.py index 7a8e6eb8d..e629fb0be 100644 --- a/astrbot/core/utils/shared_preferences.py +++ b/astrbot/core/utils/shared_preferences.py @@ -44,7 +44,7 @@ class SharedPreferences: scope: str, scope_id: str, key: str, - default: _VT = None, + default: _VT = None, # type: ignore[invalid-parameter-default] ) -> _VT: """获取指定范围和键的偏好设置""" if scope_id is not None and key is not None: @@ -72,7 +72,7 @@ class SharedPreferences: self, umo: str, key: str, - default: _VT = None, + default: _VT = None, # type: ignore[invalid-parameter-default] ) -> _VT: ... @overload @@ -103,7 +103,7 @@ class SharedPreferences: self, umo: str | None, key: str | None = None, - default: _VT = None, + default: _VT = None, # type: ignore[invalid-parameter-default] ) -> _VT | list[Preference]: """获取会话范围的偏好设置 @@ -117,12 +117,12 @@ class SharedPreferences: async def global_get(self, key: None, default: Any = None) -> list[Preference]: ... @overload - async def global_get(self, key: str, default: _VT = None) -> _VT: ... + async def global_get(self, key: str, default: _VT = None) -> _VT: ... # type: ignore[invalid-parameter-default] async def global_get( self, key: str | None, - default: _VT = None, + default: _VT = None, # type: ignore[invalid-parameter-default] ) -> _VT | list[Preference]: """获取全局范围的偏好设置 @@ -169,7 +169,7 @@ class SharedPreferences: def get( self, key: str, - default: _VT = None, + default: _VT = None, # type: ignore[invalid-parameter-default] scope: str | None = None, scope_id: str | None = "", ) -> _VT: diff --git a/astrbot/dashboard/routes/__init__.py b/astrbot/dashboard/routes/__init__.py index 4af4e4178..4ce76de31 100644 --- a/astrbot/dashboard/routes/__init__.py +++ b/astrbot/dashboard/routes/__init__.py @@ -23,6 +23,7 @@ from .static_file import StaticFileRoute from .subagent import SubAgentRoute from .t2i import T2iRoute from .tools import ToolsRoute +from .tui_chat import TUIChatRoute from .update import UpdateRoute __all__ = [ @@ -52,5 +53,6 @@ __all__ = [ "SubAgentRoute", "T2iRoute", "ToolsRoute", + "TUIChatRoute", "UpdateRoute", ] diff --git a/astrbot/dashboard/routes/tui_chat.py b/astrbot/dashboard/routes/tui_chat.py new file mode 100644 index 000000000..560cac17b --- /dev/null +++ b/astrbot/dashboard/routes/tui_chat.py @@ -0,0 +1,755 @@ +import asyncio +import json +import os +import uuid +from contextlib import asynccontextmanager +from pathlib import Path +from typing import cast + +import anyio +from quart import Response as QuartResponse +from quart import g, make_response, request, send_file + +from astrbot.core import logger +from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from astrbot.core.db import BaseDatabase +from astrbot.core.platform.message_type import MessageType +from astrbot.core.platform.sources.tui.tui_queue_mgr import tui_queue_mgr +from astrbot.core.platform.sources.webchat.message_parts_helper import ( + build_webchat_message_parts, + create_attachment_part_from_existing_file, + strip_message_parts_path_fields, + webchat_message_parts_have_content, +) +from astrbot.core.utils.active_event_registry import active_event_registry +from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.datetime_utils import to_utc_isoformat + +from .route import Response, Route, RouteContext + + +@asynccontextmanager +async def track_conversation(convs: dict, conv_id: str): + convs[conv_id] = True + try: + yield + finally: + convs.pop(conv_id, None) + + +async def _poll_tui_stream_result(back_queue, username: str): + try: + result = await asyncio.wait_for(back_queue.get(), timeout=1) + except asyncio.TimeoutError: + return None, False + except asyncio.CancelledError: + logger.debug(f"[TUI] User {username} disconnected.") + return None, True + except Exception as e: + logger.error(f"TUI stream error: {e}") + return None, False + return result, False + + +def _resolve_path(path: str) -> Path: + return Path(path).resolve(strict=False) + + +class TUIChatRoute(Route): + def __init__( + self, + context: RouteContext, + db: BaseDatabase, + core_lifecycle: AstrBotCoreLifecycle, + ) -> None: + super().__init__(context) + self.routes = { + "/tui/chat": ("POST", self.chat), + "/tui/new_session": ("GET", self.new_session), + "/tui/sessions": ("GET", self.get_sessions), + "/tui/get_session": ("GET", self.get_session), + "/tui/stop": ("POST", self.stop_session), + "/tui/delete_session": ("GET", self.delete_tui_session), + "/tui/batch_delete_sessions": ("POST", self.batch_delete_sessions), + "/tui/update_session_display_name": ( + "POST", + self.update_session_display_name, + ), + "/tui/get_file": ("GET", self.get_file), + "/tui/get_attachment": ("GET", self.get_attachment), + "/tui/post_file": ("POST", self.post_file), + } + self.core_lifecycle = core_lifecycle + self.register_routes() + self.attachments_dir = os.path.join(get_astrbot_data_path(), "attachments") + os.makedirs(self.attachments_dir, exist_ok=True) + + self.supported_imgs = ["jpg", "jpeg", "png", "gif", "webp"] + self.conv_mgr = core_lifecycle.conversation_manager + self.platform_history_mgr = core_lifecycle.platform_message_history_manager + self.db = db + self.umop_config_router = core_lifecycle.umop_config_router + + self.running_convs: dict[str, bool] = {} + + async def get_file(self): + filename = request.args.get("filename") + if not filename: + return Response().error("Missing key: filename").__dict__ + + try: + file_path = os.path.join(self.attachments_dir, os.path.basename(filename)) + resolved_file_path = _resolve_path(file_path) + resolved_base_dir = _resolve_path(self.attachments_dir) + + if not await anyio.Path(resolved_file_path).exists(): + return Response().error("File not found").__dict__ + + try: + resolved_file_path.relative_to(resolved_base_dir) + except ValueError: + return Response().error("Invalid file path").__dict__ + + filename_ext = os.path.splitext(filename)[1].lower() + if filename_ext == ".wav": + return await send_file(str(resolved_file_path), mimetype="audio/wav") + if filename_ext[1:] in self.supported_imgs: + return await send_file(str(resolved_file_path), mimetype="image/jpeg") + return await send_file(str(resolved_file_path)) + + except (FileNotFoundError, OSError): + return Response().error("File access error").__dict__ + + async def get_attachment(self): + """Get attachment file by attachment_id.""" + attachment_id = request.args.get("attachment_id") + if not attachment_id: + return Response().error("Missing key: attachment_id").__dict__ + + try: + attachment = await self.db.get_attachment_by_id(attachment_id) + if not attachment: + return Response().error("Attachment not found").__dict__ + + file_path = attachment.path + resolved_file_path = _resolve_path(file_path) + + return await send_file( + str(resolved_file_path), mimetype=attachment.mime_type + ) + + except (FileNotFoundError, OSError): + return Response().error("File access error").__dict__ + + async def post_file(self): + """Upload a file and create an attachment record, return attachment_id.""" + post_data = await request.files + if "file" not in post_data: + return Response().error("Missing key: file").__dict__ + + file = post_data["file"] + filename = file.filename or f"{uuid.uuid4()!s}" + content_type = file.content_type or "application/octet-stream" + + if content_type.startswith("image"): + attach_type = "image" + elif content_type.startswith("audio"): + attach_type = "record" + elif content_type.startswith("video"): + attach_type = "video" + else: + attach_type = "file" + + path = os.path.join(self.attachments_dir, filename) + await file.save(path) + + attachment = await self.db.insert_attachment( + path=path, + type=attach_type, + mime_type=content_type, + ) + + if not attachment: + return Response().error("Failed to create attachment").__dict__ + + filename = os.path.basename(attachment.path) + + return ( + Response() + .ok( + data={ + "attachment_id": attachment.attachment_id, + "filename": filename, + "type": attach_type, + } + ) + .__dict__ + ) + + async def _build_user_message_parts(self, message: str | list) -> list[dict]: + """Build user message parts list.""" + return await build_webchat_message_parts( + message, + get_attachment_by_id=self.db.get_attachment_by_id, + strict=False, + ) + + async def _create_attachment_from_file( + self, filename: str, attach_type: str + ) -> dict | None: + """Create attachment from local file and return message part.""" + return await create_attachment_part_from_existing_file( + filename, + attach_type=attach_type, + insert_attachment=self.db.insert_attachment, + attachments_dir=self.attachments_dir, + ) + + async def _save_bot_message( + self, + tui_conv_id: str, + text: str, + media_parts: list, + reasoning: str, + agent_stats: dict, + refs: dict, + ): + """Save bot message to history, return saved record.""" + bot_message_parts = [] + bot_message_parts.extend(media_parts) + if text: + bot_message_parts.append({"type": "plain", "text": text}) + + new_his = {"type": "bot", "message": bot_message_parts} + if reasoning: + new_his["reasoning"] = reasoning + if agent_stats: + new_his["agent_stats"] = agent_stats + if refs: + new_his["refs"] = refs + + record = await self.platform_history_mgr.insert( + platform_id="tui", + user_id=tui_conv_id, + content=new_his, + sender_id="bot", + sender_name="bot", + ) + return record + + async def chat(self, post_data: dict | None = None): + username = g.get("username", "guest") + + if post_data is None: + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + if "message" not in post_data and "files" not in post_data: + return Response().error("Missing key: message or files").__dict__ + + if "session_id" not in post_data and "conversation_id" not in post_data: + return ( + Response().error("Missing key: session_id or conversation_id").__dict__ + ) + + message = post_data["message"] + session_id = post_data.get("session_id", post_data.get("conversation_id")) + selected_provider = post_data.get("selected_provider") + selected_model = post_data.get("selected_model") + enable_streaming = post_data.get("enable_streaming", True) + + if not session_id: + return Response().error("session_id is empty").__dict__ + + tui_conv_id = session_id + + message_parts = await self._build_user_message_parts(message) + if not webchat_message_parts_have_content(message_parts): + return ( + Response() + .error("Message content is empty (reply only is not allowed)") + .__dict__ + ) + + message_id = str(uuid.uuid4()) + back_queue = tui_queue_mgr.get_or_create_back_queue( + message_id, + tui_conv_id, + ) + + async def stream(): + client_disconnected = False + accumulated_parts = [] + accumulated_text = "" + accumulated_reasoning = "" + tool_calls = {} + agent_stats = {} + refs = {} + try: + session_info = { + "type": "session_id", + "data": None, + "session_id": tui_conv_id, + } + yield f"data: {json.dumps(session_info, ensure_ascii=False)}\n\n" + + async with track_conversation(self.running_convs, tui_conv_id): + while True: + result, should_break = await _poll_tui_stream_result( + back_queue, username + ) + if should_break: + client_disconnected = True + break + if not result: + continue + + if ( + "message_id" in result + and result["message_id"] != message_id + ): + logger.warning("TUI stream message_id mismatch") + continue + + result_text = result["data"] + msg_type = result.get("type") + streaming = result.get("streaming", False) + chain_type = result.get("chain_type") + + if chain_type == "agent_stats": + stats_info = { + "type": "agent_stats", + "data": json.loads(result_text), + } + yield f"data: {json.dumps(stats_info, ensure_ascii=False)}\n\n" + agent_stats = stats_info["data"] + continue + + try: + if not client_disconnected: + yield f"data: {json.dumps(result, ensure_ascii=False)}\n\n" + except Exception as e: + if not client_disconnected: + logger.debug(f"[TUI] User {username} disconnected. {e}") + client_disconnected = True + + try: + if not client_disconnected: + await asyncio.sleep(0.05) + except asyncio.CancelledError: + logger.debug(f"[TUI] User {username} disconnected.") + client_disconnected = True + + if msg_type == "plain": + chain_type = result.get("chain_type") + if chain_type == "tool_call": + tool_call = json.loads(result_text) + tool_calls[tool_call.get("id")] = tool_call + if accumulated_text: + accumulated_parts.append( + {"type": "plain", "text": accumulated_text} + ) + accumulated_text = "" + elif chain_type == "tool_call_result": + tcr = json.loads(result_text) + tc_id = tcr.get("id") + if tc_id in tool_calls: + tool_calls[tc_id]["result"] = tcr.get("result") + tool_calls[tc_id]["finished_ts"] = tcr.get("ts") + accumulated_parts.append( + { + "type": "tool_call", + "tool_calls": [tool_calls[tc_id]], + } + ) + tool_calls.pop(tc_id, None) + elif chain_type == "reasoning": + accumulated_reasoning += result_text + elif streaming: + accumulated_text += result_text + else: + accumulated_text = result_text + elif msg_type == "image": + filename = result_text.replace("[IMAGE]", "") + part = await self._create_attachment_from_file( + filename, "image" + ) + if part: + accumulated_parts.append(part) + elif msg_type == "record": + filename = result_text.replace("[RECORD]", "") + part = await self._create_attachment_from_file( + filename, "record" + ) + if part: + accumulated_parts.append(part) + elif msg_type == "file": + filename = result_text.replace("[FILE]", "") + part = await self._create_attachment_from_file( + filename, "file" + ) + if part: + accumulated_parts.append(part) + + if msg_type == "end": + break + elif (streaming and msg_type == "complete") or not streaming: + if ( + chain_type == "tool_call" + or chain_type == "tool_call_result" + ): + continue + + saved_record = await self._save_bot_message( + tui_conv_id, + accumulated_text, + accumulated_parts, + accumulated_reasoning, + agent_stats, + refs, + ) + if saved_record and not client_disconnected: + saved_info = { + "type": "message_saved", + "data": { + "id": saved_record.id, + "created_at": to_utc_isoformat( + saved_record.created_at + ), + }, + } + try: + yield f"data: {json.dumps(saved_info, ensure_ascii=False)}\n\n" + except Exception: + pass + accumulated_parts = [] + accumulated_text = "" + accumulated_reasoning = "" + agent_stats = {} + refs = {} + except BaseException as e: + logger.exception(f"TUI stream unexpected error: {e}", exc_info=True) + finally: + tui_queue_mgr.remove_back_queue(message_id) + + chat_queue = tui_queue_mgr.get_or_create_queue(tui_conv_id) + await chat_queue.put( + ( + username, + tui_conv_id, + { + "message": message_parts, + "selected_provider": selected_provider, + "selected_model": selected_model, + "enable_streaming": enable_streaming, + "message_id": message_id, + }, + ), + ) + + message_parts_for_storage = strip_message_parts_path_fields(message_parts) + + await self.platform_history_mgr.insert( + platform_id="tui", + user_id=tui_conv_id, + content={"type": "user", "message": message_parts_for_storage}, + sender_id=username, + sender_name=username, + ) + + 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 + + async def stop_session(self): + """Stop active agent runs for a session.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + + session_id = post_data.get("session_id") + if not session_id: + return Response().error("Missing key: session_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__ + + message_type = ( + MessageType.GROUP_MESSAGE.value + if session.is_group + else MessageType.FRIEND_MESSAGE.value + ) + umo = ( + f"{session.platform_id}:{message_type}:" + f"{session.platform_id}!{username}!{session_id}" + ) + stopped_count = active_event_registry.request_agent_stop_all(umo) + + return Response().ok(data={"stopped_count": stopped_count}).__dict__ + + async def _delete_session_internal(self, session, username: str) -> None: + """Delete a single session and all its related data.""" + session_id = session.session_id + + message_type = "GroupMessage" if session.is_group else "FriendMessage" + unified_msg_origin = f"{session.platform_id}:{message_type}:{session.platform_id}!{username}!{session_id}" + await self.conv_mgr.delete_conversations_by_user_id(unified_msg_origin) + + history_list = await self.platform_history_mgr.get( + platform_id=session.platform_id, + user_id=session_id, + page=1, + page_size=100000, + ) + attachment_ids = self._extract_attachment_ids(history_list) + if attachment_ids: + await self._delete_attachments(attachment_ids) + + await self.platform_history_mgr.delete( + platform_id=session.platform_id, + user_id=session_id, + offset_sec=99999999, + ) + + try: + await self.umop_config_router.delete_route(unified_msg_origin) + except ValueError: + logger.warning( + "Failed to delete UMO route %s during session cleanup.", + unified_msg_origin, + ) + + if session.platform_id == "tui": + tui_queue_mgr.remove_queues(session_id) + + await self.db.delete_platform_session(session_id) + + async def delete_tui_session(self): + """Delete a Platform session and all its related data.""" + session_id = request.args.get("session_id") + if not session_id: + return Response().error("Missing key: session_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__ + + await self._delete_session_internal(session, username) + + return Response().ok().__dict__ + + async def batch_delete_sessions(self): + """Batch delete multiple Platform sessions.""" + post_data = await request.json + if post_data is None: + return Response().error("Missing JSON body").__dict__ + if not isinstance(post_data, dict): + return Response().error("Invalid JSON body: expected object").__dict__ + + session_ids = post_data.get("session_ids") + if not session_ids or not isinstance(session_ids, list): + return Response().error("Missing or invalid key: session_ids").__dict__ + + username = g.get("username", "guest") + sessions = await self.db.get_platform_sessions_by_ids(session_ids) + sessions_by_id = {session.session_id: session for session in sessions} + deleted_count = 0 + failed_items = [] + + for sid in session_ids: + session = sessions_by_id.get(sid) + if not session: + failed_items.append({"session_id": sid, "reason": "not found"}) + continue + if session.creator != username: + failed_items.append({"session_id": sid, "reason": "permission denied"}) + continue + + try: + await self._delete_session_internal(session, username) + deleted_count += 1 + sessions_by_id.pop(sid, None) + except Exception: + logger.warning("Failed to delete session %s", sid) + failed_items.append({"session_id": sid, "reason": "internal_error"}) + + return ( + Response() + .ok( + data={ + "deleted_count": deleted_count, + "failed_count": len(failed_items), + "failed_items": failed_items, + } + ) + .__dict__ + ) + + def _extract_attachment_ids(self, history_list) -> list[str]: + """Extract all attachment_ids from message history.""" + attachment_ids = [] + for history in history_list: + content = history.content + if not content or "message" not in content: + continue + message_parts = content.get("message", []) + for part in message_parts: + if isinstance(part, dict) and "attachment_id" in part: + attachment_ids.append(part["attachment_id"]) + return attachment_ids + + async def _delete_attachments(self, attachment_ids: list[str]) -> None: + """Delete attachments including DB records and disk files.""" + try: + attachments = await self.db.get_attachments(attachment_ids) + for attachment in attachments: + if not await anyio.Path(attachment.path).exists(): + continue + try: + await anyio.Path(attachment.path).unlink() + except OSError as e: + logger.warning( + f"Failed to delete attachment file {attachment.path}: {e}" + ) + except Exception as e: + logger.warning(f"Failed to get attachments: {e}") + + try: + await self.db.delete_attachments(attachment_ids) + except Exception as e: + logger.warning(f"Failed to delete attachments: {e}") + + async def new_session(self): + """Create a new Platform session for TUI.""" + username = g.get("username", "guest") + + session = await self.db.create_platform_session( + creator=username, + platform_id="tui", + is_group=0, + ) + + return ( + Response() + .ok( + data={ + "session_id": session.session_id, + "platform_id": session.platform_id, + } + ) + .__dict__ + ) + + async def get_sessions(self): + """Get all Platform sessions for the current user filtered by TUI platform.""" + username = g.get("username", "guest") + + platform_id = request.args.get("platform_id", "tui") + + sessions, _ = await self.db.get_platform_sessions_by_creator_paginated( + creator=username, + platform_id=platform_id, + page=1, + page_size=100, + exclude_project_sessions=True, + ) + + sessions_data = [] + for item in sessions: + session = item["session"] + + sessions_data.append( + { + "session_id": session.session_id, + "platform_id": session.platform_id, + "creator": session.creator, + "display_name": session.display_name, + "is_group": session.is_group, + "created_at": to_utc_isoformat(session.created_at), + "updated_at": to_utc_isoformat(session.updated_at), + } + ) + + return Response().ok(data=sessions_data).__dict__ + + async def get_session(self): + """Get session information and message history by session_id.""" + session_id = request.args.get("session_id") + if not session_id: + return Response().error("Missing key: session_id").__dict__ + + session = await self.db.get_platform_session_by_id(session_id) + platform_id = session.platform_id if session else "tui" + + username = g.get("username", "guest") + project_info = await self.db.get_project_by_session( + session_id=session_id, creator=username + ) + + history_ls = await self.platform_history_mgr.get( + platform_id=platform_id, + user_id=session_id, + page=1, + page_size=1000, + ) + + history_res = [history.model_dump() for history in history_ls] + + response_data = { + "history": history_res, + "is_running": self.running_convs.get(session_id, False), + } + + if project_info: + response_data["project"] = { + "project_id": project_info.project_id, + "title": project_info.title, + "emoji": project_info.emoji, + } + + return Response().ok(data=response_data).__dict__ + + async def update_session_display_name(self): + """Update a Platform session's display name.""" + post_data = await request.json + + session_id = post_data.get("session_id") + display_name = post_data.get("display_name") + + if not session_id: + return Response().error("Missing key: session_id").__dict__ + if display_name is None: + return Response().error("Missing key: display_name").__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__ + + await self.db.update_platform_session( + session_id=session_id, + display_name=display_name, + ) + + return Response().ok().__dict__ diff --git a/astrbot/dashboard/server.py b/astrbot/dashboard/server.py index f01a04479..7e87924e2 100644 --- a/astrbot/dashboard/server.py +++ b/astrbot/dashboard/server.py @@ -57,6 +57,7 @@ from .routes import ( SubAgentRoute, T2iRoute, ToolsRoute, + TUIChatRoute, UpdateRoute, ) from .routes.api_key import ALL_OPEN_API_SCOPES @@ -299,6 +300,7 @@ class AstrBotDashboard: self.platform_route = PlatformRoute(self.context, self.core_lifecycle) self.backup_route = BackupRoute(self.context, db, self.core_lifecycle) self.live_chat_route = LiveChatRoute(self.context, db, self.core_lifecycle) + self.tui_chat_route = TUIChatRoute(self.context, db, self.core_lifecycle) self.app.add_url_rule( "/api/plug/",