ForcePilot/backend/package/yuxi/channel/extensions/wechat-kf/streaming.py

import asyncio
import random


class WeChatKFStreaming:
    streaming_mode = "block"
    preview_stream_throttle_ms = 160
    preview_min_initial_chars = 18

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

    ENABLED = True
    STRATEGY = "block"
    MIN_CHARS = 200
    MAX_INTERVAL_MS = 500
    MAX_BLOCKS = 5
    FINISH_MARK = "[↓]"

    def __init__(self, outbound=None):
        self._outbound = outbound

    def create_draft_stream_session(self, target_id: str) -> object:
        return {"target_id": target_id, "mode": "block"}

    def create_block_chunker(self) -> object:
        return {"mode": "length", "min_chars": 200, "max_chars": 1200}

    @classmethod
    def should_stream(cls, estimated_length: int) -> bool:
        return estimated_length > cls.MIN_CHARS

    @classmethod
    def build_chunk(cls, text: str, is_final: bool) -> str:
        if is_final:
            return text
        return f"{text}\n{cls.FINISH_MARK}"

    async def send_block_stream(self, external_user_id: str, open_kfid: str, full_text: str, session_manager) -> list:
        if not self._outbound:
            return []

        MAX_CHUNK_CHARS = 1000
        COALESCE_THRESHOLD = 500
        MIN_DELAY_MS = 300
        MAX_DELAY_MS = 800

        session = session_manager.get_session(open_kfid, external_user_id)
        remaining = 5 - (session.msg_count_in_round if session else 0)

        chunks = self._split_content(full_text, max_chunks=remaining)
        results = []

        for chunk in chunks:
            delay = random.uniform(MIN_DELAY_MS / 1000, MAX_DELAY_MS / 1000)
            await asyncio.sleep(delay)
            result = await self._outbound.send_text(external_user_id, open_kfid, chunk)
            results.append(result)

        return results

    def _split_content(self, text: str, max_chunks: int = 5) -> list[str]:
        if not text:
            return []

        char_count = len(text)
        MAX_CHUNK_CHARS = 1000
        COALESCE_THRESHOLD = 500

        if char_count <= COALESCE_THRESHOLD:
            return [text]

        if max_chunks <= 1:
            return [text[:MAX_CHUNK_CHARS]]

        paragraphs = text.split("\n\n")
        chunks = []
        current = ""

        for para in paragraphs:
            if len(current) + len(para) + 1 <= MAX_CHUNK_CHARS:
                current = f"{current}\n\n{para}" if current else para
            else:
                if current:
                    chunks.append(current)
                    if len(chunks) >= max_chunks - 1:
                        remaining_text = "\n\n".join([current, para] + paragraphs[paragraphs.index(para) + 1 :])
                        chunks.append(remaining_text[:MAX_CHUNK_CHARS])
                        return chunks
                current = para

        if current:
            chunks.append(current)

        if len(chunks) > max_chunks:
            merged = ""
            for chunk in chunks[max_chunks - 1 :]:
                merged += chunk
            chunks = chunks[: max_chunks - 1] + [merged[:MAX_CHUNK_CHARS]]

        return chunks