From fcab00332d378dcbe07976aab1cd078b45ebe563 Mon Sep 17 00:00:00 2001 From: LIghtJUNction Date: Mon, 23 Mar 2026 02:57:29 +0800 Subject: [PATCH] docs: improve module docstrings - Add module docstring to tui_async.py describing async TUI implementation - Add new message_handler.py module with shared SSE message handling --- astrbot/cli/commands/tui_async.py | 207 +++++++++++++------- astrbot/tui/message_handler.py | 310 ++++++++++++++++++++++++++++++ 2 files changed, 447 insertions(+), 70 deletions(-) create mode 100644 astrbot/tui/message_handler.py diff --git a/astrbot/cli/commands/tui_async.py b/astrbot/cli/commands/tui_async.py index 1bc1527a5..75d5c0844 100644 --- a/astrbot/cli/commands/tui_async.py +++ b/astrbot/cli/commands/tui_async.py @@ -1,4 +1,8 @@ -"""Async TUI implementation that connects to a running AstrBot instance via HTTP API.""" +"""Async TUI implementation that connects to a running AstrBot instance via HTTP API. + +This module provides a terminal UI that connects to AstrBot via the dashboard API, +supporting streaming responses and all message types. +""" from __future__ import annotations @@ -7,9 +11,16 @@ import curses import json from dataclasses import dataclass, field from enum import Enum +from typing import Any import httpx +from astrbot.tui.message_handler import ( + ChatResponse, + MessageType, + ParsedMessage, + SSEMessageParser, +) from astrbot.tui.screen import Screen @@ -17,6 +28,8 @@ class MessageSender(Enum): USER = "user" BOT = "bot" SYSTEM = "system" + TOOL = "tool" + REASONING = "reasoning" @dataclass @@ -37,7 +50,15 @@ class TUIState: class TUIClient: - """TUI client that connects to AstrBot via HTTP API.""" + """TUI client that connects to AstrBot via HTTP API. + + Supports full streaming responses including: + - Plain text (streaming) + - Tool calls and results + - Reasoning chains + - Agent stats + - Media (images, audio, files) + """ def __init__( self, @@ -64,13 +85,14 @@ class TUIClient: # Session info self.session_id: str | None = None + self.conversation_id: str | None = None # HTTP client self._client: httpx.AsyncClient | None = None self._headers: dict[str, str] = {} - # Pending tasks - self._pending_tasks: list[asyncio.Task] = [] + # SSE parser + self._parser = SSEMessageParser() async def connect(self) -> bool: """Connect to AstrBot and authenticate.""" @@ -110,11 +132,12 @@ class TUIClient: self.state.status = f"Session error: {session_data.get('msg')}" return False - self.session_id = session_data.get("data", {}).get("session_id") - if not self.session_id: + self.conversation_id = session_data.get("data", {}).get("session_id") + if not self.conversation_id: self.state.status = "No session_id in response" return False + self.session_id = self.conversation_id self.state.connected = True self.state.status = "Connected" return True @@ -133,8 +156,49 @@ class TUIClient: await self._client.aclose() self.state.connected = False + async def load_history(self) -> None: + """Load message history for the current session.""" + if not self._client or not self.conversation_id: + return + + try: + resp = await self._client.get( + "/api/chat/get_session", + params={"session_id": self.conversation_id}, + headers=self._headers, + ) + if resp.status_code != 200: + return + + data = resp.json() + history = data.get("data", {}).get("history", []) + + for record in reversed(history): + content = record.get("content", {}) + msg_type = content.get("type") + message_parts = content.get("message", []) + + if msg_type == "user": + for part in message_parts: + if part.get("type") == "plain": + self.add_message( + MessageSender.USER, part.get("text", "") + ) + elif msg_type == "bot": + for part in message_parts: + if part.get("type") == "plain": + self.add_message(MessageSender.BOT, part.get("text", "")) + + except Exception as e: + if self.debug: + import traceback + + traceback.print_exc() + def add_message(self, sender: MessageSender, text: str) -> None: """Add a message to the chat log.""" + if not text: + return self.state.messages.append(Message(sender=sender, text=text)) if len(self.state.messages) > self._max_messages: self.state.messages = self.state.messages[-self._max_messages :] @@ -184,8 +248,7 @@ class TUIClient: # Handle Enter/Return - submit message elif key in (curses.KEY_ENTER, 10, 13): if self.state.input_buffer.strip(): - task = asyncio.create_task(self._submit_message()) - self._pending_tasks.append(task) + asyncio.create_task(self._submit_message()) return True # Handle history navigation (up/down arrows) @@ -247,8 +310,8 @@ class TUIClient: await self._process_user_message(text) async def _process_user_message(self, text: str) -> None: - """Send message to AstrBot and process the response.""" - if not self.session_id or not self._client: + """Send message to AstrBot and process the streaming response.""" + if not self.conversation_id or not self._client: self.add_system_message("Not connected to AstrBot") return @@ -256,9 +319,15 @@ class TUIClient: try: # Format umo for webchat - umo = f"webchat:FriendMessage:webchat!{self.username}!{self.session_id}" + umo = f"webchat:FriendMessage:webchat!{self.username}!{self.conversation_id}" - # Send message and stream response + # Reset parser for new stream + self._parser.reset() + + # Create a placeholder bot message that we'll update + placeholder = self.add_message(MessageSender.BOT, "...") + + # Send message and stream response using proper SSE async with self._client.stream( "POST", "/api/chat/chat", @@ -266,44 +335,65 @@ class TUIClient: json={ "umo": umo, "message": text, + "session_id": self.conversation_id, "streaming": True, }, timeout=None, ) as response: if response.status_code != 200: - self.add_system_message(f"Error: HTTP {response.status_code}") + self._update_last_bot_message(f"Error: HTTP {response.status_code}") self.state.status = "Error" return - # Process streaming response - accumulated_text = "" - + # Process streaming SSE async for line in response.aiter_lines(): - if not line.startswith("data: "): + parsed = self._parser.parse_line(line) + if parsed is None: continue - data_str = line[6:] # Remove "data: " prefix - try: - data = json.loads(data_str) - except json.JSONDecodeError: - continue + update, is_complete = self._process_parsed_message(parsed) - msg_type = data.get("type") - msg_data = data.get("data", "") + # Update display based on message type + if parsed.type == MessageType.TOOL_CALL: + tool_call = json.loads(parsed.data) + self.add_message( + MessageSender.TOOL, + f"[Tool: {tool_call.get('name', 'unknown')}]", + ) + self.state.status = "Running tool..." + elif parsed.type == MessageType.TOOL_CALL_RESULT: + try: + tcr = json.loads(parsed.data) + self.add_message( + MessageSender.TOOL, + f"[Result] {tcr.get('result', '')[:100]}...", + ) + except json.JSONDecodeError: + pass + elif parsed.type == MessageType.REASONING: + self._update_last_bot_message( + f"[Thinking] {update.reasoning[-200:]}" + ) + self.state.status = "Thinking..." + elif parsed.type == MessageType.AGENT_STATS: + self.state.status = f"Tokens: {update.agent_stats.get('total_tokens', 0)}" + elif update.text: + self._update_last_bot_message(update.text) - if msg_type == "plain": - accumulated_text += msg_data - self._update_last_bot_message(accumulated_text) - elif msg_type == "image": - self._update_last_bot_message(f"[Image: {msg_data}]") - elif msg_type == "record": - self._update_last_bot_message(f"[Audio: {msg_data}]") - elif msg_type == "file": - self._update_last_bot_message(f"[File: {msg_data}]") - elif msg_type in ("complete", "end"): + if is_complete: break - self.state.status = "Ready" + # Final status + if update.reasoning: + self.add_message(MessageSender.REASONING, f"[Reasoning]\n{update.reasoning}") + + for tool_display in update.get_tool_calls_display(): + self.add_message(MessageSender.TOOL, tool_display) + + if update.error: + self.add_message(MessageSender.SYSTEM, f"Error: {update.error}") + + self.state.status = "Ready" except asyncio.CancelledError: self.state.status = "Cancelled" @@ -315,6 +405,12 @@ class TUIClient: traceback.print_exc() + def _process_parsed_message( + self, msg: ParsedMessage + ) -> tuple[ChatResponse, bool]: + """Process a parsed message and return updated response state.""" + return self._parser.process_message(msg) + def _update_last_bot_message(self, text: str) -> None: """Update the last bot message with new text (for streaming).""" for i in range(len(self.state.messages) - 1, -1, -1): @@ -351,8 +447,12 @@ class TUIClient: self.add_system_message(f"Failed to connect: {self.state.status}") else: self.add_system_message("Connected to AstrBot!") - self.add_system_message("Type your message and press Enter to send.") - self.add_system_message("Press ESC or Ctrl+C to exit.") + # Load history + await self.load_history() + + # Welcome message + self.add_system_message("Type your message and press Enter to send.") + self.add_system_message("Press ESC or Ctrl+C to exit.") # Initial render self.render() @@ -379,39 +479,6 @@ class TUIClient: await self.disconnect() -def _run_tui_curses( - screen: curses.window, - host: str, - api_key: str | None, - username: str, - password: str, - debug: bool, -) -> None: - """Curses wrapper for the async TUI.""" - screen.clear() - screen.refresh() - - scr = Screen(screen) - client = TUIClient( - screen=scr, - host=host, - api_key=api_key, - username=username, - password=password, - debug=debug, - ) - - # Create a new event loop for this thread since curses runs in its own thread - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - loop.run_until_complete(client.run_event_loop(screen)) - except KeyboardInterrupt: - pass - finally: - loop.close() - - def run_tui_async( debug: bool = False, host: str = "http://localhost:6185", diff --git a/astrbot/tui/message_handler.py b/astrbot/tui/message_handler.py new file mode 100644 index 000000000..01a6bebb4 --- /dev/null +++ b/astrbot/tui/message_handler.py @@ -0,0 +1,310 @@ +"""Shared SSE message handler for AstrBot clients (WebChat, TUI, etc). + +This module provides a unified way to parse and handle SSE messages from the +AstrBot chat API, supporting all message types including streaming responses. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + + +class MessageType(Enum): + """SSE message types from AstrBot API.""" + + SESSION_ID = "session_id" + PLAIN = "plain" + IMAGE = "image" + RECORD = "record" + FILE = "file" + TOOL_CALL = "tool_call" + TOOL_CALL_RESULT = "tool_call_result" + REASONING = "reasoning" + AGENT_STATS = "agent_stats" + AUDIO_CHUNK = "audio_chunk" + COMPLETE = "complete" + END = "end" + MESSAGE_SAVED = "message_saved" + ERROR = "error" + + +@dataclass +class ToolCall: + """Represents a tool call in progress.""" + + id: str + name: str + arguments: str | None = None + result: str | None = None + finished_ts: float | None = None + + +@dataclass +class ParsedMessage: + """A parsed SSE message with type and data.""" + + type: MessageType + data: str + raw: dict[str, Any] = field(default_factory=dict) + chain_type: str | None = None + streaming: bool = False + message_id: str | None = None + + +@dataclass +class ChatResponse: + """Complete chat response accumulated from SSE stream.""" + + text: str = "" + reasoning: str = "" + tool_calls: dict[str, ToolCall] = field(default_factory=dict) + agent_stats: dict[str, Any] = field(default_factory=dict) + refs: dict[str, Any] = field(default_factory=dict) + media_parts: list[dict[str, Any]] = field(default_factory=list) + complete: bool = False + session_id: str | None = None + saved_message_id: str | None = None + error: str | None = None + + def get_display_text(self) -> str: + """Get the main text content for display.""" + return self.text + + def get_reasoning_display(self) -> str: + """Get reasoning content formatted for display.""" + if not self.reasoning: + return "" + return f"[Reasoning]\n{self.reasoning}" + + def get_tool_calls_display(self) -> list[str]: + """Get tool calls formatted for display.""" + results = [] + for tc in self.tool_calls.values(): + if tc.result: + results.append(f"[Tool: {tc.name}]\n{tc.result}") + else: + results.append(f"[Tool: {tc.name}] (running...)") + return results + + def get_stats_display(self) -> str: + """Get agent stats formatted for display.""" + if not self.agent_stats: + return "" + parts = [] + for key, value in self.agent_stats.items(): + parts.append(f"{key}: {value}") + return " | ".join(parts) + + +class SSEMessageParser: + """Parse SSE messages from AstrBot chat API. + + Usage: + parser = SSEMessageParser() + async for msg in parser.parse_stream(response): + handle_message(msg) + """ + + def __init__(self) -> None: + self._tool_calls: dict[str, ToolCall] = {} + self._accumulated_text: str = "" + self._accumulated_reasoning: str = "" + self._accumulated_parts: list[dict[str, Any]] = [] + + def reset(self) -> None: + """Reset parser state for a new stream.""" + self._tool_calls = {} + self._accumulated_text = "" + self._accumulated_reasoning = "" + self._accumulated_parts = [] + + def parse_line(self, line: str) -> ParsedMessage | None: + """Parse a single SSE data line. + + Args: + line: A line starting with "data: " + + Returns: + ParsedMessage if valid, None if skip-worthy + """ + if not line.startswith("data: "): + return None + + data_str = line[6:] # Remove "data: " prefix + if not data_str: + return None + + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + + msg_type_str = data.get("type", "") + msg_type = self._get_message_type(msg_type_str) + msg_data = data.get("data", "") + chain_type = data.get("chain_type") + streaming = data.get("streaming", False) + message_id = data.get("message_id") + + return ParsedMessage( + type=msg_type, + data=msg_data, + raw=data, + chain_type=chain_type, + streaming=streaming, + message_id=message_id, + ) + + def _get_message_type(self, type_str: str) -> MessageType: + """Map string type to MessageType enum.""" + try: + return MessageType(type_str) + except ValueError: + return MessageType.PLAIN + + def process_message( + self, msg: ParsedMessage + ) -> tuple[ChatResponse, bool]: + """Process a parsed message and update accumulated response. + + Args: + msg: The parsed message + + Returns: + tuple of (accumulated_response, is_complete) + """ + response = ChatResponse() + + if msg.type == MessageType.SESSION_ID: + response.session_id = msg.raw.get("session_id") + return response, False + + if msg.type == MessageType.AGENT_STATS: + try: + response.agent_stats = json.loads(msg.data) + except json.JSONDecodeError: + pass + return response, False + + if msg.type == MessageType.REASONING: + self._accumulated_reasoning += msg.data + response.reasoning = self._accumulated_reasoning + return response, False + + if msg.type == MessageType.TOOL_CALL: + try: + tool_call = json.loads(msg.data) + tc = ToolCall( + id=tool_call.get("id", ""), + name=tool_call.get("name", ""), + arguments=tool_call.get("arguments"), + ) + self._tool_calls[tc.id] = tc + self._accumulated_parts.append( + {"type": "plain", "text": self._accumulated_text} + ) + self._accumulated_text = "" + except json.JSONDecodeError: + pass + response.tool_calls = self._tool_calls + return response, False + + if msg.type == MessageType.TOOL_CALL_RESULT: + try: + tcr = json.loads(msg.data) + tc_id = tcr.get("id") + if tc_id in self._tool_calls: + self._tool_calls[tc_id].result = tcr.get("result") + self._tool_calls[tc_id].finished_ts = tcr.get("ts") + self._accumulated_parts.append( + {"type": "tool_call", "tool_calls": [self._tool_calls[tc_id].__dict__]} + ) + self._tool_calls.pop(tc_id, None) + except json.JSONDecodeError: + pass + response.tool_calls = self._tool_calls + return response, False + + if msg.type == MessageType.PLAIN: + if msg.chain_type == "tool_call": + pass # Already handled above + elif msg.chain_type == "reasoning": + self._accumulated_reasoning += msg.data + response.reasoning = self._accumulated_reasoning + elif msg.streaming: + self._accumulated_text += msg.data + else: + self._accumulated_text = msg.data + response.text = self._accumulated_text + return response, False + + if msg.type == MessageType.IMAGE: + filename = msg.data.replace("[IMAGE]", "") + self._accumulated_parts.append({"type": "image", "filename": filename}) + response.media_parts = self._accumulated_parts + return response, False + + if msg.type == MessageType.RECORD: + filename = msg.data.replace("[RECORD]", "") + self._accumulated_parts.append({"type": "record", "filename": filename}) + response.media_parts = self._accumulated_parts + return response, False + + if msg.type == MessageType.FILE: + filename = msg.data.replace("[FILE]", "") + self._accumulated_parts.append({"type": "file", "filename": filename}) + response.media_parts = self._accumulated_parts + return response, False + + if msg.type == MessageType.COMPLETE: + response.text = self._accumulated_text + response.reasoning = self._accumulated_reasoning + response.tool_calls = self._tool_calls + response.complete = True + self.reset() + return response, True + + if msg.type == MessageType.END: + response.text = self._accumulated_text + response.complete = True + self.reset() + return response, True + + if msg.type == MessageType.MESSAGE_SAVED: + response.saved_message_id = msg.raw.get("data", {}).get("id") + return response, False + + return response, False + + +async def parse_sse_stream(async_iterable, callback) -> ChatResponse: + """Parse SSE stream and call callback for each message update. + + This is a convenience function for processing SSE streams. + + Args: + async_iterable: Async iterable of SSE lines (e.g., response.aiter_lines()) + callback: Async function called with (ChatResponse, is_complete) + + Returns: + Final ChatResponse when stream completes + """ + parser = SSEMessageParser() + final_response = ChatResponse() + + async for line in async_iterable: + msg = parser.parse_line(line) + if msg is None: + continue + + response, is_complete = parser.process_message(msg) + await callback(response, is_complete) + final_response = response + + if is_complete: + break + + return final_response