Kris
552aef767c
新增QQ Bot适配器完整代码栈，包含： 1. 基础适配器入口与工具类封装 2. 会话管理、重试队列与流量控制 3. 命令系统与内置指令（ping/help/status等） 4. 富媒体消息处理与格式转换 5. 引用存储与审批管理 6. 凭证备份与会话持久化 7. 健康检查与交互回调系统
344 lines
12 KiB
Python
344 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, auto
|
|
from collections.abc import AsyncGenerator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamEventType(Enum):
|
|
FLUSH = auto()
|
|
COMPLETE = auto()
|
|
ERROR = auto()
|
|
|
|
|
|
class FlushStrategy(Enum):
|
|
PER_CHAR = auto()
|
|
INTERVAL = auto()
|
|
BACKLOG = auto()
|
|
|
|
|
|
@dataclass
|
|
class StreamEvent:
|
|
event_type: StreamEventType
|
|
content: str = ""
|
|
full_content: str = ""
|
|
error: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class StreamCheckpoint:
|
|
chat_id: str
|
|
msg_id: str = ""
|
|
content: str = ""
|
|
sent_seq: int = 0
|
|
timestamp: float = field(default_factory=time.time)
|
|
|
|
@property
|
|
def is_stale(self) -> bool:
|
|
return time.time() - self.timestamp > 600
|
|
|
|
|
|
class FlushController:
|
|
strategy: FlushStrategy = FlushStrategy.BACKLOG
|
|
flush_interval: float = 0.6
|
|
backlog_threshold: int = 30
|
|
max_retries: int = 3
|
|
retry_delay: float = 1.0
|
|
seq_counter: int = 0
|
|
_buffer: list[str] = field(default_factory=list, repr=False)
|
|
_last_flush: float = 0.0
|
|
_accumulated: str = ""
|
|
|
|
def reset(self) -> None:
|
|
self._buffer = []
|
|
self._last_flush = 0.0
|
|
self._accumulated = ""
|
|
self.seq_counter = 0
|
|
|
|
def feed(self, chunk: str) -> list[str]:
|
|
self._buffer.append(chunk)
|
|
self._accumulated += chunk
|
|
|
|
if self.strategy == FlushStrategy.PER_CHAR:
|
|
return self._flush_all()
|
|
|
|
if self.strategy == FlushStrategy.BACKLOG and len(self._accumulated) >= self.backlog_threshold:
|
|
return self._flush_all()
|
|
|
|
now = time.monotonic()
|
|
if self.strategy == FlushStrategy.INTERVAL and now - self._last_flush >= self.flush_interval:
|
|
return self._flush_all()
|
|
|
|
return []
|
|
|
|
def flush_remaining(self) -> tuple[list[str], str]:
|
|
if self.strategy == FlushStrategy.PER_CHAR:
|
|
return [], self._accumulated
|
|
batches = self._flush_all() if self._buffer else []
|
|
return batches, self._accumulated
|
|
|
|
def _flush_all(self) -> list[str]:
|
|
if not self._buffer:
|
|
return []
|
|
batches = list(self._buffer)
|
|
self._buffer = []
|
|
self._accumulated = ""
|
|
self._last_flush = time.monotonic()
|
|
return batches
|
|
|
|
|
|
class C2CStreamingController:
|
|
def __init__(
|
|
self,
|
|
send_message_fn,
|
|
retry_delay: float = 1.0,
|
|
max_retries: int = 3,
|
|
flush_interval: float = 0.6,
|
|
) -> None:
|
|
self._send_message_fn = send_message_fn
|
|
self._retry_delay = retry_delay
|
|
self._max_retries = max_retries
|
|
self._flush_interval = flush_interval
|
|
self._flush_controller = FlushController(strategy=FlushStrategy.BACKLOG, flush_interval=flush_interval)
|
|
self._any_chunk_delivered: bool = False
|
|
self._static_fallback_msg: str = ""
|
|
self._active_checkpoints: dict[str, StreamCheckpoint] = {}
|
|
|
|
@property
|
|
def flush_controller(self) -> FlushController:
|
|
return self._flush_controller
|
|
|
|
async def stream(
|
|
self,
|
|
chat_id: str,
|
|
msg_id: str,
|
|
content_generator: AsyncGenerator[str, None],
|
|
event_id: str = "",
|
|
) -> bool:
|
|
self._flush_controller.reset()
|
|
self._any_chunk_delivered = False
|
|
self._static_fallback_msg = ""
|
|
|
|
checkpoint = StreamCheckpoint(chat_id=chat_id, msg_id=msg_id)
|
|
self._active_checkpoints[chat_id] = checkpoint
|
|
|
|
collected = ""
|
|
in_media_interrupt = False
|
|
media_buffer = ""
|
|
|
|
try:
|
|
async for chunk in content_generator:
|
|
if not chunk:
|
|
continue
|
|
|
|
media_splits = self._split_on_media_tags(chunk)
|
|
for segment, is_media in media_splits:
|
|
if is_media:
|
|
in_media_interrupt = True
|
|
media_buffer += segment
|
|
else:
|
|
if in_media_interrupt and media_buffer:
|
|
checkpoint.content = collected
|
|
await self._flush_and_end_stream(chat_id, msg_id, checkpoint, event_id)
|
|
await self._send_media_interruption(chat_id, media_buffer)
|
|
media_buffer = ""
|
|
in_media_interrupt = False
|
|
|
|
new_msg_id = f"{msg_id}_resume_{checkpoint.sent_seq}"
|
|
checkpoint.msg_id = new_msg_id
|
|
self._flush_controller.reset()
|
|
|
|
collected += segment
|
|
batches = self._flush_controller.feed(segment)
|
|
for batch in batches:
|
|
self._flush_controller.seq_counter += 1
|
|
success = await self._send_stream_chunk(
|
|
chat_id,
|
|
checkpoint.msg_id or msg_id,
|
|
batch,
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
)
|
|
if success:
|
|
self._any_chunk_delivered = True
|
|
else:
|
|
logger.warning(
|
|
"C2CStreaming: flush failed for seq %d", self._flush_controller.seq_counter
|
|
)
|
|
|
|
self._static_fallback_msg = collected
|
|
|
|
flush_batches, full = self._flush_controller.flush_remaining()
|
|
for batch in flush_batches:
|
|
self._flush_controller.seq_counter += 1
|
|
success = await self._send_stream_chunk(
|
|
chat_id,
|
|
checkpoint.msg_id or msg_id,
|
|
batch,
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
)
|
|
if success:
|
|
self._any_chunk_delivered = True
|
|
|
|
self._flush_controller.seq_counter += 1
|
|
success = await self._send_stream_chunk(
|
|
chat_id,
|
|
checkpoint.msg_id or msg_id,
|
|
"",
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
is_end=True,
|
|
)
|
|
self._active_checkpoints.pop(chat_id, None)
|
|
return success
|
|
|
|
except Exception as e:
|
|
logger.exception("C2CStreaming: stream error for %s", chat_id)
|
|
self._active_checkpoints.pop(chat_id, None)
|
|
return await self._cancel_stream(chat_id, msg_id, event_id, str(e))
|
|
|
|
async def _send_stream_chunk(
|
|
self,
|
|
chat_id: str,
|
|
msg_id: str,
|
|
content: str,
|
|
msg_seq: int,
|
|
event_id: str = "",
|
|
is_end: bool = False,
|
|
) -> bool:
|
|
for attempt in range(self._max_retries):
|
|
try:
|
|
payload = {
|
|
"content": content,
|
|
"msg_type": 0,
|
|
"msg_id": msg_id,
|
|
"msg_seq": msg_seq,
|
|
"stream": {"state": 2 if is_end else 1},
|
|
}
|
|
if event_id:
|
|
payload["event_id"] = event_id
|
|
|
|
response = await self._send_message_fn(chat_id, payload)
|
|
if response is not None:
|
|
cp = self._active_checkpoints.get(chat_id)
|
|
if cp:
|
|
cp.sent_seq = msg_seq
|
|
cp.timestamp = time.time()
|
|
return response is not None
|
|
except Exception:
|
|
if attempt < self._max_retries - 1:
|
|
await asyncio.sleep(self._retry_delay)
|
|
else:
|
|
logger.exception("C2CStreaming: send chunk failed after %d retries", self._max_retries)
|
|
return False
|
|
|
|
@staticmethod
|
|
def _split_on_media_tags(text: str) -> list[tuple[str, bool]]:
|
|
import re
|
|
|
|
from .media_tags import _MEDIA_TAG_RE, _IMG_TAG_RE, _VIDEO_TAG_RE
|
|
|
|
combined_re = re.compile(f"({_MEDIA_TAG_RE.pattern}|{_IMG_TAG_RE.pattern}|{_VIDEO_TAG_RE.pattern})")
|
|
|
|
results: list[tuple[str, bool]] = []
|
|
last_end = 0
|
|
for match in combined_re.finditer(text):
|
|
if match.start() > last_end:
|
|
results.append((text[last_end : match.start()], False))
|
|
results.append((match.group(0), True))
|
|
last_end = match.end()
|
|
|
|
if last_end < len(text):
|
|
results.append((text[last_end:], False))
|
|
|
|
return results
|
|
|
|
async def _flush_and_end_stream(
|
|
self, chat_id: str, msg_id: str, checkpoint: StreamCheckpoint, event_id: str
|
|
) -> None:
|
|
flush_batches, _ = self._flush_controller.flush_remaining()
|
|
for batch in flush_batches:
|
|
self._flush_controller.seq_counter += 1
|
|
await self._send_stream_chunk(
|
|
chat_id,
|
|
msg_id,
|
|
batch,
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
)
|
|
|
|
self._flush_controller.seq_counter += 1
|
|
await self._send_stream_chunk(
|
|
chat_id,
|
|
msg_id,
|
|
"",
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
is_end=True,
|
|
)
|
|
|
|
async def _send_media_interruption(self, chat_id: str, media_tags: str) -> None:
|
|
try:
|
|
payload = {
|
|
"content": media_tags,
|
|
"msg_type": 0,
|
|
"msg_id": "",
|
|
"msg_seq": 0,
|
|
}
|
|
await self._send_message_fn(chat_id, payload)
|
|
except Exception:
|
|
logger.exception("C2CStreaming: media interruption send failed")
|
|
|
|
async def _cancel_stream(self, chat_id: str, msg_id: str, event_id: str, reason: str) -> bool:
|
|
delivered_fallback = False
|
|
if self._static_fallback_msg and not self._any_chunk_delivered:
|
|
try:
|
|
static_payload = {
|
|
"content": self._static_fallback_msg[:2000],
|
|
"msg_type": 0,
|
|
"msg_id": msg_id,
|
|
"msg_seq": 0,
|
|
}
|
|
if event_id:
|
|
static_payload["event_id"] = event_id
|
|
await self._send_message_fn(chat_id, static_payload)
|
|
delivered_fallback = True
|
|
logger.info("C2CStreaming: delivered static fallback message for %s", chat_id)
|
|
except Exception:
|
|
logger.exception("C2CStreaming: static fallback delivery failed for %s", chat_id)
|
|
|
|
if not delivered_fallback:
|
|
try:
|
|
payload = {
|
|
"content": reason,
|
|
"msg_type": 0,
|
|
"msg_id": msg_id,
|
|
"msg_seq": 0,
|
|
"stream": {"state": 0},
|
|
}
|
|
if event_id:
|
|
payload["event_id"] = event_id
|
|
await self._send_message_fn(chat_id, payload)
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
def get_checkpoint(self, chat_id: str) -> StreamCheckpoint | None:
|
|
return self._active_checkpoints.get(chat_id)
|
|
|
|
def has_pending_stream(self, chat_id: str) -> bool:
|
|
cp = self._active_checkpoints.get(chat_id)
|
|
return cp is not None and not cp.is_stale
|
|
|
|
def cleanup_stale_checkpoints(self) -> int:
|
|
stale = [cid for cid, cp in self._active_checkpoints.items() if cp.is_stale]
|
|
for cid in stale:
|
|
self._active_checkpoints.pop(cid, None)
|
|
return len(stale)
|