AstrBot 3.5.6__py3-none-any.whl → 4.7.0__py3-none-any.whl

astrbot/dashboard/routes/chat.py

@@ -1,13 +1,28 @@
-import uuid
+import asyncio
 import json
 import os
-from .route import Route, Response, RouteContext
-from astrbot.core import web_chat_queue, web_chat_back_queue
-from quart import request, Response as QuartResponse, g, make_response
-from astrbot.core.db import BaseDatabase
-import asyncio
+import uuid
+from contextlib import asynccontextmanager
+
+from quart import Response as QuartResponse
+from quart import g, make_response, request
+
 from astrbot.core import logger
 from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
+from astrbot.core.db import BaseDatabase
+from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queue_mgr
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+
+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)
 
 
 class ChatRoute(Route):
@@ -20,38 +35,30 @@ class ChatRoute(Route):
         super().__init__(context)
         self.routes = {
             "/chat/send": ("POST", self.chat),
-            "/chat/listen": ("GET", self.listener),
-            "/chat/new_conversation": ("GET", self.new_conversation),
-            "/chat/conversations": ("GET", self.get_conversations),
-            "/chat/get_conversation": ("GET", self.get_conversation),
-            "/chat/delete_conversation": ("GET", self.delete_conversation),
+            "/chat/new_session": ("GET", self.new_session),
+            "/chat/sessions": ("GET", self.get_sessions),
+            "/chat/get_session": ("GET", self.get_session),
+            "/chat/delete_session": ("GET", self.delete_webchat_session),
+            "/chat/update_session_display_name": (
+                "POST",
+                self.update_session_display_name,
+            ),
             "/chat/get_file": ("GET", self.get_file),
             "/chat/post_image": ("POST", self.post_image),
             "/chat/post_file": ("POST", self.post_file),
-            "/chat/status": ("GET", self.status),
         }
-        self.db = db
         self.core_lifecycle = core_lifecycle
         self.register_routes()
