agent-management/backend/embedding-bridge/backend/indexing/text_splitter.py

"""文档三级分块服务（接收已解析文本，不直接读取文件）"""
import re
import unicodedata
from typing import Dict, List

from backend.indexing.semantic_chunker import SemanticChunker

# 编译非打印 C0/C1 控制字符的正则（保留常规排版字：\t, \n, \r）
_CONTROL_CHAR_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
# 编译零宽字符和不可见格式化控制字符（零宽空白、BOM 标记、左右强排标志等）
_INVISIBLE_CHAR_RE = re.compile(r"[\u200b-\u200d\ufeff\u200f\u202a-\u202e]")


def sanitize_text(text: str) -> str:
    """
    企业级标准文本净化器 (Text Sanitizer)。
    1. 规范化 (Normalization)：统一转换为标准 NFC 格式。
    2. 剔除/替换不合法及不可见字节：过滤 NUL、零宽字符、BOM 标签等。
    3. 清洗非打印字符及乱码：剔除 C0/C1 控制符号，剥离 Unicode PUA 私有使用区乱码符号。
    4. 编码收敛防爆：利用 utf-8 ignore 安全剥离孤立的 UTF-16 代理项。
    """
    if not text:
        return ""

    text = unicodedata.normalize("NFC", text)
    text = _INVISIBLE_CHAR_RE.sub("", text)
    text = _CONTROL_CHAR_RE.sub("", text)
    text = re.sub(r"[\ue000-\uf8ff]", "", text)

    try:
        cleaned = text.encode("utf-8", "ignore").decode("utf-8", "ignore")
    except Exception:
        chars = []
        for char in text:
            if 0xD800 <= ord(char) <= 0xDFFF:
                continue
            chars.append(char)
        cleaned = "".join(chars)

    return cleaned


# 分隔符按语义粒度分组，越靠前越优先作为分块边界。
# 切分时保留分隔符作为独立原子，保证拼接后原文不变。
_SENTENCE_SEPARATORS = ["\n\n", "。", "！", "？", "；", "\n"]
_CLAUSE_SEPARATORS = ["，", "、"]
_WORD_SEPARATORS = [" ", ""]

DEFAULT_SEPARATORS = _SENTENCE_SEPARATORS + _CLAUSE_SEPARATORS + _WORD_SEPARATORS


