ForcePilot/backend/package/yuxi/channels/adapters/wechat/streaming_fallback.py

from __future__ import annotations

import re
from typing import TYPE_CHECKING

from yuxi.channels.models import DeliveryResult
from yuxi.utils.logging_config import logger

if TYPE_CHECKING:
    from .adapter import WeChatAdapter

_PARAGRAPH_SPLITTER = re.compile(r"\n{2,}")
_SENTENCE_BOUNDARY = re.compile(r"[。！？.!?\n]")


class ParagraphChunker:
    def __init__(self, chunk_size: int = 2000, min_chunk_size: int = 100):
        self.chunk_size = chunk_size
        self.min_chunk_size = min_chunk_size
        self._buffer: list[str] = []

    def feed(self, text: str) -> list[str]:
        self._buffer.append(text)
        accumulated = "".join(self._buffer)
        results: list[str] = []

        if "\n\n" in accumulated:
            paragraphs = _PARAGRAPH_SPLITTER.split(accumulated)
            if len(paragraphs) > 1:
                for para in paragraphs[:-1]:
                    stripped = para.strip()
                    if len(stripped) >= self.min_chunk_size:
                        results.append(stripped)
                self._buffer = [paragraphs[-1]]

        return results

    def flush(self) -> list[str]:
        if not self._buffer:
            return []
        accumulated = "".join(self._buffer)
        self._buffer = []
        if not accumulated.strip():
            return []

        if len(accumulated) <= self.chunk_size:
            return [accumulated.strip()]

        results = []
        remaining = accumulated
        while len(remaining) > self.chunk_size:
            split_at = self._find_split_point(remaining, self.chunk_size)
            results.append(remaining[:split_at].strip())
            remaining = remaining[split_at:].strip()
        if remaining.strip():
            results.append(remaining.strip())
        return results

    def _find_split_point(self, text: str, limit: int) -> int:
        text_slice = text[:limit]
        for sep in ("\n\n", "\n", "。", "！", "？", ". ", "! ", "? "):
            idx = text_slice.rfind(sep)
            if idx > limit // 2:
                return idx + len(sep)
        return limit


class ProgressPrefixFormatter:
    def __init__(self, template: str = "({current}/{total})", enabled: bool = True):
        self._template = template
        self.enabled = enabled

    def format(self, current: int, total: int) -> str:
        if not self.enabled:
            return ""
        return self._template.format(current=current, total=total)


async def send_stream_chunked(
    adapter: WeChatAdapter,
    chat_id: str,
    text: str,
    finished: bool = False,
) -> DeliveryResult:
    from yuxi.channels.models import ChannelIdentity, ChannelResponse

    identity = ChannelIdentity(
        channel_id=adapter.channel_id,
        channel_type=adapter.channel_type,
        channel_chat_id=chat_id,
        channel_user_id="",
    )

    response = ChannelResponse(identity=identity, content=text)
    try:
        return await adapter.send(response)
    except Exception as e:
        logger.error(f"[WeChat] Chunked send failed for {chat_id}: {e}")
        return DeliveryResult(success=False, error=str(e))


async def send_with_typing(
    adapter: WeChatAdapter,
    chat_id: str,
    text: str,
    typing_enabled: bool = True,
    is_wecom: bool = False,
) -> DeliveryResult:
    if typing_enabled and is_wecom:
        try:
            await adapter.heartbeat_adapter.send_typing_indicator(adapter.send, chat_id)
        except Exception:
            pass

    from yuxi.channels.models import ChannelIdentity, ChannelResponse

    identity = ChannelIdentity(
        channel_id=adapter.channel_id,
        channel_type=adapter.channel_type,
        channel_chat_id=chat_id,
        channel_user_id="",
    )

    response = ChannelResponse(identity=identity, content=text)
    try:
        return await adapter.send(response)
    except Exception as e:
        return DeliveryResult(success=False, error=str(e))