from __future__ import annotations import asyncio import logging from collections.abc import AsyncGenerator, Callable from dataclasses import dataclass, field from typing import Any from yuxi.channels.adapters.qqbot.c2c_stream import ( C2CStreamingController, FlushController, FlushStrategy, ) logger = logging.getLogger(__name__) class ParagraphChunker: def __init__(self, flush_per_paragraph: bool = True, min_chunk_size: int = 10): self.flush_per_paragraph = flush_per_paragraph self.min_chunk_size = min_chunk_size self._buffer: list[str] = [] def feed(self, text: str) -> list[str]: results: list[str] = [] self._buffer.append(text) if not self.flush_per_paragraph: return results accumulated = "".join(self._buffer) if "\n\n" in accumulated: paragraphs = accumulated.split("\n\n") if len(paragraphs) > 1: for para in paragraphs[:-1]: if len(para.strip()) >= self.min_chunk_size: results.append(para + "\n\n") self._buffer = [paragraphs[-1]] return results def flush(self) -> str: if not self._buffer: return "" result = "".join(self._buffer) self._buffer = [] return result async def stream_content( content_generator: AsyncGenerator[str, None], c2c_ctrl: C2CStreamingController | None, chunker: ParagraphChunker | None = None, ) -> AsyncGenerator[str, None]: if chunker is None: chunker = ParagraphChunker() try: async for chunk in content_generator: if not chunk: continue paragraphs = chunker.feed(chunk) for para in paragraphs: yield para if c2c_ctrl is not None: c2c_batches = c2c_ctrl.flush_controller.feed(chunk) for batch in c2c_batches: pass await asyncio.sleep(0) remaining = chunker.flush() if remaining: yield remaining except asyncio.CancelledError: logger.debug("Streaming cancelled") remaining = chunker.flush() if remaining: yield remaining except Exception: logger.exception("Streaming error") remaining = chunker.flush() if remaining: yield remaining __all__ = [ "ParagraphChunker", "stream_content", "send_blocks_stream", "C2CStreamingController", "FlushController", "FlushStrategy", "MediaAwareStreamer", "StreamMediaContext", ] @dataclass class StreamMediaContext: stream_active: bool = True media_queue: list[tuple[str, str]] = field(default_factory=list) interrupt_count: int = 0 def interrupt(self) -> None: self.stream_active = False self.interrupt_count += 1 def restore(self) -> None: self.stream_active = True class MediaAwareStreamer: def __init__( self, send_text_fn: Callable[..., Any], send_media_fn: Callable[..., Any], chat_id: str = "", max_interrupts: int = 10, ): self._send_text_fn = send_text_fn self._send_media_fn = send_media_fn self._chat_id = chat_id self._max_interrupts = max_interrupts self._context = StreamMediaContext() from .media_tags import parse_media_tags self._parse_media_tags = parse_media_tags @property def context(self) -> StreamMediaContext: return self._context async def feed(self, chunk: str) -> None: from .media_tags import has_media_tags if not has_media_tags(chunk): if self._context.stream_active: await self._send_text_fn(self._chat_id, chunk) return if self._context.interrupt_count >= self._max_interrupts: clean = self._parse_media_tags(chunk).text if clean: await self._send_text_fn(self._chat_id, clean) return parsed = self._parse_media_tags(chunk) if parsed.text: self._context.interrupt() await self._send_text_fn(self._chat_id, parsed.text) self._context.restore() for item in parsed.media_items: try: await self._send_media_fn( self._chat_id, media_type=item.media_type, reference=item.reference, is_url=item.is_url, ) except Exception: logger.exception("MediaAwareStreamer: failed to send media %s", item) async def flush(self) -> None: pass async def stream_with_media_handling( content_generator: AsyncGenerator[str, None], send_text_fn: Callable[..., Any], send_media_fn: Callable[..., Any], chat_id: str = "", c2c_ctrl: C2CStreamingController | None = None, chunker: ParagraphChunker | None = None, max_interrupts: int = 10, ) -> AsyncGenerator[str, None]: if chunker is None: chunker = ParagraphChunker() media_streamer = MediaAwareStreamer( send_text_fn=send_text_fn, send_media_fn=send_media_fn, chat_id=chat_id, max_interrupts=max_interrupts, ) try: async for chunk in content_generator: if not chunk: continue await media_streamer.feed(chunk) paragraphs = chunker.feed(chunk) for para in paragraphs: yield para if c2c_ctrl is not None: c2c_batches = c2c_ctrl.flush_controller.feed(chunk) for _batch in c2c_batches: pass await asyncio.sleep(0) remaining = chunker.flush() if remaining: yield remaining except asyncio.CancelledError: logger.debug("Streaming cancelled") remaining = chunker.flush() if remaining: yield remaining except Exception: logger.exception("Streaming error") remaining = chunker.flush() if remaining: yield remaining async def send_blocks_stream( chat_id: str, text: str, send_fn, channel_id: str = "qqbot", channel_type=None, chunk_size: int = 1, parallelism: int = 1, ) -> None: from yuxi.channels.models import ( ChannelIdentity, ChannelResponse, ChannelType, DeliveryResult, ) ct = channel_type or ChannelType.QQ_BOT identity = ChannelIdentity( channel_id=channel_id, channel_type=ct, channel_user_id="", channel_chat_id=chat_id, ) paragraphs = text.split("\n\n") for para in paragraphs: if not para.strip(): continue response = ChannelResponse(identity=identity, content=para) try: result = await send_fn(response) if isinstance(result, DeliveryResult) and not result.success: logger.warning("send_blocks_stream: failed to send para: %s", result.error) except Exception: logger.exception("send_blocks_stream: error sending paragraph")