ForcePilot/backend/package/yuxi/channels/infra/text_chunker.py

from __future__ import annotations

import re
from collections.abc import Iterator
from dataclasses import dataclass
from enum import StrEnum


class ChunkMode(StrEnum):
    LENGTH = "length"
    NEWLINE = "newline"
    SENTENCE = "sentence"
    MARKDOWN = "markdown"


@dataclass
class TextChunk:
    text: str
    index: int
    is_last: bool


_NATURAL_BOUNDARIES = [
    "\n\n",
    "\n",
    "。",
    "！",
    "？",
    "；",
    "：",
    ". ",
    "! ",
    "? ",
]

_MD_FENCE = re.compile(r"^```")
_MD_TABLE_SEP = re.compile(r"^\|[-:| ]+\|")
_MD_ORDERED_LIST = re.compile(r"^\d+\.\s")
_MD_UNORDERED_LIST = re.compile(r"^[-*+]\s")
_MD_HEADING = re.compile(r"^#{1,6}\s")


class TextChunker:
    def __init__(self, max_length: int = 2000, mode: ChunkMode = ChunkMode.NEWLINE):
        self.max_length = max_length
        self.mode: ChunkMode = mode

    def chunk(self, text: str) -> list[str]:
        if len(text) <= self.max_length:
            return [text] if text else []
        if self.mode == ChunkMode.LENGTH:
            return [c for c in self._chunk_by_length(text) if c]
        if self.mode == ChunkMode.SENTENCE:
            return [c for c in self._chunk_by_sentence(text) if c]
        if self.mode == ChunkMode.MARKDOWN:
            return [c for c in self._chunk_markdown(text) if c]
        return [c for c in self._chunk_by_newline(text) if c]

    def chunk_with_index(self, text: str) -> Iterator[tuple[int, int, str]]:
        chunks = self.chunk(text)
        total = len(chunks)
        for i, chunk in enumerate(chunks):
            yield i + 1, total, chunk

    def chunk_as_objects(self, text: str) -> list[TextChunk]:
        chunks = self.chunk(text)
        return [TextChunk(text=c, index=i, is_last=(i == len(chunks) - 1)) for i, c in enumerate(chunks)]

    def _chunk_by_length(self, text: str) -> list[str]:
        chunks: list[str] = []
        remaining = text
        while len(remaining) > self.max_length:
            split_at = self._find_split_point(remaining)
            chunks.append(remaining[:split_at].rstrip())
            remaining = remaining[split_at:].lstrip()
        if remaining:
            chunks.append(remaining)
        return chunks

    def _chunk_by_newline(self, text: str) -> list[str]:
        paragraphs = re.split(r"\n{2,}", text)
        chunks: list[str] = []
        for para in paragraphs:
            if len(para) <= self.max_length:
                if para:
                    chunks.append(para)
                continue
            for sub in self._chunk_by_length(para):
                chunks.append(sub)
        return chunks

    def _chunk_by_sentence(self, text: str) -> list[str]:
        sentences = re.split(r"(?<=[。！？.!?])\s*", text)
        chunks: list[str] = []
        buffer = ""
        for sentence in sentences:
            if not sentence.strip():
                buffer += sentence
                continue
            candidate = buffer + sentence
            if len(candidate) <= self.max_length:
                buffer = candidate
            else:
                if buffer:
                    chunks.append(buffer)
                if len(sentence) > self.max_length:
                    for sub in self._chunk_by_length(sentence):
                        chunks.append(sub)
                    buffer = ""
                else:
                    buffer = sentence
        if buffer:
            chunks.append(buffer)
        return chunks

    def _chunk_markdown(self, text: str) -> list[str]:
        segments = self._segment_markdown(text)
        chunks: list[str] = []
        buffer = ""
        for seg in segments:
            candidate = buffer + ("\n" if buffer else "") + seg if buffer else seg
            if len(candidate) <= self.max_length:
                buffer = candidate
            else:
                if buffer:
                    chunks.append(buffer)
                if len(seg) > self.max_length:
                    for sub in self._chunk_overlong_md(seg):
                        chunks.append(sub)
                    buffer = ""
                else:
                    buffer = seg
        if buffer:
            chunks.append(buffer)
        return chunks

    def _segment_markdown(self, text: str) -> list[str]:
        lines = text.split("\n")
        segments: list[str] = []
        current: list[str] = []
        in_fence = False
        in_table = False
        in_list = False

        def flush():
            nonlocal current, in_table, in_list
            if current:
                segments.append("\n".join(current))
                current = []
            in_table = False
            in_list = False

        for line in lines:
            is_fence = bool(_MD_FENCE.match(line))
            is_table_sep = bool(_MD_TABLE_SEP.match(line))
            is_ordered = bool(_MD_ORDERED_LIST.match(line))
            is_unordered = bool(_MD_UNORDERED_LIST.match(line))
            is_heading = bool(_MD_HEADING.match(line))

            if is_fence:
                if in_fence:
                    current.append(line)
                    flush()
                    in_fence = False
                    continue
                flush()
                in_fence = True
                current.append(line)
                continue

            if in_fence:
                current.append(line)
                continue

            if is_table_sep and current and "|" in current[-1]:
                flush()
                in_table = True
                current.append(line)
                continue

            if in_table and "|" in line:
                current.append(line)
                continue

            if in_table:
                flush()

            if is_ordered or is_unordered:
                if not in_list:
                    flush()
                    in_list = True
                current.append(line)
                continue

            if is_heading:
                flush()
                current.append(line)
                continue

            if in_list:
                flush()

            if line.strip() == "":
                flush()
                continue

            current.append(line)

            if len("\n".join(current)) > self.max_length * 2:
                flush()

        flush()
        return segments

    def _chunk_overlong_md(self, text: str) -> list[str]:
        lines = text.split("\n")
        chunks: list[str] = []
        buffer = ""

        for line in lines:
            candidate = buffer + ("\n" if buffer else "") + line if buffer else line
            if len(candidate) <= self.max_length:
                buffer = candidate
                continue

            if buffer:
                chunks.append(buffer)
                buffer = ""

            if len(line) <= self.max_length:
                buffer = line
            else:
                for sub in self._chunk_by_length(line):
                    if len(sub) <= self.max_length:
                        if buffer:
                            chunks.append(buffer)
                        buffer = sub
                    else:
                        if buffer:
                            chunks.append(buffer)
                            buffer = ""
                        chunks.append(sub)

        if buffer:
            chunks.append(buffer)
        return chunks

    def _find_split_point(self, text: str) -> int:
        limit = self.max_length
        slice_end = min(limit, len(text))
        for sep in _NATURAL_BOUNDARIES:
            idx = text.rfind(sep, 0, slice_end)
            if idx > slice_end // 2:
                return idx + len(sep.rstrip())
        space_idx = text.rfind(" ", 0, slice_end)
        if space_idx > slice_end // 2:
            return space_idx + 1
        while slice_end > 0 and (ord(text[slice_end - 1]) & 0xC0) == 0x80:
            slice_end -= 1
        return slice_end if slice_end > 0 else limit