feat(channels/utils): 添加流式消息处理工具模块
新增了SmartCoalesceBuffer、StreamStateMachine等流式消息处理相关工具类， 提供配置解析、缓冲合并、状态管理和分块发送能力，完善流式消息处理链路。
This commit is contained in:
parent
fb8bd90329
commit
4a5051e7d1
15
backend/package/yuxi/channels/utils/__init__.py
Normal file
15
backend/package/yuxi/channels/utils/__init__.py
Normal file
@ -0,0 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from yuxi.channels.utils.streaming_fallback import (
|
||||
SmartCoalesceBuffer,
|
||||
StreamingConfig,
|
||||
StreamStateMachine,
|
||||
UniversalChunkedSender,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SmartCoalesceBuffer",
|
||||
"StreamingConfig",
|
||||
"StreamStateMachine",
|
||||
"UniversalChunkedSender",
|
||||
]
|
||||
401
backend/package/yuxi/channels/utils/streaming_fallback.py
Normal file
401
backend/package/yuxi/channels/utils/streaming_fallback.py
Normal file
@ -0,0 +1,401 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channels.models import DeliveryResult
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
STREAM_CONFIG_DEFAULTS = {
|
||||
"mode": "partial",
|
||||
"preview": {
|
||||
"enabled": True,
|
||||
"chunk": {
|
||||
"min_chars": 50,
|
||||
"max_chars": 200,
|
||||
"break_preference": "newline",
|
||||
},
|
||||
"tool_progress": True,
|
||||
},
|
||||
"block": {
|
||||
"enabled": True,
|
||||
"coalesce": False,
|
||||
"coalesce_min_chars": 1500,
|
||||
"coalesce_idle_ms": 1000,
|
||||
},
|
||||
"typing_indicator": {
|
||||
"enabled": True,
|
||||
"duration_ms": 5000,
|
||||
},
|
||||
"fallback": {
|
||||
"on_edit_failure": "send_final",
|
||||
"on_rate_limit": "send_final",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingConfig:
|
||||
mode: str = "off"
|
||||
preview_enabled: bool = True
|
||||
preview_chunk_min_chars: int = 50
|
||||
preview_chunk_max_chars: int = 200
|
||||
preview_chunk_break_preference: str = "newline"
|
||||
preview_tool_progress: bool = True
|
||||
block_enabled: bool = True
|
||||
block_coalesce: bool = False
|
||||
block_coalesce_min_chars: int = 1500
|
||||
block_coalesce_idle_ms: int = 1000
|
||||
typing_indicator_enabled: bool = True
|
||||
typing_indicator_duration_ms: int = 5000
|
||||
fallback_on_edit_failure: str = "send_final"
|
||||
fallback_on_rate_limit: str = "send_final"
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any] | None) -> StreamingConfig:
|
||||
if not config:
|
||||
return cls()
|
||||
|
||||
streaming = config.get("streaming", {})
|
||||
if not isinstance(streaming, dict):
|
||||
streaming = {"mode": str(streaming)} if streaming else {}
|
||||
|
||||
mode = streaming.get("mode", config.get("streaming_mode", "off"))
|
||||
|
||||
preview = streaming.get("preview", {})
|
||||
if not isinstance(preview, dict):
|
||||
preview = {}
|
||||
chunk = preview.get("chunk", {})
|
||||
if not isinstance(chunk, dict):
|
||||
chunk = {}
|
||||
|
||||
block = streaming.get("block", {})
|
||||
if not isinstance(block, dict):
|
||||
block = {}
|
||||
|
||||
typing_cfg = streaming.get("typing_indicator", {})
|
||||
if not isinstance(typing_cfg, dict):
|
||||
typing_cfg = {}
|
||||
|
||||
fallback = streaming.get("fallback", {})
|
||||
if not isinstance(fallback, dict):
|
||||
fallback = {}
|
||||
|
||||
return cls(
|
||||
mode=mode,
|
||||
preview_enabled=preview.get("enabled", True),
|
||||
preview_chunk_min_chars=chunk.get("min_chars", 50),
|
||||
preview_chunk_max_chars=chunk.get("max_chars", 200),
|
||||
preview_chunk_break_preference=chunk.get("break_preference", "newline"),
|
||||
preview_tool_progress=preview.get("tool_progress", True),
|
||||
block_enabled=block.get("enabled", True),
|
||||
block_coalesce=block.get("coalesce", False),
|
||||
block_coalesce_min_chars=block.get("coalesce_min_chars", 1500),
|
||||
block_coalesce_idle_ms=block.get("coalesce_idle_ms", 1000),
|
||||
typing_indicator_enabled=typing_cfg.get("enabled", True),
|
||||
typing_indicator_duration_ms=typing_cfg.get("duration_ms", 5000),
|
||||
fallback_on_edit_failure=fallback.get("on_edit_failure", "send_final"),
|
||||
fallback_on_rate_limit=fallback.get("on_rate_limit", "send_final"),
|
||||
)
|
||||
|
||||
|
||||
SendFn = Callable[[str], Awaitable[DeliveryResult]]
|
||||
TypingFn = Callable[[str, bool], Awaitable[DeliveryResult | None]]
|
||||
EditFn = Callable[[str, str, str, bool], Awaitable[DeliveryResult]]
|
||||
|
||||
|
||||
class SmartCoalesceBuffer:
|
||||
def __init__(
|
||||
self,
|
||||
min_chars: int = 1500,
|
||||
idle_ms: int = 1000,
|
||||
max_wait_ms: int = 5000,
|
||||
):
|
||||
self._min_chars = min_chars
|
||||
self._idle_ms = idle_ms
|
||||
self._max_wait_ms = max_wait_ms
|
||||
self._buffer: str = ""
|
||||
self._last_feed: float = 0.0
|
||||
self._first_feed: float = 0.0
|
||||
self._has_pending: bool = False
|
||||
|
||||
def feed(self, text: str) -> str | None:
|
||||
now = time.monotonic()
|
||||
if not self._buffer:
|
||||
self._first_feed = now
|
||||
self._buffer = text
|
||||
self._last_feed = now
|
||||
self._has_pending = True
|
||||
return None
|
||||
|
||||
self._buffer += text
|
||||
self._last_feed = now
|
||||
|
||||
if len(self._buffer) >= self._min_chars:
|
||||
result = self._buffer
|
||||
self._buffer = ""
|
||||
self._has_pending = False
|
||||
return result
|
||||
|
||||
elapsed_since_first = (now - self._first_feed) * 1000
|
||||
if elapsed_since_first >= self._max_wait_ms:
|
||||
result = self._buffer
|
||||
self._buffer = ""
|
||||
self._has_pending = False
|
||||
return result
|
||||
|
||||
return None
|
||||
|
||||
def peek(self) -> str:
|
||||
return self._buffer
|
||||
|
||||
def flush(self) -> str:
|
||||
result = self._buffer
|
||||
self._buffer = ""
|
||||
self._has_pending = False
|
||||
return result
|
||||
|
||||
@property
|
||||
def pending(self) -> bool:
|
||||
return self._has_pending
|
||||
|
||||
def reset(self) -> None:
|
||||
self._buffer = ""
|
||||
self._first_feed = 0.0
|
||||
self._last_feed = 0.0
|
||||
self._has_pending = False
|
||||
|
||||
|
||||
class StreamPhase(StrEnum):
|
||||
INIT = "init"
|
||||
TYPING = "typing"
|
||||
STREAMING = "streaming"
|
||||
FINALIZING = "finalizing"
|
||||
ERROR = "error"
|
||||
FALLBACK_SEND = "fallback_send"
|
||||
DONE = "done"
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamContext:
|
||||
chat_id: str
|
||||
msg_id: str = ""
|
||||
phase: StreamPhase = StreamPhase.INIT
|
||||
accumulated_text: str = ""
|
||||
started_at: float = field(default_factory=time.monotonic)
|
||||
last_update_at: float = 0.0
|
||||
error: str = ""
|
||||
fallback_used: bool = False
|
||||
|
||||
|
||||
class StreamStateMachine:
|
||||
def __init__(self):
|
||||
self._ctx: dict[str, StreamContext] = {}
|
||||
|
||||
def get(self, chat_id: str) -> StreamContext | None:
|
||||
return self._ctx.get(chat_id)
|
||||
|
||||
def start(self, chat_id: str, msg_id: str = "") -> StreamContext:
|
||||
ctx = StreamContext(chat_id=chat_id, msg_id=msg_id)
|
||||
self._ctx[chat_id] = ctx
|
||||
return ctx
|
||||
|
||||
async def transition_typing(self, chat_id: str, typing_fn: TypingFn | None = None) -> None:
|
||||
ctx = self._ctx.get(chat_id)
|
||||
if ctx is None:
|
||||
return
|
||||
ctx.phase = StreamPhase.TYPING
|
||||
if typing_fn:
|
||||
try:
|
||||
await typing_fn(chat_id, True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def transition_streaming(
|
||||
self,
|
||||
chat_id: str,
|
||||
send_fn: SendFn | None = None,
|
||||
initial_text: str = "",
|
||||
) -> None:
|
||||
ctx = self._ctx.get(chat_id)
|
||||
if ctx is None:
|
||||
return
|
||||
ctx.phase = StreamPhase.STREAMING
|
||||
ctx.accumulated_text = initial_text
|
||||
ctx.last_update_at = time.monotonic()
|
||||
|
||||
async def transition_error(
|
||||
self,
|
||||
chat_id: str,
|
||||
error: str,
|
||||
fallback_fn: SendFn | None = None,
|
||||
final_text: str = "",
|
||||
) -> DeliveryResult:
|
||||
ctx = self._ctx.get(chat_id)
|
||||
if ctx is None:
|
||||
return DeliveryResult(success=False, error=error)
|
||||
|
||||
ctx.phase = StreamPhase.ERROR
|
||||
ctx.error = error
|
||||
ctx.fallback_used = True
|
||||
|
||||
if fallback_fn and final_text:
|
||||
logger.warning(f"Stream error for {chat_id}: {error}, falling back to send_final")
|
||||
result = await fallback_fn(final_text)
|
||||
ctx.phase = StreamPhase.DONE
|
||||
return result
|
||||
|
||||
return DeliveryResult(success=False, error=error)
|
||||
|
||||
async def finalize(
|
||||
self,
|
||||
chat_id: str,
|
||||
typing_fn: TypingFn | None = None,
|
||||
) -> StreamContext | None:
|
||||
ctx = self._ctx.get(chat_id)
|
||||
if ctx is None:
|
||||
return None
|
||||
|
||||
ctx.phase = StreamPhase.FINALIZING
|
||||
|
||||
if typing_fn:
|
||||
try:
|
||||
await typing_fn(chat_id, False)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
ctx.phase = StreamPhase.DONE
|
||||
return ctx
|
||||
|
||||
def cleanup(self, chat_id: str) -> StreamContext | None:
|
||||
return self._ctx.pop(chat_id, None)
|
||||
|
||||
def is_active(self, chat_id: str) -> bool:
|
||||
ctx = self._ctx.get(chat_id)
|
||||
return ctx is not None and ctx.phase not in (
|
||||
StreamPhase.DONE,
|
||||
StreamPhase.ERROR,
|
||||
)
|
||||
|
||||
def clear(self) -> None:
|
||||
self._ctx.clear()
|
||||
|
||||
@property
|
||||
def active_count(self) -> int:
|
||||
return sum(1 for ctx in self._ctx.values() if ctx.phase not in (StreamPhase.DONE, StreamPhase.ERROR))
|
||||
|
||||
|
||||
class UniversalChunkedSender:
|
||||
def __init__(
|
||||
self,
|
||||
send_fn: SendFn,
|
||||
chunk_size: int = 2000,
|
||||
prefix_format: str = "",
|
||||
rate_limit_delay_ms: int = 0,
|
||||
max_chunks: int = 10,
|
||||
):
|
||||
self._send_fn = send_fn
|
||||
self._chunk_size = max(1, chunk_size)
|
||||
self._prefix_format = prefix_format
|
||||
self._rate_limit_delay_ms = rate_limit_delay_ms
|
||||
self._max_chunks = max_chunks
|
||||
self._sent_chunks: dict[str, int] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def send_stream(
|
||||
self,
|
||||
chat_id: str,
|
||||
text: str,
|
||||
finished: bool = False,
|
||||
) -> list[DeliveryResult]:
|
||||
results: list[DeliveryResult] = []
|
||||
|
||||
async with self._lock:
|
||||
sent_count = self._sent_chunks.get(chat_id, 0)
|
||||
total_estimate = sent_count + 1 if not finished else sent_count + 1
|
||||
|
||||
content = self._build_content(text, sent_count + 1, total_estimate)
|
||||
|
||||
result = await self._send_fn(content)
|
||||
results.append(result)
|
||||
|
||||
if result.success:
|
||||
async with self._lock:
|
||||
self._sent_chunks[chat_id] = sent_count + 1
|
||||
|
||||
if finished:
|
||||
async with self._lock:
|
||||
self._sent_chunks.pop(chat_id, None)
|
||||
|
||||
if self._prefix_format:
|
||||
summary = self._prefix_format.format_map({"total": sent_count + 1})
|
||||
if summary and summary != content:
|
||||
summary_result = await self._send_fn(summary)
|
||||
results.append(summary_result)
|
||||
|
||||
return results
|
||||
|
||||
async def send_text_chunked(
|
||||
self,
|
||||
chat_id: str,
|
||||
text: str,
|
||||
finished: bool = False,
|
||||
) -> list[DeliveryResult]:
|
||||
results: list[DeliveryResult] = []
|
||||
paragraphs = text.split("\n\n")
|
||||
accumulated = ""
|
||||
chunk_idx = 0
|
||||
|
||||
for para in paragraphs:
|
||||
accumulated += para + "\n\n"
|
||||
if len(accumulated) >= self._chunk_size:
|
||||
chunk_idx += 1
|
||||
if chunk_idx > self._max_chunks:
|
||||
accumulated = accumulated[: self._chunk_size]
|
||||
chunk_result = await self._send_fn(accumulated)
|
||||
results.append(chunk_result)
|
||||
accumulated = ""
|
||||
continue
|
||||
|
||||
content = accumulated.strip()
|
||||
if self._prefix_format:
|
||||
content = f"{content}\n\n{self._prefix_format}"
|
||||
chunk_result = await self._send_fn(content)
|
||||
results.append(chunk_result)
|
||||
accumulated = ""
|
||||
|
||||
if self._rate_limit_delay_ms > 0:
|
||||
await asyncio.sleep(self._rate_limit_delay_ms / 1000.0)
|
||||
|
||||
if accumulated.strip():
|
||||
chunk_idx += 1
|
||||
chunk_result = await self._send_fn(accumulated.strip())
|
||||
results.append(chunk_result)
|
||||
|
||||
return results
|
||||
|
||||
def _build_content(self, text: str, chunk_num: int, total: int) -> str:
|
||||
if not self._prefix_format:
|
||||
return text
|
||||
|
||||
prefix = self._prefix_format.format_map({"chunk": str(chunk_num), "total": str(total)})
|
||||
|
||||
return f"{prefix}\n{text}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SmartCoalesceBuffer",
|
||||
"StreamPhase",
|
||||
"StreamContext",
|
||||
"StreamStateMachine",
|
||||
"UniversalChunkedSender",
|
||||
"StreamingConfig",
|
||||
"STREAM_CONFIG_DEFAULTS",
|
||||
"SendFn",
|
||||
"TypingFn",
|
||||
"EditFn",
|
||||
]
|
||||
Loading…
Reference in New Issue
Block a user