from __future__ import annotations import asyncio import re import time from typing import Any from telegram import Bot from telegram.error import BadRequest, RetryAfter from yuxi.channels.adapters.telegram.format import markdown_to_html, strip_all_tags from yuxi.channels.adapters.telegram.stream.lane_delivery import LaneDelivery, LaneType from yuxi.channels.adapters.telegram.stream.reasoning_lane_coordinator import ReasoningLaneCoordinator from yuxi.utils.logging_config import logger _STREAM_STORAGE: dict[str, dict] = {} _STREAM_TTL_SECONDS = 300 _STREAM_STORAGE_LOCK = asyncio.Lock() _lane_delivery = LaneDelivery() _reasoning_coordinator = ReasoningLaneCoordinator() _BREAK_BOUNDARY_SENTENCE = re.compile(r"[。！？.!?\n]") _BREAK_BOUNDARY_WORD = re.compile(r"\s") def streaming_mode_from_config(config: dict[str, Any]) -> str: streaming_cfg = config.get("streaming", {}) if isinstance(streaming_cfg, dict): return streaming_cfg.get("mode", config.get("streaming_mode", "partial")) if isinstance(streaming_cfg, str): return streaming_cfg return config.get("streaming_mode", "partial") def _resolve_stream_config(config: dict[str, Any]) -> dict[str, Any]: streaming_cfg = config.get("streaming", {}) if not isinstance(streaming_cfg, dict): streaming_cfg = {} preview_cfg = streaming_cfg.get("preview", {}) if not isinstance(preview_cfg, dict): preview_cfg = {} chunk_cfg = preview_cfg.get("chunk", {}) if not isinstance(chunk_cfg, dict): chunk_cfg = {} block_cfg = streaming_cfg.get("block", {}) if not isinstance(block_cfg, dict): block_cfg = {} return { "preview_chunk_min_chars": chunk_cfg.get("min_chars", chunk_cfg.get("minChars", 50)), "preview_chunk_max_chars": chunk_cfg.get("max_chars", chunk_cfg.get("maxChars", 200)), "preview_chunk_break_preference": chunk_cfg.get( "break_preference", chunk_cfg.get("breakPreference", "newline") ), "preview_tool_progress": preview_cfg.get("tool_progress", preview_cfg.get("toolProgress", True)), "block_enabled": block_cfg.get("enabled", True), "block_coalesce": block_cfg.get("coalesce", False), } def _cleanup_expired(now: float) -> None: expired = [key for key, state in _STREAM_STORAGE.items() if now - state.get("created_at", 0) > _STREAM_TTL_SECONDS] for key in expired: del _STREAM_STORAGE[key] def _format_block(content: str, max_chars: int = 4096) -> str: if len(content) <= max_chars: return content return content[: max_chars - 4] + "..." def _format_preview_chunk(content: str, stream_cfg: dict[str, Any]) -> str: min_chars = stream_cfg.get("preview_chunk_min_chars", 50) max_chars = stream_cfg.get("preview_chunk_max_chars", 200) break_pref = stream_cfg.get("preview_chunk_break_preference", "newline") if len(content) <= max_chars: return content if break_pref == "any": return content[:max_chars] + "..." lo = max(min_chars, max_chars // 2) hi = max_chars if break_pref == "sentence": pattern = _BREAK_BOUNDARY_SENTENCE elif break_pref == "word": pattern = _BREAK_BOUNDARY_WORD else: pattern = re.compile(r"\n") candidates = [m.start() for m in pattern.finditer(content, lo, hi)] if candidates: return content[: candidates[-1] + 1].rstrip() + "..." return content[:max_chars] + "..." def _format_progress(content: str, state: dict) -> str: counter = state.get("edit_counter", 0) + 1 state["edit_counter"] = counter dots = "." * ((counter % 3) + 1) if len(content) > 3800: content = content[:3800] return f"{content}\n\n\u23f3 {dots}" def _format_tool_progress(content: str, stream_cfg: dict[str, Any]) -> str: if not stream_cfg.get("preview_tool_progress", True): return content return f"\U0001f6e0\ufe0f {content}" async def stream_send_chunked( bot: Bot, chat_id: str, content: str, message_id: str | None = None, finished: bool = False, config: dict[str, Any] | None = None, ) -> str | None: config = config or {} stream_cfg = _resolve_stream_config(config) mode = streaming_mode_from_config(config) edit_interval = config.get("stream_edit_interval_ms", 500) / 1000 stream_key = f"{chat_id}:{message_id}" if message_id else None async with _STREAM_STORAGE_LOCK: state = _STREAM_STORAGE.get(stream_key) if stream_key else None _cleanup_expired(time.monotonic()) if finished and state: async with _STREAM_STORAGE_LOCK: del _STREAM_STORAGE[stream_key] if state.get("message_id"): try: final_content = content if mode == "progress": final_content = _format_progress(content, state).rstrip("\n\u23f3 .") elif mode == "block" and stream_cfg.get("block_coalesce"): final_content = state.get("accumulated", "") + content formatted = markdown_to_html(final_content) await bot.edit_message_text( chat_id=chat_id, message_id=state["message_id"], text=formatted, parse_mode="HTML", ) except BadRequest: try: clean = strip_all_tags(content) await bot.edit_message_text( chat_id=chat_id, message_id=state["message_id"], text=clean, ) except Exception: pass if state.get("reasoning_state"): combined = _reasoning_coordinator.get_combined_text(chat_id) if combined and state["reasoning_state"] != stream_key: _reasoning_coordinator.cleanup(chat_id) return None if not message_id: try: if mode == "block": if stream_cfg.get("block_coalesce"): preview = _format_preview_chunk(content, stream_cfg) else: preview = _format_block(content, 100) + "..." if len(content) > 100 else content elif mode == "progress": preview = content[:50] + "..." if len(content) > 50 else content else: preview = _format_preview_chunk(content, stream_cfg) msg = await bot.send_message(chat_id=chat_id, text=preview) now = time.monotonic() async with _STREAM_STORAGE_LOCK: _STREAM_STORAGE[f"{chat_id}:{str(msg.message_id)}"] = { "message_id": msg.message_id, "last_edit": now, "counter": 0, "edit_counter": 0, "created_at": now, "mode": mode, "accumulated": "", } return str(msg.message_id) except Exception: logger.exception("[Telegram] Failed to create stream placeholder") return None if not state: return message_id if stream_cfg.get("block_coalesce"): state["accumulated"] = (state.get("accumulated", "") + "\n" + content).strip() return message_id counter = state.get("counter", 0) now = time.monotonic() if counter < 3 and (now - state.get("last_edit", 0)) < edit_interval: state["counter"] = counter + 1 return message_id state["counter"] = 0 state["last_edit"] = now display_content = content if mode == "block": display_content = _format_block(content) elif mode == "progress": display_content = _format_progress(content, state) try: formatted = markdown_to_html(display_content) await bot.edit_message_text( chat_id=chat_id, message_id=state["message_id"], text=formatted, parse_mode="HTML", ) except RetryAfter as e: await asyncio.sleep(e.retry_after) except BadRequest: pass except Exception: logger.exception(f"[Telegram] Stream edit failed for {chat_id}") return message_id async def stream_send_with_lane( bot: Bot, chat_id: str, content: str, message_id: str | None = None, finished: bool = False, config: dict[str, Any] | None = None, ) -> str | None: config = config or {} lane_enabled = ( config.get("streaming", {}).get("lane_enabled", False) if isinstance(config.get("streaming"), dict) else False ) if not lane_enabled: return await stream_send_chunked(bot, chat_id, content, message_id, finished, config) reasoning_enabled = _lane_delivery.is_reasoning_enabled() if reasoning_enabled: lane = _lane_delivery.extract_lane(content) if lane == LaneType.REASONING: state = _reasoning_coordinator.get_state(chat_id) if not state: _reasoning_coordinator.start_reasoning(chat_id) _reasoning_coordinator.append_reasoning(chat_id, content) if not message_id: return await stream_send_chunked(bot, chat_id, content, None, False, config) return message_id if lane == LaneType.MAIN: combined = _reasoning_coordinator.get_combined_text(chat_id) if combined: content = combined if finished: _reasoning_coordinator.complete(chat_id) return await stream_send_chunked(bot, chat_id, content, message_id, finished, config)