ForcePilot/backend/package/yuxi/channel/extensions/telegram/streaming.py

from __future__ import annotations

import asyncio
import logging

logger = logging.getLogger(__name__)


class TelegramStreaming:
    streaming_mode = "partial"
    preview_stream_throttle_ms = 160
    preview_min_initial_chars = 18

    block_streaming_enabled = True
    block_streaming_break = "text_end"
    block_streaming_chunk_min_chars = 800
    block_streaming_chunk_max_chars = 1200
    block_streaming_chunk_break_preference = "paragraph"
    block_streaming_coalesce_defaults = {
        "min_chars": 400,
        "max_chars": 800,
        "idle_ms": 1000,
    }

    def create_draft_stream_session(self, target_id: str) -> dict:
        return {"target_id": target_id, "mode": "partial", "message_id": None}

    def create_block_chunker(self) -> dict:
        return {
            "mode": "length",
            "min_chars": self.block_streaming_chunk_min_chars,
            "max_chars": self.block_streaming_chunk_max_chars,
        }

    def create_preview_stream(
        self,
        outbound,
        target_id: str,
        account_id: str,
        *,
        thread_id: str | None = None,
        throttle_ms: int | None = None,
        min_initial_chars: int | None = None,
        chunk_limit: int = 4096,
    ):
        return PreviewStream(
            outbound=outbound,
            target_id=target_id,
            account_id=account_id,
            thread_id=thread_id,
            throttle_ms=throttle_ms or self.preview_stream_throttle_ms,
            min_initial_chars=min_initial_chars or self.preview_min_initial_chars,
            chunk_limit=chunk_limit,
        )


class PreviewStream:
    def __init__(
        self,
        outbound,
        target_id: str,
        account_id: str,
        *,
        thread_id: str | None = None,
        throttle_ms: int = 160,
        min_initial_chars: int = 18,
        chunk_limit: int = 4096,
    ):
        self._outbound = outbound
        self._target_id = target_id
        self._account_id = account_id
        self._thread_id = thread_id
        self._throttle_ms = throttle_ms
        self._min_initial_chars = min_initial_chars
        self._chunk_limit = chunk_limit

        self._buffer = ""
        self._draft_message_id: str | None = None
        self._sent_length = 0
        self._last_edit_time = 0.0
        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._done = asyncio.Event()
        self._final_text = ""

    async def feed(self, token: str) -> None:
        self._buffer += token
        await self._queue.put(token)

    async def finalize(self, final_text: str) -> None:
        self._final_text = final_text
        self._done.set()

        if self._draft_message_id:
            try:
                await self._outbound.edit_message(
                    self._target_id,
                    self._draft_message_id,
                    final_text,
                    account_id=self._account_id,
                )
            except Exception:
                logger.exception("PreviewStream finalize edit failed for msg %s", self._draft_message_id)
        else:
            try:
                await self._outbound.send_text_raw(
                    self._target_id,
                    final_text,
                    account_id=self._account_id,
                    thread_id=self._thread_id,
                )
            except Exception:
                logger.exception("PreviewStream finalize send failed")

    async def run(self) -> None:
        loop = asyncio.get_running_loop()

        while not self._done.is_set():
            try:
                await asyncio.wait_for(self._queue.get(), timeout=0.5)
            except TimeoutError:
                continue

            current_len = len(self._buffer)

            if self._draft_message_id is None:
                if current_len >= self._min_initial_chars:
                    try:
                        self._draft_message_id = await self._outbound.send_text_raw(
                            self._target_id,
                            self._buffer,
                            account_id=self._account_id,
                            thread_id=self._thread_id,
                        )
                        self._last_edit_time = loop.time()
                        self._sent_length = current_len
                    except Exception:
                        logger.exception("PreviewStream initial send failed")
                continue

            if current_len <= self._sent_length:
                continue

            now = loop.time()
            if now - self._last_edit_time >= self._throttle_ms / 1000:
                try:
                    display = self._buffer[:self._chunk_limit]
                    await self._outbound.edit_message(
                        self._target_id,
                        self._draft_message_id,
                        display,
                        account_id=self._account_id,
                    )
                    self._last_edit_time = now
                    self._sent_length = current_len
                except Exception:
                    pass