-        self.imgs_dir = "data/webchat/imgs"
+        self.imgs_dir = os.path.join(get_astrbot_data_path(), "webchat", "imgs")
+        os.makedirs(self.imgs_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.curr_user_cid = {}
-        self.curr_chat_sse = {}
-
-    async def status(self):
-        has_llm_enabled = (
-            self.core_lifecycle.provider_manager.curr_provider_inst is not None
-        )
-        has_stt_enabled = (
-            self.core_lifecycle.provider_manager.curr_stt_provider_inst is not None
-        )
-        return (
-            Response()
-            .ok(data={"llm_enabled": has_llm_enabled, "stt_enabled": has_stt_enabled})
-            .__dict__
-        )
+        self.running_convs: dict[str, bool] = {}
 
     async def get_file(self):
         filename = request.args.get("filename")
@@ -59,16 +66,24 @@ class ChatRoute(Route):
             return Response().error("Missing key: filename").__dict__
 
         try:
-            with open(os.path.join(self.imgs_dir, filename), "rb") as f:
-                if filename.endswith(".wav"):
+            file_path = os.path.join(self.imgs_dir, os.path.basename(filename))
+            real_file_path = os.path.realpath(file_path)
+            real_imgs_dir = os.path.realpath(self.imgs_dir)
+
+            if not real_file_path.startswith(real_imgs_dir):
+                return Response().error("Invalid file path").__dict__
+
+            with open(real_file_path, "rb") as f:
+                filename_ext = os.path.splitext(filename)[1].lower()
+
+                if filename_ext == ".wav":
                     return QuartResponse(f.read(), mimetype="audio/wav")
-                elif filename.split(".")[-1] in self.supported_imgs:
+                if filename_ext[1:] in self.supported_imgs:
                     return QuartResponse(f.read(), mimetype="image/jpeg")
-                else:
-                    return QuartResponse(f.read())
+                return QuartResponse(f.read())
 
-        except FileNotFoundError:
-            return Response().error("File not found").__dict__
+        except (FileNotFoundError, OSError):
+            return Response().error("File access error").__dict__
 
     async def post_image(self):
         post_data = await request.files
@@ -88,8 +103,7 @@ class ChatRoute(Route):
             return Response().error("Missing key: file").__dict__
 
         file = post_data["file"]
-        filename = f"{str(uuid.uuid4())}"
-        print(file)
+        filename = f"{uuid.uuid4()!s}"
         # 通过文件格式判断文件类型
         if file.content_type.startswith("audio"):
             filename += ".wav"
@@ -106,112 +120,125 @@ class ChatRoute(Route):
         if "message" not in post_data and "image_url" not in post_data:
             return Response().error("Missing key: message or image_url").__dict__
 
-        if "conversation_id" not in post_data:
-            return Response().error("Missing key: conversation_id").__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"]
-        conversation_id = post_data["conversation_id"]
+        # conversation_id = post_data["conversation_id"]
+        session_id = post_data.get("session_id", post_data.get("conversation_id"))
         image_url = post_data.get("image_url")
         audio_url = post_data.get("audio_url")
+        selected_provider = post_data.get("selected_provider")
+        selected_model = post_data.get("selected_model")
+        enable_streaming = post_data.get("enable_streaming", True)  # 默认为 True
+
         if not message and not image_url and not audio_url:
             return (
                 Response()
                 .error("Message and image_url and audio_url are empty")
                 .__dict__
             )
-        if not conversation_id:
-            return Response().error("conversation_id is empty").__dict__
+        if not session_id:
+            return Response().error("session_id is empty").__dict__
 
-        self.curr_user_cid[username] = conversation_id
+        # 追加用户消息
+        webchat_conv_id = session_id
 
-        await web_chat_queue.put(
-            (
-                username,
-                conversation_id,
-                {
-                    "message": message,
-                    "image_url": image_url,  # list
-                    "audio_url": audio_url,
-                },
-            )
-        )
+        # 获取会话特定的队列
+        back_queue = webchat_queue_mgr.get_or_create_back_queue(webchat_conv_id)
 
-        # 持久化
-        conversation = self.db.get_conversation_by_user_id(username, conversation_id)
-        try:
-            history = json.loads(conversation.history)
-        except BaseException as e:
-            print(e)
-            history = []
         new_his = {"type": "user", "message": message}
         if image_url:
             new_his["image_url"] = image_url
         if audio_url:
             new_his["audio_url"] = audio_url
-        history.append(new_his)
-        self.db.update_conversation(
-            username, conversation_id, history=json.dumps(history)
+        await self.platform_history_mgr.insert(
+            platform_id="webchat",
+            user_id=webchat_conv_id,
+            content=new_his,
+            sender_id=username,
+            sender_name=username,
         )
 
-        return Response().ok().__dict__
-
-    async def listener(self):
-        """一直保持长连接"""
-
-        username = g.get("username", "guest")
-
-        if username in self.curr_chat_sse:
-            return Response().error("Already connected").__dict__
+        async def stream():
+            client_disconnected = False
 
-        self.curr_chat_sse[username] = None
+            try:
+                async with track_conversation(self.running_convs, webchat_conv_id):
+                    while True:
+                        try:
+                            result = await asyncio.wait_for(back_queue.get(), timeout=1)
+                        except asyncio.TimeoutError:
+                            continue
+                        except asyncio.CancelledError:
+                            logger.debug(f"[WebChat] 用户 {username} 断开聊天长连接。")
+                            client_disconnected = True
+                        except Exception as e:
+                            logger.error(f"WebChat stream error: {e}")
+
+                        if not result:
+                            continue
+
+                        result_text = result["data"]
+                        type = result.get("type")
+                        streaming = result.get("streaming", False)
 
-        heartbeat = json.dumps({"type": "heartbeat", "data": "ping"})
+                        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"[WebChat] 用户 {username} 断开聊天长连接。 {e}",
+                                )
+                            client_disconnected = True
 
-        async def stream():
-            try:
-                yield f"data: {heartbeat}\n\n"  # 心跳包
-                while True:
-                    try:
-                        result = await asyncio.wait_for(
-                            web_chat_back_queue.get(), timeout=10
-                        )  # 设置超时时间为5秒
-                    except asyncio.TimeoutError:
-                        yield f"data: {heartbeat}\n\n"  # 心跳包
-                        continue
-
-                    if not result:
-                        continue
-
-                    result_text = result["data"]
-                    type = result.get("type")
-                    cid = result.get("cid")
-                    streaming = result.get("streaming", False)
-                    if cid != self.curr_user_cid.get(username):
-                        # 丢弃
-                        continue
-                    yield f"data: {json.dumps(result, ensure_ascii=False)}\n\n"
-                    await asyncio.sleep(0.05)
-
-                    if streaming and type != "end":
-                        continue
-
-                    if result_text:
-                        conversation = self.db.get_conversation_by_user_id(
-                            username, cid
-                        )
                         try:
-                            history = json.loads(conversation.history)
-                        except BaseException as e:
-                            print(e)
-                            history = []
-                        history.append({"type": "bot", "message": result_text})
-                        self.db.update_conversation(
-                            username, cid, history=json.dumps(history)
-                        )
-            except BaseException as _:
-                logger.debug(f"用户 {username} 断开聊天长连接。")
-                self.curr_chat_sse.pop(username)
-                return
+                            if not client_disconnected:
+                                await asyncio.sleep(0.05)
+                        except asyncio.CancelledError:
+                            logger.debug(f"[WebChat] 用户 {username} 断开聊天长连接。")
+                            client_disconnected = True
+
+                        if type == "end":
+                            break
+                        elif (
+                            (streaming and type == "complete")
+                            or not streaming
+                            or type == "break"
+                        ):
+                            # 追加机器人消息
+                            new_his = {"type": "bot", "message": result_text}
+                            if "reasoning" in result:
+                                new_his["reasoning"] = result["reasoning"]
+                            await self.platform_history_mgr.insert(
+                                platform_id="webchat",
+                                user_id=webchat_conv_id,
+                                content=new_his,
+                                sender_id="bot",
+                                sender_name="bot",
+                            )
+            except BaseException as e:
+                logger.exception(f"WebChat stream unexpected error: {e}", exc_info=True)
+
+        # 将消息放入会话特定的队列
+        chat_queue = webchat_queue_mgr.get_or_create_queue(webchat_conv_id)
+        await chat_queue.put(
+            (
+                username,
+                webchat_conv_id,
+                {
+                    "message": message,
+                    "image_url": image_url,  # list
+                    "audio_url": audio_url,
+                    "selected_provider": selected_provider,
+                    "selected_model": selected_model,
+                    "enable_streaming": enable_streaming,
+                },
+            ),
+        )
 
         response = await make_response(
             stream(),
@@ -222,37 +249,166 @@ class ChatRoute(Route):
                 "Connection": "keep-alive",
             },
         )
