from __future__ import annotations import asyncio from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from yuxi.channel.extensions.qqbot.types import QQBotChatType, QueuedMessage if TYPE_CHECKING: from yuxi.channel.common.attachment_cache import AttachmentCache @dataclass class PipelineContext: msg: QueuedMessage account_id: str config: dict = field(default_factory=dict) blocked: bool = False skip_reason: str | None = None metadata: dict = field(default_factory=dict) class QQBotPipeline: def __init__(self): self._stages: list[Any] = [] def add_stage(self, stage: Any) -> None: self._stages.append(stage) async def process(self, msg: QueuedMessage, account_id: str, config: dict) -> PipelineContext: ctx = PipelineContext(msg=msg, account_id=account_id, config=config) for stage in self._stages: if ctx.blocked: break try: ctx = await stage(ctx) except Exception: pass return ctx class DedupeStage: def __init__(self, ttl_seconds: int = 300): self._seen: dict[str, float] = {} self._ttl = ttl_seconds async def __call__(self, ctx: PipelineContext) -> PipelineContext: self._cleanup() msg_id = ctx.msg.msg_id if msg_id in self._seen: ctx.blocked = True ctx.skip_reason = "duplicate" return ctx import time self._seen[msg_id] = time.time() return ctx def _cleanup(self) -> None: import time now = time.time() expired = [k for k, v in self._seen.items() if now - v > self._ttl] for k in expired: del self._seen[k] class AccessStage: def __init__(self, check_fn: Any | None = None): self._check_fn = check_fn async def __call__(self, ctx: PipelineContext) -> PipelineContext: if self._check_fn: allowed = await self._check_fn(ctx.msg.sender_id, ctx.msg.chat_type) if not allowed: ctx.blocked = True ctx.skip_reason = "access_denied" return ctx class ContentStage: def __init__(self, attachment_cache: AttachmentCache | None = None): self._attachment_cache = attachment_cache async def __call__(self, ctx: PipelineContext) -> PipelineContext: session_id = ctx.msg.sender_id if self._attachment_cache is not None and ctx.msg.attachments: if not ctx.msg.content.strip(): for att in ctx.msg.attachments: await self._attachment_cache.add( session_id, att.url, att.content_type or "", filename=att.filename, ) ctx.blocked = True ctx.skip_reason = "attachment_cached" return ctx cached = await self._attachment_cache.consume(session_id) if cached: cached_tags = _build_cached_attachment_tags(cached) ctx.msg.content = f"{ctx.msg.content}\n{cached_tags}" if ctx.msg.content else cached_tags if ctx.msg.attachments: ctx.metadata["has_attachments"] = True ctx.metadata["attachment_count"] = len(ctx.msg.attachments) ctx.msg.content = _inject_attachment_tags(ctx.msg.content, ctx.msg.attachments) return ctx def _build_cached_attachment_tags(cached: list) -> str: tags: list[str] = [] for att in cached: ct = (att.content_type or "").lower() if ct.startswith("image/"): tag = f"[图片: {att.url}]" elif ct.startswith("video/"): tag = f"[视频: {att.url}]" elif ct.startswith("audio/") or ct.startswith("voice/"): tag = f"[语音: {att.url}]" else: label = att.filename or att.url tag = f"[文件: {label}]" tags.append(tag) return "\n".join(tags) def _inject_attachment_tags(content: str, attachments: list) -> str: tags: list[str] = [] for att in attachments: ct = (att.content_type or "").lower() if ct.startswith("image/"): tag = f"[图片: {att.url}]" elif ct.startswith("video/"): tag = f"[视频: {att.url}]" elif ct.startswith("audio/") or ct.startswith("voice/"): tag = f"[语音: {att.url}]" else: label = att.filename or att.url tag = f"[文件: {label}]" tags.append(tag) if tags: tag_section = "\n".join(tags) content = f"{content}\n{tag_section}" if content else tag_section return content class GroupGateStage: def __init__(self, gate_fn: Any | None = None, history_fn: Any | None = None): self._gate_fn = gate_fn self._history_fn = history_fn async def __call__(self, ctx: PipelineContext) -> PipelineContext: if ctx.msg.chat_type not in (QQBotChatType.GROUP, QQBotChatType.GUILD): return ctx if self._gate_fn: decision = await self._gate_fn(ctx) if decision and decision != "process": ctx.blocked = True ctx.skip_reason = decision if self._history_fn and decision in ("skip_no_mention", "drop_other_mention"): await self._history_fn(ctx) return ctx if self._history_fn and not ctx.blocked: await self._history_fn(ctx) return ctx class AssemblyStage: async def __call__(self, ctx: PipelineContext) -> PipelineContext: if ctx.msg.sender_name and ctx.msg.chat_type == QQBotChatType.GROUP: ctx.metadata["agent_content"] = f"[{ctx.msg.sender_name}]: {ctx.msg.content}" else: ctx.metadata["agent_content"] = ctx.msg.content return ctx class EnvelopeStage: def __init__(self, dispatch_fn: Any | None = None): self._dispatch_fn = dispatch_fn async def __call__(self, ctx: PipelineContext) -> PipelineContext: if ctx.blocked or not self._dispatch_fn: return ctx await self._dispatch_fn(ctx) return ctx class RefIndexStage: def __init__(self, index_fn: Any | None = None): self._index_fn = index_fn async def __call__(self, ctx: PipelineContext) -> PipelineContext: if self._index_fn: await self._index_fn(ctx) return ctx def create_default_pipeline( dispatch_fn=None, access_check_fn=None, gate_fn=None, history_fn=None, ref_index_fn=None, attachment_cache: AttachmentCache | None = None, ) -> QQBotPipeline: pipeline = QQBotPipeline() pipeline.add_stage(DedupeStage()) pipeline.add_stage(AccessStage(check_fn=access_check_fn)) pipeline.add_stage(ContentStage(attachment_cache=attachment_cache)) pipeline.add_stage(GroupGateStage(gate_fn=gate_fn, history_fn=history_fn)) pipeline.add_stage(AssemblyStage()) pipeline.add_stage(RefIndexStage(index_fn=ref_index_fn)) pipeline.add_stage(EnvelopeStage(dispatch_fn=dispatch_fn)) return pipeline