YG_FT/backend/app/modules/data_process/document_chunking.py

"""基于 Docling 与 LlamaIndex 的文档切分实现。"""

from __future__ import annotations

import os
import re
import threading
import unicodedata
from dataclasses import dataclass
from functools import lru_cache
from io import BytesIO
from typing import Any, Literal

import tiktoken
from docling_core.transforms.chunker.hierarchical_chunker import ChunkingSerializerProvider
from llama_index.core import Document
from llama_index.core.base.embeddings.base import BaseEmbedding
from llama_index.core.node_parser import SemanticSplitterNodeParser, SentenceSplitter

from app.modules.data_process.algorithms import normalize_text

ChunkMethod = Literal["layout_hybrid", "semantic", "fixed"]

_PAGE_FURNITURE = re.compile(
    r"(?m)^\s*(?:第\s*\d+\s*页\s*共\s*\d+\s*页|[-—–]?\s*\d+\s*[/／]\s*\d+\s*[-—–]?)\s*$"
)
_COMPACT_CHARACTER = re.compile(r"[\w\u3400-\u4dbf\u4e00-\u9fff]", re.UNICODE)
_CONVERTER_LOCK = threading.Lock()


@dataclass(frozen=True, slots=True)
class DocumentChunk:
    """切片正文及其在原文件中的可追溯信息。"""

    original_content: str
    contextualized_content: str
    source_start: int | None
    source_end: int | None
    source_start_line: int | None
    source_end_line: int | None
    token_count: int
    heading_path: tuple[str, ...] = ()
    source_pages: tuple[int, ...] = ()
    doc_item_refs: tuple[str, ...] = ()
    source_bboxes: tuple[dict[str, Any], ...] = ()


def _sentence_chunks(text: str) -> list[str]:
    """提供稳定的中英文句界，避免 LlamaIndex 默认分词器下载额外资源。"""

    boundary = re.compile(
        r".*?(?:\n\s*\n|[。！？!?；;](?:[\"'”’）】》]*)|\.(?:\s+|$)|$)",
        re.DOTALL,
    )
    return [part for part in boundary.findall(text) if part]