class SemanticTextSplitter:
    """
    语义优先的文本分块器。

    核心保证：
    1. 分块后所有 chunk 按顺序拼接，等于原始输入。
    2. 每个 chunk 的末尾尽量落在语义边界（段落 > 句子 > 子句 > 词）。
    3. 重叠区只在边界处截取，不会从句子中间截断。
    """

    def __init__(
        self,
        chunk_size: int = 800,
        chunk_overlap: int = 100,
        separators: List[str] | None = None,
        use_semantic_chunking: bool = False,
        semantic_similarity_threshold: float = 0.6,
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = separators or DEFAULT_SEPARATORS
        self.use_semantic_chunking = use_semantic_chunking
        self.semantic_chunker = None
        if use_semantic_chunking:
            # overlap 在语义段落边界处由 SemanticChunker 处理，保证承上启下句子完整
            self.semantic_chunker = SemanticChunker(
                similarity_threshold=semantic_similarity_threshold,
                max_chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
            )

    def split_text(self, text: str) -> List[str]:
        """对单段文本执行分块。"""
        text = sanitize_text(text)
        if not text:
            return []
        if len(text) <= self.chunk_size:
            return [text]

        # 若启用语义分块，先按话题边界粗分，再在每个语义段落内做规则细分。
        # 保证所有段落拼接仍等于原文。
        if self.use_semantic_chunking and self.semantic_chunker is not None:
            semantic_chunks = self.semantic_chunker.split_text(text)
            result: List[str] = []
            for chunk in semantic_chunks:
                if len(chunk) <= self.chunk_size:
                    result.append(chunk)
                else:
                    atoms = self._split_into_atoms(chunk)
                    result.extend(self._merge_atoms(atoms))
            return result

        # 未启用语义分块：按标点规则切分。
        atoms = self._split_into_atoms(text)
        return self._merge_atoms(atoms)

    def _split_into_atoms(self, text: str) -> List[str]:
        """
        递归地在最粗可用的分隔符处切分文本，保留分隔符作为独立原子。
        例如 "。" 切分后，标点本身会作为一个原子项，保证后续拼接不丢标点。
        """
        return self._split_recursive(text, self.separators.copy())

    def _split_recursive(self, text: str, separators: List[str]) -> List[str]:
        if len(text) <= self.chunk_size or not separators:
            return [text] if text else []

        separator = separators.pop(0)
        if separator == "":
            # 兜底：按字符切分
            return list(text)

        parts = text.split(separator)
        if len(parts) <= 1:
            return self._split_recursive(text, separators)

        result: List[str] = []
        for i, part in enumerate(parts):
            if i > 0:
                # 分隔符本身作为独立原子保留
                result.append(separator)
            if not part:
                continue
            if len(part) > self.chunk_size:
                result.extend(self._split_recursive(part, separators.copy()))
            else:
                result.append(part)

        return result

    def _merge_atoms(self, atoms: List[str]) -> List[str]:
        """
        把原子片段合并成接近 chunk_size 的块。

        合并规则：
        1. 原子不可拆分；
        2. 中间 chunk 的末尾尽量落在语义分隔符上；
        3. 因回退而移出当前 chunk 的非边界原子，会保留到下一个 chunk 中，保证原文不丢失；
        4. 若启用 overlap，下一个 chunk 以当前 chunk 末尾的边界感知重叠区开头。
        """
        if not atoms:
            return []

        chunks: List[str] = []
        current: List[str] = []
        current_len = 0

        def last_sep_index(seq: List[str]) -> int:
            for idx in range(len(seq) - 1, -1, -1):
                if seq[idx] in self.separators:
                    return idx
            return -1

        for atom in atoms:
            atom_len = len(atom)

            # 加入当前原子会超过 chunk_size，则先结束当前 chunk
            if current and current_len + atom_len > self.chunk_size:
                sep_idx = last_sep_index(current)
                if sep_idx >= 0:
                    chunk_atoms = current[: sep_idx + 1]
                    pending = current[sep_idx + 1 :] + [atom]
                else:
                    # 当前窗口内没有可用分隔符，直接结束整个窗口
                    chunk_atoms = current
                    pending = [atom]

                chunks.append("".join(chunk_atoms))

                if self.chunk_overlap > 0 and chunk_atoms:
                    overlap_atoms = self._extract_overlap_atoms(chunk_atoms)
                    current = overlap_atoms + pending
                else:
                    current = pending
                current_len = sum(len(a) for a in current)
                continue

            current.append(atom)
            current_len += atom_len

        if current:
            chunks.append("".join(current))

        return chunks

    def _extract_overlap_atoms(self, atoms: List[str]) -> List[str]:
        """
        从一组原子末尾回退，截取不超过 chunk_overlap 且以语义边界结尾的重叠区。
        优先让整个重叠区以句子或子句分隔符结尾。
        """
        if self.chunk_overlap <= 0 or not atoms:
            return []

        overlap: List[str] = []
        length = 0
        for atom in reversed(atoms):
            overlap.insert(0, atom)
            length += len(atom)
            if length >= self.chunk_overlap:
                break

        # 让 overlap 以语义边界结束
        while len(overlap) > 1:
            if overlap[-1] in self.separators:
                break
            removed_len = len(overlap[-1])
            if length - removed_len >= self.chunk_overlap // 3:
                length -= removed_len
                overlap.pop()
            else:
                break

        return overlap

    @staticmethod
    def _atoms_length(atoms: List[str]) -> int:
        return sum(len(a) for a in atoms)


class HierarchicalTextSplitter:
    """三级层次化分块（L1 / L2 / L3）"""

    def __init__(self, chunk_size: int = 800, chunk_overlap: int = 100):
        # L1 是对原文的粗粒度切分，L1 之间不允许 overlap，保证所有 L1 拼接等于原文。
        # L2/L3 在父级内部进行切分，启用基于 embedding 相似度的语义边界；
        # 允许 L2/L3 在语义边界处保留 overlap，使承上启下的句子同时出现在前后块中。
        level_1_size = max(2000, chunk_size * 3)
        level_2_size = max(1000, chunk_size * 2)
        level_3_size = max(600, chunk_size)

        self._splitter_level_1 = SemanticTextSplitter(
            chunk_size=level_1_size,
            chunk_overlap=0,
            separators=DEFAULT_SEPARATORS,
            use_semantic_chunking=False,
        )
        self._splitter_level_2 = SemanticTextSplitter(
            chunk_size=level_2_size,
            chunk_overlap=chunk_overlap,
            separators=DEFAULT_SEPARATORS,
            use_semantic_chunking=True,
        )
        self._splitter_level_3 = SemanticTextSplitter(
            chunk_size=level_3_size,
            chunk_overlap=chunk_overlap,
            separators=DEFAULT_SEPARATORS,
            use_semantic_chunking=True,
        )

    @staticmethod
    def _build_chunk_id(document_id: int, level: int, index: int) -> str:
        return f"{document_id}::l{level}::{index}"

    def split_text(
        self,
        text: str,
        document_id: int,
        filename: str = "",
        file_type: str = "",
        file_path: str = "",
        page_number: int = 0,
    ) -> List[Dict]:
        """对单段文本执行三级分块，返回 L1/L2/L3 全部分块。"""
        text = sanitize_text(text)
        if not text:
            return []

        base_doc = {
            "filename": sanitize_text(filename),
            "file_path": sanitize_text(file_path),
            "file_type": sanitize_text(file_type),
            "page_number": int(page_number),
        }

        root_chunks: List[Dict] = []
        level_1_counter = 0
        level_2_counter = 0
        level_3_counter = 0
        chunk_idx = 0

        level_1_texts = self._splitter_level_1.split_text(text)
        for level_1_text in level_1_texts:
            if not level_1_text:
                continue
            level_1_id = self._build_chunk_id(document_id, 1, level_1_counter)
            level_1_counter += 1

            level_1_chunk = {
                **base_doc,
                "text": level_1_text,
                "chunk_id": level_1_id,
                "parent_chunk_id": "",
                "root_chunk_id": level_1_id,
                "chunk_level": 1,
                "chunk_idx": chunk_idx,
            }
            chunk_idx += 1
            root_chunks.append(level_1_chunk)

            level_2_texts = self._splitter_level_2.split_text(level_1_text)
            for level_2_text in level_2_texts:
                if not level_2_text:
                    continue
                level_2_id = self._build_chunk_id(document_id, 2, level_2_counter)
                level_2_counter += 1

                level_2_chunk = {
                    **base_doc,
                    "text": level_2_text,
                    "chunk_id": level_2_id,
                    "parent_chunk_id": level_1_id,
                    "root_chunk_id": level_1_id,
                    "chunk_level": 2,
                    "chunk_idx": chunk_idx,
                }
                chunk_idx += 1
                root_chunks.append(level_2_chunk)

                level_3_texts = self._splitter_level_3.split_text(level_2_text)
                for level_3_text in level_3_texts:
                    if not level_3_text:
                        continue
                    level_3_id = self._build_chunk_id(document_id, 3, level_3_counter)
                    level_3_counter += 1
                    root_chunks.append({
                        **base_doc,
                        "text": level_3_text,
                        "chunk_id": level_3_id,
                        "parent_chunk_id": level_2_id,
                        "root_chunk_id": level_1_id,
                        "chunk_level": 3,
                        "chunk_idx": chunk_idx,
                    })
                    chunk_idx += 1

        return root_chunks