ForcePilot/backend/package/yuxi/channels/adapters/qqbot/streaming.py

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")