diff --git a/backend/package/yuxi/channels/utils/__init__.py b/backend/package/yuxi/channels/utils/__init__.py new file mode 100644 index 00000000..71116000 --- /dev/null +++ b/backend/package/yuxi/channels/utils/__init__.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from yuxi.channels.utils.streaming_fallback import ( + SmartCoalesceBuffer, + StreamingConfig, + StreamStateMachine, + UniversalChunkedSender, +) + +__all__ = [ + "SmartCoalesceBuffer", + "StreamingConfig", + "StreamStateMachine", + "UniversalChunkedSender", +] diff --git a/backend/package/yuxi/channels/utils/streaming_fallback.py b/backend/package/yuxi/channels/utils/streaming_fallback.py new file mode 100644 index 00000000..27a0f357 --- /dev/null +++ b/backend/package/yuxi/channels/utils/streaming_fallback.py @@ -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", +]