@lru_cache(maxsize=1)
def _tokenizer() -> tiktoken.Encoding:
    """加载 cl100k_base 编码器，优先在线下载，失败时使用本地缓存以支持离线环境。"""
    import os
    import base64

    # 先设置缓存目录环境变量
    offline_cache = os.path.expanduser("~/.cache/tiktoken")
    os.environ.setdefault("TIKTOKEN_CACHE_DIR", offline_cache)

    try:
        # 尝试标准方式加载
        return tiktoken.get_encoding("cl100k_base")
    except Exception:
        # 如果失败，尝试手动从本地文件构造
        try:
            from pathlib import Path

            local_file = Path(offline_cache) / "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
            if not local_file.exists():
                # 尝试另一个可能的文件名
                local_file = Path(offline_cache) / "cl100k_base.tiktoken"

            if local_file.exists():
                # 读取 BPE 文件内容
                with open(local_file, "rb") as f:
                    contents = f.read()

                # 解析 BPE 文件
                mergeable_ranks = {}
                for line in contents.splitlines():
                    if line:
                        token, rank = line.split()
                        mergeable_ranks[base64.b64decode(token)] = int(rank)

                # 构造 Encoding 对象（模块顶部已 import tiktoken，
                # 此处不能再 import tiktoken.core，否则会把 tiktoken
                # 变成局部变量，使函数开头的 tiktoken.get_encoding 抛
                # UnboundLocalError）
                return tiktoken.core.Encoding(
                    name="cl100k_base",
                    pat_str=r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""",
                    mergeable_ranks=mergeable_ranks,
                    special_tokens={
                        "<|endoftext|>": 100257,
                        "<|fim_prefix|>": 100258,
                        "<|fim_middle|>": 100259,
                        "<|fim_suffix|>": 100260,
                        "<|endofprompt|>": 100276,
                    },
                )
        except Exception:
            pass

        raise RuntimeError(
            f"无法加载 cl100k_base 编码器\n"
            f"请确保以下任一条件满足:\n"
            f"1. 服务器可以访问网络\n"
            f"2. 本地存在缓存文件: {offline_cache}/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
        )


def _text_chunks(
    text: str,
    *,
    chunk_size: int,
    chunk_overlap: int,
) -> list[DocumentChunk]:
    normalized = normalize_text(text)
    if not normalized:
        return []
    splitter = SentenceSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        tokenizer=_tokenizer().encode,
        chunking_tokenizer_fn=_sentence_chunks,
        include_metadata=False,
        include_prev_next_rel=False,
    )
    nodes = splitter.get_nodes_from_documents([Document(text=normalized)])
    return _nodes_to_chunks(nodes, normalized)


def chunk_fixed_text(
    text: str,
    *,
    chunk_size: int,
    chunk_overlap: int,
) -> list[DocumentChunk]:
    """使用 LlamaIndex SentenceSplitter 按句界控制固定 Token 长度。"""

    return _text_chunks(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)


@lru_cache(maxsize=1)
def _semantic_embedding_model() -> BaseEmbedding:
    # 模型可在部署环境覆盖；默认模型体积较小且适合中英文语义边界判断。
    from llama_index.embeddings.huggingface import HuggingFaceEmbedding

    return HuggingFaceEmbedding(
        model_name=os.getenv("DATA_PROCESS_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5"),
        device=os.getenv("DATA_PROCESS_EMBEDDING_DEVICE", "cpu"),
        trust_remote_code=False,
    )


def chunk_semantic_text(
    text: str,
    *,
    chunk_size: int,
    chunk_overlap: int,
    breakpoint_percentile_threshold: int,
    embed_model: BaseEmbedding | None = None,
) -> list[DocumentChunk]:
    """使用 LlamaIndex SemanticSplitter 识别主题跳变，再限制最大长度。"""

    normalized = normalize_text(text)
    if not normalized:
        return []
    splitter = SemanticSplitterNodeParser.from_defaults(
        embed_model=embed_model or _semantic_embedding_model(),
        breakpoint_percentile_threshold=breakpoint_percentile_threshold,
        buffer_size=1,
        sentence_splitter=_sentence_chunks,
        include_metadata=False,
        include_prev_next_rel=False,
    )
    semantic_nodes = splitter.get_nodes_from_documents([Document(text=normalized)])
    result: list[DocumentChunk] = []
    search_from = 0
    for node in semantic_nodes:
        content = node.get_content().strip()
        if not content:
            continue
        start = _locate_text(normalized, content, search_from)
        if start is None:
            start = _locate_text(normalized, content, 0)
        if start is None:
            continue
        if len(_tokenizer().encode(content)) <= chunk_size:
            result.append(_make_text_chunk(normalized, start, start + len(content)))
        else:
            for child in _text_chunks(
                content,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
            ):
                if child.source_start is None or child.source_end is None:
                    continue
                result.append(
                    _make_text_chunk(
                        normalized,
                        start + child.source_start,
                        start + child.source_end,
                    )
                )
        search_from = start + len(content)
    return result


def _nodes_to_chunks(nodes: list[Any], source_text: str) -> list[DocumentChunk]:
    chunks: list[DocumentChunk] = []
    search_from = 0
    for node in nodes:
        content = node.get_content().strip()
        if not content:
            continue
        raw_start = getattr(node, "start_char_idx", None)
        raw_end = getattr(node, "end_char_idx", None)
        if (
            isinstance(raw_start, int)
            and isinstance(raw_end, int)
            and source_text[raw_start:raw_end].strip() == content
        ):
            start = raw_start + len(source_text[raw_start:raw_end]) - len(source_text[raw_start:raw_end].lstrip())
        else:
            start = _locate_text(source_text, content, search_from)
            if start is None:
                start = _locate_text(source_text, content, 0)
        if start is None:
            continue
        end = start + len(content)
        chunks.append(_make_text_chunk(source_text, start, end))
        search_from = max(search_from, end)
    return chunks


def _locate_text(source: str, content: str, start: int) -> int | None:
    position = source.find(content, start)
    return position if position >= 0 else None


def _make_text_chunk(source: str, start: int, end: int) -> DocumentChunk:
    content = source[start:end]
    return DocumentChunk(
        original_content=content,
        contextualized_content=content,
        source_start=start,
        source_end=end,
        source_start_line=source.count("\n", 0, start) + 1,
        source_end_line=source.count("\n", 0, max(start, end - 1)) + 1,
        token_count=len(_tokenizer().encode(content)),
    )


@lru_cache(maxsize=1)
def _document_converter():
    from docling.document_converter import DocumentConverter

    return DocumentConverter()


class _MarkdownSerializerProvider(ChunkingSerializerProvider):
    def get_serializer(self, doc: Any):
        from docling_core.transforms.chunker.hierarchical_chunker import ChunkingDocSerializer
        from docling_core.transforms.serializer.markdown import (
            MarkdownParams,
            MarkdownTableSerializer,
        )
        from docling_core.types.doc import DocItemLabel

        excluded = {
            DocItemLabel.DOCUMENT_INDEX,
            DocItemLabel.PAGE_HEADER,
            DocItemLabel.PAGE_FOOTER,
        }
        return ChunkingDocSerializer(
            doc=doc,
            table_serializer=MarkdownTableSerializer(),
            params=MarkdownParams(
                labels=set(DocItemLabel) - excluded,
                compact_tables=True,
                image_placeholder="",
                escape_html=False,
                escape_underscores=False,
            ),
        )


def _clean_layout_text(value: str) -> str:
    return normalize_text(_PAGE_FURNITURE.sub("", value)).strip()


def _compact_with_offsets(value: str) -> tuple[str, list[int]]:
    compact: list[str] = []
    offsets: list[int] = []
    for index, character in enumerate(unicodedata.normalize("NFKC", value)):
        if _COMPACT_CHARACTER.fullmatch(character):
            compact.append(character.casefold())
            offsets.append(index)
    return "".join(compact), offsets


def _project_layout_span(
    source_text: str,
    content: str,
    *,
    compact_source: str,
    source_offsets: list[int],
    compact_start: int,
) -> tuple[int | None, int | None, int]:
    compact_content, _ = _compact_with_offsets(content)
    if len(compact_content) < 4:
        return None, None, compact_start
    position = compact_source.find(compact_content, compact_start)
    if position < 0:
        position = compact_source.find(compact_content)
    if position < 0:
        return None, None, compact_start
    start = source_offsets[position]
    end = source_offsets[position + len(compact_content) - 1] + 1
    while start > 0 and source_text[start - 1] not in "\r\n":
        start -= 1
    while end < len(source_text) and source_text[end] not in "\r\n":
        end += 1
    return start, end, position + len(compact_content)


def chunk_layout_document(
    raw: bytes,
    *,
    filename: str,
    source_text: str,
    chunk_size: int,
) -> list[DocumentChunk]:
    """使用 Docling HybridChunker 按版面层级、列表与表格边界切分。"""

    from docling.chunking import HybridChunker
    from docling.datamodel.base_models import DocumentStream
    from docling.exceptions import BaseError as DoclingError
    from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
    from docling_core.types.doc import DocItemLabel

    try:
        with _CONVERTER_LOCK:
            conversion = _document_converter().convert(
                DocumentStream(name=filename, stream=BytesIO(raw))
            )
    except DoclingError as exc:
        raise ValueError(f"文档版面解析失败: {exc}") from exc
    chunker = HybridChunker(
        tokenizer=OpenAITokenizer(tokenizer=_tokenizer(), max_tokens=chunk_size),
        serializer_provider=_MarkdownSerializerProvider(),
        merge_peers=True,
        repeat_table_header=True,
    )
    compact_source, source_offsets = _compact_with_offsets(source_text)
    compact_start = 0
    result: list[DocumentChunk] = []
    excluded = {
        DocItemLabel.DOCUMENT_INDEX,
        DocItemLabel.PAGE_HEADER,
        DocItemLabel.PAGE_FOOTER,
    }
    for raw_chunk in chunker.chunk(conversion.document):
        doc_items = tuple(raw_chunk.meta.doc_items or ())
        if doc_items and all(item.label in excluded for item in doc_items):
            continue
        content = _clean_layout_text(raw_chunk.text)
        if not content:
            continue
        contextualized = _clean_layout_text(chunker.contextualize(raw_chunk)) or content
        start, end, compact_start = _project_layout_span(
            source_text,
            content,
            compact_source=compact_source,
            source_offsets=source_offsets,
            compact_start=compact_start,
        )
        original = source_text[start:end] if start is not None and end is not None else content
        pages: set[int] = set()
        refs: list[str] = []
        bboxes: list[dict[str, Any]] = []
        for item in doc_items:
            refs.append(str(item.self_ref))
            for provenance in item.prov or ():
                pages.add(int(provenance.page_no))
                bbox = provenance.bbox
                bboxes.append(
                    {
                        "page": int(provenance.page_no),
                        "left": float(bbox.l),
                        "top": float(bbox.t),
                        "right": float(bbox.r),
                        "bottom": float(bbox.b),
                        "origin": str(bbox.coord_origin.value),
                    }
                )
        result.append(
            DocumentChunk(
                original_content=original,
                contextualized_content=contextualized,
                source_start=start,
                source_end=end,
                source_start_line=(source_text.count("\n", 0, start) + 1 if start is not None else None),
                source_end_line=(
                    source_text.count("\n", 0, max(start or 0, (end or 1) - 1)) + 1
                    if end is not None
                    else None
                ),
                token_count=len(_tokenizer().encode(contextualized)),
                heading_path=tuple(str(item) for item in (raw_chunk.meta.headings or ())),
                source_pages=tuple(sorted(pages)),
                doc_item_refs=tuple(refs),
                source_bboxes=tuple(bboxes),
            )
        )
    return result


def merge_short_chunks(
    chunks: list[DocumentChunk],
    *,
    source_text: str,
    min_token_count: int,
    max_token_count: int,
) -> list[DocumentChunk]:
    """在不突破长度上限的前提下，把过短块并入相邻内容。"""

    result: list[DocumentChunk] = []
    index = 0
    while index < len(chunks):
        current = chunks[index]
        if current.token_count >= min_token_count:
            result.append(current)
            index += 1
            continue
        if index + 1 < len(chunks):
            combined = _combine_chunks(current, chunks[index + 1], source_text)
            if combined.token_count <= max_token_count:
                result.append(combined)
                index += 2
                continue
        if result:
            combined = _combine_chunks(result[-1], current, source_text)
            if combined.token_count <= max_token_count:
                result[-1] = combined
                index += 1
                continue
        result.append(current)
        index += 1
    return result


def _combine_chunks(
    left: DocumentChunk,
    right: DocumentChunk,
    source_text: str,
) -> DocumentChunk:
    contextualized = "\n\n".join(
        part for part in (left.contextualized_content, right.contextualized_content) if part
    )
    start = left.source_start
    end = right.source_end
    has_contiguous_source = (
        start is not None
        and left.source_end is not None
        and right.source_start is not None
        and end is not None
        and left.source_end <= right.source_start
    )
    original = (
        source_text[start:end]
        if has_contiguous_source and start is not None and end is not None
        else "\n\n".join(
            part for part in (left.original_content, right.original_content) if part
        )
    )
    if not has_contiguous_source:
        start = None
        end = None
    return DocumentChunk(
        original_content=original,
        contextualized_content=contextualized,
        source_start=start,
        source_end=end,
        source_start_line=left.source_start_line if start is not None else None,
        source_end_line=right.source_end_line if end is not None else None,
        token_count=len(_tokenizer().encode(contextualized)),
        heading_path=left.heading_path or right.heading_path,
        source_pages=tuple(sorted(set(left.source_pages) | set(right.source_pages))),
        doc_item_refs=left.doc_item_refs + right.doc_item_refs,
        source_bboxes=left.source_bboxes + right.source_bboxes,
    )