Kris
1f78c44b03
这是一个批量整理提交，包含以下主要改动： 1. 删除多处冗余的空行和未使用的导入 2. 修复文件末尾缺少换行符的问题 3. 调整部分模块的导入顺序与代码排版 4. 修复部分配置默认值与策略逻辑 5. 新增多个功能模块与辅助工具 6. 完善异常处理与日志记录 7. 修复速率限制、消息缓存、权限校验等逻辑bug 8. 废弃部分旧有API与配置项并添加警告提示
216 lines
7.5 KiB
Python
216 lines
7.5 KiB
Python
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 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._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)
|
|
|
|
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)
|