from __future__ import annotations 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 update_message, send_message, delete_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._update_interval_ms = update_interval_ms self._coalesce_min_chars = coalesce_min_chars self._coalesce_idle_ms = coalesce_idle_ms 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, "") def register_message(self, chat_id: str, message_name: str, text: str) -> None: self._messages[chat_id] = message_name self._texts[chat_id] = text self._last_update[chat_id] = time.monotonic() def register_first_chunk(self, chat_id: str, message_name: str, text: str) -> None: self._messages[chat_id] = message_name self._texts[chat_id] = text self._last_update[chat_id] = time.monotonic() 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}") def append_text(self, chat_id: str, chunk: str) -> str: 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: new_chars_since_last = 0 return new_chars_since_last >= self._coalesce_min_chars return False def mark_update(self, chat_id: str) -> None: self._last_update[chat_id] = time.monotonic() async def create_typing_message( self, chat_service: Resource, chat_id: str, bot_name: str = "Bot", ) -> str | None: 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: 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 = self._typing_messages.pop(chat_id, None) if msg_id and chat_id not in self._typing_replaced: try: await delete_message(chat_service, msg_id) except Exception as e: logger.warning(f"Failed to delete typing message {msg_id}: {e}") 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) self._messages.pop(chat_id, None) self._texts.pop(chat_id, None) self._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) message_name = self._messages.pop(chat_id, None) text = self._texts.pop(chat_id, None) self._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: 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: typing_id = self._typing_messages.pop(chat_id, None) if typing_id: try: await delete_message(chat_service, typing_id) except Exception: pass self._typing_replaced.add(chat_id) body = {"text": text if finished else text + " ⏳"} result = await update_message(chat_service, message_name, body) if result.success: if finished: self._messages.pop(chat_id, None) self._texts.pop(chat_id, None) self._last_update.pop(chat_id, None) await self.clear_typing_message(chat_service, chat_id) else: self._last_update[chat_id] = time.monotonic() return result def clear(self) -> None: self._messages.clear() self._texts.clear() self._last_update.clear() self._typing_messages.clear() self._typing_replaced.clear() @property def pending_count(self) -> int: return len(self._messages)