mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-18 02:00:09 +08:00
feat: add TUI platform adapter and chat route
- Add TUI platform adapter (tui_adapter, tui_event, tui_queue_mgr) - Add TUI chat route for dashboard - Add cmd_run_tui command - Fix shared_preferences TypeVar syntax - Various platform adapter and provider updates
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
307
astrbot/cli/commands/cmd_run_tui.py
Normal file
307
astrbot/cli/commands/cmd_run_tui.py
Normal file
@@ -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()
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
"""实例化一个平台"""
|
||||
# 动态导入
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
5
astrbot/core/platform/sources/tui/__init__.py
Normal file
5
astrbot/core/platform/sources/tui/__init__.py
Normal file
@@ -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"]
|
||||
239
astrbot/core/platform/sources/tui/tui_adapter.py
Normal file
239
astrbot/core/platform/sources/tui/tui_adapter.py
Normal file
@@ -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()
|
||||
203
astrbot/core/platform/sources/tui/tui_event.py
Normal file
203
astrbot/core/platform/sources/tui/tui_event.py
Normal file
@@ -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)
|
||||
164
astrbot/core/platform/sources/tui/tui_queue_mgr.py
Normal file
164
astrbot/core/platform/sources/tui/tui_queue_mgr.py
Normal file
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
755
astrbot/dashboard/routes/tui_chat.py
Normal file
755
astrbot/dashboard/routes/tui_chat.py
Normal file
@@ -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__
|
||||
@@ -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/<path:subpath>",
|
||||
|
||||
Reference in New Issue
Block a user