Kris
87fc26cd15
- 调整依赖导入顺序与清理冗余空行 - 修复setup wizard步骤存储逻辑 - 新增消息编辑/转发标记、卡片控件支持 - 完善目标校验、媒体上传、配对存储功能 - 优化流式响应与错误处理逻辑 - 增强SSRF防护与配置灵活性 - 新增/扩展内置命令与卡片处理能力 - 调整媒体大小限制与类型参数
219 lines
7.6 KiB
Python
219 lines
7.6 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 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)
|