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

astrbot/core/platform/sources/webchat/webchat_adapter.py

@@ -1,53 +1,97 @@
-import time
 import asyncio
-import uuid
 import os
-from typing import Awaitable, Any
+import time
+import uuid
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from astrbot import logger
+from astrbot.core.message.components import Image, Plain, Record
+from astrbot.core.message.message_event_result import MessageChain
 from astrbot.core.platform import (
-    Platform,
     AstrBotMessage,
     MessageMember,
     MessageType,
+    Platform,
     PlatformMetadata,
 )
-from astrbot.core.message.message_event_result import MessageChain
-from astrbot.core.message.components import Plain, Image, Record  # noqa: F403
-from astrbot import logger
-from astrbot.core import web_chat_queue
-from .webchat_event import WebChatMessageEvent
 from astrbot.core.platform.astr_message_event import MessageSesion
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
+
 from ...register import register_platform_adapter
+from .webchat_event import WebChatMessageEvent
+from .webchat_queue_mgr import WebChatQueueMgr, webchat_queue_mgr
 
 
 class QueueListener:
-    def __init__(self, queue: asyncio.Queue, callback: callable) -> None:
-        self.queue = queue
+    def __init__(self, webchat_queue_mgr: WebChatQueueMgr, callback: Callable) -> None:
+        self.webchat_queue_mgr = webchat_queue_mgr
         self.callback = callback
+        self.running_tasks = set()
+
+    async def listen_to_queue(self, conversation_id: str):
+        """Listen to a specific conversation queue"""
+        queue = self.webchat_queue_mgr.get_or_create_queue(conversation_id)
+        while True:
+            try:
+                data = await queue.get()
+                await self.callback(data)
+            except Exception as e:
+                logger.error(
+                    f"Error processing message from conversation {conversation_id}: {e}",
+                )
+                break
 
     async def run(self):
+        """Monitor for new conversation queues and start listeners"""
+        monitored_conversations = set()
+
         while True:
-            data = await self.queue.get()
-            await self.callback(data)
+            # Check for new conversations
+            current_conversations = set(self.webchat_queue_mgr.queues.keys())
+            new_conversations = current_conversations - monitored_conversations
+
+            # Start listeners for new conversations
+            for conversation_id in new_conversations:
+                task = asyncio.create_task(self.listen_to_queue(conversation_id))
+                self.running_tasks.add(task)
+                task.add_done_callback(self.running_tasks.discard)
+                monitored_conversations.add(conversation_id)
+                logger.debug(f"Started listener for conversation: {conversation_id}")
+
+            # Clean up monitored conversations that no longer exist
+            removed_conversations = monitored_conversations - current_conversations
+            monitored_conversations -= removed_conversations
+
+            await asyncio.sleep(1)  # Check for new conversations every second
 
 
 @register_platform_adapter("webchat", "webchat")
 class WebChatAdapter(Platform):
     def __init__(
-        self, platform_config: dict, platform_settings: dict, event_queue: asyncio.Queue
+        self,
+        platform_config: dict,
+        platform_settings: dict,
+        event_queue: asyncio.Queue,
     ) -> None:
         super().__init__(event_queue)
 
         self.config = platform_config
         self.settings = platform_settings
         self.unique_session = platform_settings["unique_session"]
-        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.metadata = PlatformMetadata(
-            name="webchat", description="webchat", id=self.config.get("id")
+            name="webchat",
+            description="webchat",
+            id="webchat",
         )
 
     async def send_by_session(
-        self, session: MessageSesion, message_chain: MessageChain
+        self,
+        session: MessageSesion,
+        message_chain: MessageChain,
     ):
         await WebChatMessageEvent._send(message_chain, session.session_id)
         await super().send_by_session(session, message_chain)
@@ -57,7 +101,6 @@ class WebChatAdapter(Platform):
 
         abm = AstrBotMessage()
         abm.self_id = "webchat"
-        abm.tag = "webchat"
         abm.sender = MessageMember(username, username)
 
         abm.type = MessageType.FRIEND_MESSAGE
@@ -73,13 +116,13 @@ class WebChatAdapter(Platform):
             if isinstance(payload["image_url"], list):
                 for img in payload["image_url"]:
                     abm.message.append(
-                        Image.fromFileSystem(os.path.join(self.imgs_dir, img))
+                        Image.fromFileSystem(os.path.join(self.imgs_dir, img)),
                     )
             else:
                 abm.message.append(
                     Image.fromFileSystem(
-                        os.path.join(self.imgs_dir, payload["image_url"])
-                    )
+                        os.path.join(self.imgs_dir, payload["image_url"]),
+                    ),
                 )
         if payload["audio_url"]:
             if isinstance(payload["audio_url"], list):
@@ -103,7 +146,7 @@ class WebChatAdapter(Platform):
             abm = await self.convert_message(data)
             await self.handle_msg(abm)
 
-        bot = QueueListener(web_chat_queue, callback)
+        bot = QueueListener(webchat_queue_mgr, callback)
         return bot.run()
 
     def meta(self) -> PlatformMetadata:
@@ -117,6 +160,13 @@ class WebChatAdapter(Platform):
             session_id=message.session_id,
         )
 
+        _, _, payload = message.raw_message  # type: ignore
+        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)
+        )
+
         self.commit_event(message_event)
 
     async def terminate(self):

