Merge pull request #1256 from Raven95676/better-stream

perf: 为不支持流式输出的平台提供fallback。

Author: Soulter

astrbot/core/config/default.py

@@ -53,6 +53,7 @@
         "max_context_length": -1,
         "dequeue_context_length": 1,
         "streaming_response": False,
+        "streaming_segmented": False,
     },
     "provider_stt_settings": {
         "enable": False,
@@ -1028,6 +1029,11 @@
                         "type": "bool",
                         "hint": "启用后，将会流式输出 LLM 的响应。目前仅支持 OpenAI API提供商 以及 Telegram、QQ Official 私聊 两个平台",
                     },
+                    "streaming_segmented": {
+                        "description": "不支持流式回复的平台分段输出",
+                        "type": "bool",
+                        "hint": "启用后，若平台不支持流式回复，会分段输出。目前仅支持 aiocqhttp 和 gewechat 两个平台，不支持或无需使用流式分段输出的平台会静默忽略此选项",
+                    },
                 },
             },
             "persona": {

astrbot/core/pipeline/respond/stage.py

@@ -146,9 +146,12 @@ async def process(
 
         if result.result_content_type == ResultContentType.STREAMING_RESULT:
             # 流式结果直接交付平台适配器处理
+            use_fallback = self.config.get("provider_settings", {}).get(
+                "streaming_segmented", False
+            )
             logger.info(f"应用流式输出({event.get_platform_name()})")
             await event._pre_send()
-            await event.send_streaming(result.async_stream)
+            await event.send_streaming(result.async_stream, use_fallback)
             await event._post_send()
             return
         elif len(result.chain) > 0:
@@ -159,7 +162,7 @@ async def process(
                         # 支持 File 消息段的路径映射。
                         component.file = path_Mapping(mappings, component.file)
                         event.get_result().chain[idx] = component
-            
+
             await event._pre_send()
 
             # 检查消息链是否为空

astrbot/core/platform/astr_message_event.py

@@ -1,5 +1,6 @@
 import abc
 import asyncio
+import re
 import hashlib
 import uuid
 from dataclasses import dataclass
@@ -207,9 +208,26 @@ def is_admin(self) -> bool:
         """
         return self.role == "admin"
 
-    async def send_streaming(self, generator: AsyncGenerator[MessageChain, None]):
+    async def process_buffer(self, buffer: str, pattern: re.Pattern) -> str:
+        """
+        将消息缓冲区中的文本按指定正则表达式分割后发送至消息平台，作为不支持流式输出平台的Fallback。
+        """
+        while True:
+            match = re.search(pattern, buffer)
+            if not match:
+                break
+            matched_text = match.group()
+            await self.send(MessageChain([Plain(matched_text)]))
+            buffer = buffer[match.end() :]
+            await asyncio.sleep(1.5)  # 限速
+        return buffer
+
+    async def send_streaming(
+        self, generator: AsyncGenerator[MessageChain, None], use_fallback: bool = False
+    ):
         """发送流式消息到消息平台，使用异步生成器。
         目前仅支持: telegram，qq official 私聊。
+        Fallback仅支持 aiocqhttp, gewechat。
         """
         asyncio.create_task(
             Metric.upload(msg_event_tick=1, adapter_name=self.platform_meta.name)

astrbot/core/platform/sources/aiocqhttp/aiocqhttp_message_event.py

@@ -1,9 +1,10 @@
 import asyncio
-import typing
+import re
+from typing import AsyncGenerator, Dict, List
+from aiocqhttp import CQHttp
 from astrbot.api.event import AstrMessageEvent, MessageChain
+from astrbot.api.message_components import At, Image, Node, Nodes, Plain, Record
 from astrbot.api.platform import Group, MessageMember
-from astrbot.api.message_components import Plain, Image, Record, At, Node, Nodes
-from aiocqhttp import CQHttp
 
 
 class AiocqhttpMessageEvent(AstrMessageEvent):
@@ -82,18 +83,39 @@ async def send(self, message: MessageChain):
 
         await super().send(message)
 
-    async def send_streaming(self, generator):
-        buffer = None
-        async for chain in generator:
+    async def send_streaming(
+        self, generator: AsyncGenerator, use_fallback: bool = False
+    ):
+        if not use_fallback:
+            buffer = None
+            async for chain in generator:
+                if not buffer:
+                    buffer = chain
+                else:
+                    buffer.chain.extend(chain.chain)
             if not buffer:
-                buffer = chain
-            else:
-                buffer.chain.extend(chain.chain)
-        if not buffer:
-            return
-        buffer.squash_plain()
-        await self.send(buffer)
-        return await super().send_streaming(generator)
+                return
+            buffer.squash_plain()
+            await self.send(buffer)
+            return await super().send_streaming(generator, use_fallback)
+
+        buffer = ""
+        pattern = re.compile(r"[^。？！~…]+[。？！~…]+")
+
+        async for chain in generator:
+            if isinstance(chain, MessageChain):
+                for comp in chain.chain:
+                    if isinstance(comp, Plain):
+                        buffer += comp.text
+                        if any(p in buffer for p in "。？！~…"):
+                            buffer = await self.process_buffer(buffer, pattern)
+                    else:
+                        await self.send(MessageChain(chain=[comp]))
+                        await asyncio.sleep(1.5)  # 限速
+
+        if buffer.strip():
+            await self.send(MessageChain([Plain(buffer)]))
+        return await super().send_streaming(generator, use_fallback)
 
     async def get_group(self, group_id=None, **kwargs):
         if isinstance(group_id, str) and group_id.isdigit():
@@ -108,7 +130,7 @@ async def get_group(self, group_id=None, **kwargs):
             group_id=group_id,
         )
 
-        members: typing.List[typing.Dict] = await self.bot.call_action(
+        members: List[Dict] = await self.bot.call_action(
             "get_group_member_list",
             group_id=group_id,
         )

astrbot/core/platform/sources/dingtalk/dingtalk_event.py

@@ -61,7 +61,7 @@ async def send(self, message: MessageChain):
         await self.send_with_client(self.client, message)
         await super().send(message)
 
-    async def send_streaming(self, generator):
+    async def send_streaming(self, generator, use_fallback: bool = False):
         buffer = None
         async for chain in generator:
             if not buffer:
@@ -72,4 +72,4 @@ async def send_streaming(self, generator):
             return
         buffer.squash_plain()
         await self.send(buffer)
-        return await super().send_streaming(generator)
+        return await super().send_streaming(generator, use_fallback)

astrbot/core/platform/sources/gewechat/gewechat_event.py

@@ -1,8 +1,11 @@
+import asyncio
+import re
 import wave
 import uuid
 import traceback
 import os
 
+from typing import AsyncGenerator
 from astrbot.core.utils.io import save_temp_img, download_file
 from astrbot.core.utils.tencent_record_helper import wav_to_tencent_silk
 from astrbot.api import logger
@@ -217,15 +220,36 @@ async def get_group(self, group_id=None, **kwargs):
             members=members,
         )
 
-    async def send_streaming(self, generator):
-        buffer = None
-        async for chain in generator:
+    async def send_streaming(
+        self, generator: AsyncGenerator, use_fallback: bool = False
+    ):
+        if not use_fallback:
+            buffer = None
+            async for chain in generator:
+                if not buffer:
+                    buffer = chain
+                else:
+                    buffer.chain.extend(chain.chain)
             if not buffer:
-                buffer = chain
-            else:
-                buffer.chain.extend(chain.chain)
-        if not buffer:
-            return
-        buffer.squash_plain()
-        await self.send(buffer)
-        return await super().send_streaming(generator)
+                return
+            buffer.squash_plain()
+            await self.send(buffer)
+            return await super().send_streaming(generator, use_fallback)
+
+        buffer = ""
+        pattern = re.compile(r"[^。？！~…]+[。？！~…]+")
+
+        async for chain in generator:
+            if isinstance(chain, MessageChain):
+                for comp in chain.chain:
+                    if isinstance(comp, Plain):
+                        buffer += comp.text
+                        if any(p in buffer for p in "。？！~…"):
+                            buffer = await self.process_buffer(buffer, pattern)
+                    else:
+                        await self.send(MessageChain(chain=[comp]))
+                        await asyncio.sleep(1.5)  # 限速
+
+        if buffer.strip():
+            await self.send(MessageChain([Plain(buffer)]))
+        return await super().send_streaming(generator, use_fallback)

astrbot/core/platform/sources/lark/lark_event.py

@@ -104,7 +104,7 @@ async def send(self, message: MessageChain):
 
         await super().send(message)
 
-    async def send_streaming(self, generator):
+    async def send_streaming(self, generator, use_fallback: bool = False):
         buffer = None
         async for chain in generator:
             if not buffer:
@@ -115,4 +115,4 @@ async def send_streaming(self, generator):
             return
         buffer.squash_plain()
         await self.send(buffer)
-        return await super().send_streaming(generator)
+        return await super().send_streaming(generator, use_fallback)

astrbot/core/platform/sources/qqofficial/qqofficial_message_event.py

@@ -33,7 +33,7 @@ async def send(self, message: MessageChain):
         else:
             self.send_buffer.chain.extend(message.chain)
 
-    async def send_streaming(self, generator):
+    async def send_streaming(self, generator, use_fallback: bool = False):
         """流式输出仅支持消息列表私聊"""
         stream_payload = {"state": 1, "id": None, "index": 0, "reset": False}
         last_edit_time = 0  # 上次编辑消息的时间
@@ -66,7 +66,7 @@ async def send_streaming(self, generator):
             logger.error(f"发送流式消息时出错: {e}", exc_info=True)
             self.send_buffer = None
 
-        return await super().send_streaming(generator)
+        return await super().send_streaming(generator, use_fallback)
 
     async def _post_send(self, stream: dict = None):
         if not self.send_buffer:
@@ -97,7 +97,7 @@ async def _post_send(self, stream: dict = None):
             "msg_id": self.message_obj.message_id,
         }
 
-        if not isinstance(source, (botpy.message.Message,botpy.message.DirectMessage)):
+        if not isinstance(source, (botpy.message.Message, botpy.message.DirectMessage)):
             payload["msg_seq"] = random.randint(1, 10000)
 
         match type(source):

astrbot/core/platform/sources/telegram/tg_event.py

@@ -91,7 +91,7 @@ async def send(self, message: MessageChain):
             await self.send_with_client(self.client, message, self.get_sender_id())
         await super().send(message)
 
-    async def send_streaming(self, generator):
+    async def send_streaming(self, generator, use_fallback: bool = False):
         message_thread_id = None
 
         if self.get_message_type() == MessageType.GROUP_MESSAGE:
@@ -183,16 +183,14 @@ async def send_streaming(self, generator):
                         text=markdown_text,
                         chat_id=payload["chat_id"],
                         message_id=message_id,
-                        parse_mode="MarkdownV2"
+                        parse_mode="MarkdownV2",
                     )
                 except Exception as e:
                     logger.warning(f"Markdown转换失败，使用普通文本: {e!s}")
                     await self.client.edit_message_text(
-                        text=delta,
-                        chat_id=payload["chat_id"],
-                        message_id=message_id
+                        text=delta, chat_id=payload["chat_id"], message_id=message_id
                     )
         except Exception as e:
             logger.warning(f"编辑消息失败(streaming): {e!s}")
 
-        return await super().send_streaming(generator)
+        return await super().send_streaming(generator, use_fallback)

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

@@ -106,7 +106,7 @@ async def send(self, message: MessageChain):
         )
         await super().send(message)
 
-    async def send_streaming(self, generator):
+    async def send_streaming(self, generator, use_fallback: bool = False):
         final_data = ""
         async for chain in generator:
             final_data += await WebChatMessageEvent._send(
@@ -121,4 +121,4 @@ async def send_streaming(self, generator):
                 "cid": self.session_id.split("!")[-1],
             }
         )
-        await super().send_streaming(generator)
+        await super().send_streaming(generator, use_fallback)