-        response.timeout = None
+        response.timeout = None  # fix SSE auto disconnect issue
         return response
 
-    async def delete_conversation(self):
+    async def delete_webchat_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")
-        conversation_id = request.args.get("conversation_id")
-        if not conversation_id:
-            return Response().error("Missing key: conversation_id").__dict__
 
-        self.db.delete_conversation(username, conversation_id)
+        # 验证会话是否存在且属于当前用户
+        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 = "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)
+
+        # 删除消息历史
+        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 as exc:
+            logger.warning(
+                "Failed to delete UMO route %s during session cleanup: %s",
+                unified_msg_origin,
+                exc,
+            )
+
+        # 清理队列（仅对 webchat）
+        if session.platform_id == "webchat":
+            webchat_queue_mgr.remove_queues(session_id)
+
+        # 删除会话
+        await self.db.delete_platform_session(session_id)
+
         return Response().ok().__dict__
 
-    async def new_conversation(self):
+    async def new_session(self):
+        """Create a new Platform session (default: webchat)."""
         username = g.get("username", "guest")
-        conversation_id = str(uuid.uuid4())
-        self.db.new_conversation(username, conversation_id)
-        return Response().ok(data={"conversation_id": conversation_id}).__dict__
 
-    async def get_conversations(self):
-        username = g.get("username", "guest")
-        conversations = self.db.get_conversations(username)
-        return Response().ok(data=conversations).__dict__
+        # 获取可选的 platform_id 参数，默认为 webchat
+        platform_id = request.args.get("platform_id", "webchat")
+
+        # 创建新会话
+        session = await self.db.create_platform_session(
+            creator=username,
+            platform_id=platform_id,
+            is_group=0,
+        )
+
+        return (
+            Response()
+            .ok(
+                data={
+                    "session_id": session.session_id,
+                    "platform_id": session.platform_id,
+                }
+            )
+            .__dict__
+        )
 
-    async def get_conversation(self):
+    async def get_sessions(self):
+        """Get all Platform sessions for the current user."""
         username = g.get("username", "guest")
-        conversation_id = request.args.get("conversation_id")
-        if not conversation_id:
-            return Response().error("Missing key: conversation_id").__dict__
 
-        conversation = self.db.get_conversation_by_user_id(username, conversation_id)
+        # 获取可选的 platform_id 参数
+        platform_id = request.args.get("platform_id")
+
+        sessions = await self.db.get_platform_sessions_by_creator(
+            creator=username,
+            platform_id=platform_id,
+            page=1,
+            page_size=100,  # 暂时返回前100个
+        )
+
+        # 转换为字典格式，并添加额外信息
+        sessions_data = []
+        for session in sessions:
+            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": session.created_at.astimezone().isoformat(),
+                    "updated_at": session.updated_at.astimezone().isoformat(),
+                }
+            )
+
+        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__
+
+        # 获取会话信息以确定 platform_id
+        session = await self.db.get_platform_session_by_id(session_id)
+        platform_id = session.platform_id if session else "webchat"
+
+        # Get platform message history using session_id
+        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]
+
+        return (
+            Response()
+            .ok(
+                data={
+                    "history": history_res,
+                    "is_running": self.running_convs.get(session_id, False),
+                },
+            )
+            .__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")
 
-        self.curr_user_cid[username] = conversation_id
+        # 验证会话是否存在且属于当前用户
+        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__
+
+        # 更新 display_name
+        await self.db.update_platform_session(
+            session_id=session_id,
+            display_name=display_name,
+        )
 
-        return Response().ok(data=conversation).__dict__
+        return Response().ok().__dict__