astrbot/core/platform/sources/webchat/webchat_event.py

@@ -1,13 +1,16 @@
+import base64
 import os
 import uuid
-import base64
+
 from astrbot.api import logger
 from astrbot.api.event import AstrMessageEvent, MessageChain
-from astrbot.api.message_components import Plain, Image, Record
+from astrbot.api.message_components import Image, Plain, Record
+from astrbot.core.utils.astrbot_path import get_astrbot_data_path
 from astrbot.core.utils.io import download_image_by_url
-from astrbot.core import web_chat_back_queue
 
-imgs_dir = "data/webchat/imgs"
+from .webchat_queue_mgr import webchat_queue_mgr
+
+imgs_dir = os.path.join(get_astrbot_data_path(), "webchat", "imgs")
 
 
 class WebChatMessageEvent(AstrMessageEvent):
@@ -17,13 +20,18 @@ class WebChatMessageEvent(AstrMessageEvent):
 
     @staticmethod
     async def _send(message: MessageChain, session_id: str, streaming: bool = False):
+        cid = session_id.split("!")[-1]
+        web_chat_back_queue = webchat_queue_mgr.get_or_create_back_queue(cid)
         if not message:
             await web_chat_back_queue.put(
-                {"type": "end", "data": "", "streaming": False}
+                {
+                    "type": "end",
+                    "data": "",
+                    "streaming": False,
+                },  # end means this request is finished
             )
             return ""
 
-        cid = session_id.split("!")[-1]
         data = ""
         for comp in message.chain:
             if isinstance(comp, Plain):
@@ -34,7 +42,8 @@ class WebChatMessageEvent(AstrMessageEvent):
                         "cid": cid,
                         "data": data,
                         "streaming": streaming,
-                    }
+                        "chain_type": message.type,
+                    },
                 )
             elif isinstance(comp, Image):
                 # save image to local
@@ -63,7 +72,7 @@ class WebChatMessageEvent(AstrMessageEvent):
                         "cid": cid,
                         "data": data,
                         "streaming": streaming,
-                    }
+                    },
                 )
             elif isinstance(comp, Record):
                 # save record to local
@@ -87,7 +96,7 @@ class WebChatMessageEvent(AstrMessageEvent):
                         "cid": cid,
                         "data": data,
                         "streaming": streaming,
-                    }
+                    },
                 )
             else:
                 logger.debug(f"webchat 忽略: {comp.type}")
@@ -96,29 +105,44 @@ class WebChatMessageEvent(AstrMessageEvent):
 
     async def send(self, message: MessageChain):
         await WebChatMessageEvent._send(message, session_id=self.session_id)
-        await web_chat_back_queue.put(
-            {
-                "type": "end",
-                "data": "",
-                "streaming": False,
-                "cid": self.session_id.split("!")[-1],
-            }
-        )
         await super().send(message)
 
     async def send_streaming(self, generator, use_fallback: bool = False):
         final_data = ""
+        reasoning_content = ""
+        cid = self.session_id.split("!")[-1]
+        web_chat_back_queue = webchat_queue_mgr.get_or_create_back_queue(cid)
         async for chain in generator:
-            final_data += await WebChatMessageEvent._send(
-                chain, session_id=self.session_id, streaming=True
+            if chain.type == "break" and final_data:
+                # 分割符
+                await web_chat_back_queue.put(
+                    {
+                        "type": "break",  # break means a segment end
+                        "data": final_data,
+                        "streaming": True,
+                        "cid": cid,
+                    },
+                )
+                final_data = ""
+                continue
+
+            r = await WebChatMessageEvent._send(
+                chain,
+                session_id=self.session_id,
+                streaming=True,
             )
+            if chain.type == "reasoning":
+                reasoning_content += chain.get_plain_text()
+            else:
+                final_data += r
 
         await web_chat_back_queue.put(
             {
-                "type": "end",
+                "type": "complete",  # complete means we return the final result
                 "data": final_data,
+                "reasoning": reasoning_content,
                 "streaming": True,
-                "cid": self.session_id.split("!")[-1],
-            }
+                "cid": cid,
+            },
         )
         await super().send_streaming(generator, use_fallback)

astrbot/core/platform/sources/webchat/webchat_queue_mgr.py

@@ -0,0 +1,35 @@
+import asyncio
+
+
+class WebChatQueueMgr:
+    def __init__(self) -> None:
+        self.queues = {}
+        """Conversation ID to asyncio.Queue mapping"""
+        self.back_queues = {}
+        """Conversation ID to asyncio.Queue mapping for responses"""
+
+    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()
+        return self.queues[conversation_id]
+
+    def get_or_create_back_queue(self, conversation_id: str) -> asyncio.Queue:
+        """Get or create a back queue for the given conversation ID"""
+        if conversation_id not in self.back_queues:
+            self.back_queues[conversation_id] = asyncio.Queue()
+        return self.back_queues[conversation_id]
+
+    def remove_queues(self, conversation_id: str):
+        """Remove queues for the given conversation ID"""
+        if conversation_id in self.queues:
+            del self.queues[conversation_id]
+        if conversation_id in self.back_queues:
+            del self.back_queues[conversation_id]
+
+    def has_queue(self, conversation_id: str) -> bool:
+        """Check if a queue exists for the given conversation ID"""
+        return conversation_id in self.queues
+
+
+webchat_queue_mgr = WebChatQueueMgr()