ForcePilot/backend/package/yuxi/channels/adapters/telegram/chunking.py

from __future__ import annotations

import re
from typing import Any, Literal

ChunkMode = Literal["length", "newline", "paragraph"]

TELEGRAM_MAX_TEXT_LENGTH = 4096
_SENTENCE_BOUNDARY = re.compile(r"[。！？.!?\n]")


def _chunk_by_length(text: str, limit: int) -> list[str]:
    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = min(start + limit, len(text))
        if end < len(text):
            split_at = _find_split_point(text, limit, start, end)
            chunks.append(text[start:split_at].strip())
            start = split_at
        else:
            chunks.append(text[start:end].strip())
            start = end
    return chunks or [text]


def _chunk_by_newline(text: str, limit: int) -> list[str]:
    chunks: list[str] = []
    lines = text.split("\n")
    current = ""
    for line in lines:
        if len(line) > limit:
            if current:
                chunks.append(current)
            sub_chunks = _chunk_by_length(line, limit)
            if sub_chunks:
                current = sub_chunks.pop()
                chunks.extend(sub_chunks)
            else:
                current = ""
        elif len(current) + len(line) + (1 if current else 0) <= limit:
            current = f"{current}\n{line}" if current else line
        else:
            chunks.append(current)
            current = line
    if current:
        chunks.append(current)
    return chunks or [text]


def _chunk_by_paragraph(text: str, limit: int) -> list[str]:
    if len(text) <= limit:
        return [text]

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

    for para in paragraphs:
        if len(current) + len(para) + 2 <= limit:
            current = f"{current}\n\n{para}" if current else para
        else:
            if current:
                chunks.append(current)
            if len(para) > limit:
                sub_chunks = _chunk_long_paragraph(para, limit)
                if sub_chunks:
                    if sub_chunks[-1]:
                        current = sub_chunks.pop()
                    else:
                        sub_chunks.pop()
                        current = ""
                    chunks.extend(sub_chunks)
                else:
                    current = ""
            else:
                current = para

    if current:
        chunks.append(current)
    return chunks or [text]


def chunk_text(
    text: str,
    limit: int = TELEGRAM_MAX_TEXT_LENGTH,
    mode: ChunkMode = "paragraph",
) -> list[str]:
    modes: dict[ChunkMode, callable] = {
        "length": _chunk_by_length,
        "newline": _chunk_by_newline,
        "paragraph": _chunk_by_paragraph,
    }
    chunker = modes.get(mode, _chunk_by_paragraph)
    return chunker(text, limit)


def chunk_mode_from_config(config: dict[str, Any] | None = None) -> ChunkMode:
    cfg = config or {}
    streaming_cfg = cfg.get("streaming", {})
    if isinstance(streaming_cfg, dict):
        return streaming_cfg.get("chunk_mode", "paragraph")
    return cfg.get("chunking_mode", "paragraph")


def _chunk_long_paragraph(text: str, limit: int) -> list[str]:
    chunks: list[str] = []
    while len(text) > limit:
        split_at = _find_split_point(text, limit, 0, limit)
        chunks.append(text[:split_at].rstrip())
        text = text[split_at:].lstrip()
    chunks.append(text)
    return chunks


def _find_split_point(text: str, limit: int, lo: int = 0, hi: int | None = None) -> int:
    look_start = max(lo, lo or limit // 2)
    look_end = hi or limit
    candidates = [m.start() for m in _SENTENCE_BOUNDARY.finditer(text, look_start, look_end)]
    if candidates:
        return candidates[-1] + 1
    newline = text.rfind("\n", look_start, look_end)
    if newline != -1:
        return newline + 1
    space = text.rfind(" ", look_start, look_end)
    if space != -1:
        return space + 1
    return look_end