from __future__ import annotations import asyncio import time from typing import TYPE_CHECKING if TYPE_CHECKING: from googleapiclient.discovery import Resource from yuxi.channels.models import DeliveryResult from yuxi.utils.logging_config import logger from .send import delete_message, send_message, update_message DEFAULT_UPDATE_INTERVAL_MS = 800 DEFAULT_COALESCE_MIN_CHARS = 1500 DEFAULT_COALESCE_IDLE_MS = 1000 class StreamManager: def __init__( self, update_interval_ms: int = DEFAULT_UPDATE_INTERVAL_MS, coalesce_min_chars: int = DEFAULT_COALESCE_MIN_CHARS, coalesce_idle_ms: int = DEFAULT_COALESCE_IDLE_MS, ): self._messages: dict[str, str] = {} self._texts: dict[str, str] = {} self._last_update: dict[str, float] = {} self._chars_at_last_update: dict[str, int] = {} self._update_interval_ms = update_interval_ms self._coalesce_min_chars = coalesce_min_chars self._coalesce_idle_ms = coalesce_idle_ms self._lock = asyncio.Lock() self._typing_messages: dict[str, str] = {} self._typing_replaced: set[str] = set() @property def update_interval_ms(self) -> int: return self._update_interval_ms def has_pending(self, chat_id: str) -> bool: return chat_id in self._messages def get_message_name(self, chat_id: str) -> str | None: return self._messages.get(chat_id) def get_accumulated_text(self, chat_id: str) -> str: return self._texts.get(chat_id, "") async def register_message(self, chat_id: str, message_name: str, text: str) -> None: async with self._lock: self._messages[chat_id] = message_name self._texts[chat_id] = text self._last_update[chat_id] = time.monotonic() self._chars_at_last_update[chat_id] = len(text) async def register_first_chunk(self, chat_id: str, message_name: str, text: str) -> None: async with self._lock: self._messages[chat_id] = message_name self._texts[chat_id] = text self._last_update[chat_id] = time.monotonic() self._chars_at_last_update[chat_id] = len(text) typing_id = self._typing_messages.pop(chat_id, None) if typing_id: self._typing_replaced.add(chat_id) logger.debug(f"Stream: replaced typing message {typing_id} with first chunk {message_name}") async def append_text(self, chat_id: str, chunk: str) -> str: async with self._lock: current = self._texts.get(chat_id, "") current += chunk self._texts[chat_id] = current return current def should_update(self, chat_id: str) -> bool: last = self._last_update.get(chat_id, 0) elapsed_ms = (time.monotonic() - last) * 1000 if elapsed_ms >= self._update_interval_ms: return True if self._coalesce_idle_ms > 0: idle_threshold = max(self._coalesce_idle_ms, self._update_interval_ms) if elapsed_ms >= idle_threshold: current_text = self._texts.get(chat_id, "") last_snapshot = self._chars_at_last_update.get(chat_id, 0) new_chars_since_last = len(current_text) - last_snapshot return new_chars_since_last >= self._coalesce_min_chars return False async def mark_update(self, chat_id: str) -> None: async with self._lock: self._last_update[chat_id] = time.monotonic() self._chars_at_last_update[chat_id] = len(self._texts.get(chat_id, "")) async def create_typing_message( self, chat_service: Resource, chat_id: str, bot_name: str = "Bot", ) -> str | None: async with self._lock: if chat_id in self._typing_messages: return self._typing_messages[chat_id] try: body = {"text": f"*{bot_name} is typing...*"} result = await send_message(chat_service, chat_id, body) if result.success and result.message_id: async with self._lock: self._typing_messages[chat_id] = result.message_id return result.message_id except Exception as e: logger.warning(f"Failed to create typing message in {chat_id}: {e}") return None async def clear_typing_message( self, chat_service: Resource, chat_id: str, ) -> None: msg_id = None replaced = False async with self._lock: msg_id = self._typing_messages.pop(chat_id, None) replaced = chat_id in self._typing_replaced if msg_id and not replaced: try: await delete_message(chat_service, msg_id) except Exception as e: logger.warning(f"Failed to delete typing message {msg_id}: {e}") async with self._lock: self._typing_replaced.discard(chat_id) async def handle_error( self, chat_service: Resource, chat_id: str, ) -> None: await self.clear_typing_message(chat_service, chat_id) async with self._lock: self._messages.pop(chat_id, None) self._texts.pop(chat_id, None) self._last_update.pop(chat_id, None) self._chars_at_last_update.pop(chat_id, None) async def finalize( self, chat_service: Resource, chat_id: str, ) -> DeliveryResult: await self.clear_typing_message(chat_service, chat_id) async with self._lock: message_name = self._messages.pop(chat_id, None) text = self._texts.pop(chat_id, None) self._last_update.pop(chat_id, None) self._chars_at_last_update.pop(chat_id, None) if not message_name or text is None: return DeliveryResult(success=False, error="No pending stream message") return await update_message(chat_service, message_name, {"text": text}) async def send_update( self, chat_service: Resource, chat_id: str, finished: bool = False, ) -> DeliveryResult | None: async with self._lock: message_name = self._messages.get(chat_id) text = self._texts.get(chat_id) if not message_name or text is None: return None if chat_id in self._typing_messages and text: self._typing_messages.pop(chat_id, None) self._typing_replaced.add(chat_id) if self._coalesce_min_chars > 0 and not finished: body = {"text": text + "\n\n_正在生成中..._"} else: body = {"text": text if finished else text + " ⏳"} result = await update_message(chat_service, message_name, body) if result.success: async with self._lock: if finished: self._messages.pop(chat_id, None) self._texts.pop(chat_id, None) self._last_update.pop(chat_id, None) self._chars_at_last_update.pop(chat_id, None) else: self._last_update[chat_id] = time.monotonic() if finished: await self.clear_typing_message(chat_service, chat_id) return result async def clear(self) -> None: async with self._lock: self._messages.clear() self._texts.clear() self._last_update.clear() self._chars_at_last_update.clear() self._typing_messages.clear() self._typing_replaced.clear() @property def pending_count(self) -> int: return len(self._messages)