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

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

from __future__ import annotations

import logging
import re
import threading
import time
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
from app.modules.data_process.algorithms.embedding import semantic_embedding_model

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*$"
)
# Docling 的 markdown 序列化会给列表项补上自动编号，而 Word 的编号存放在
# numbering.xml 中，python-docx 抽取的正文不含这些编号；紧凑匹配前剥掉
# 行首编号，否则带列表的切片会整体定位失败。
_LIST_MARKER_PREFIX = re.compile(
    r"(?m)^[ \t>]*(?:(?:\d{1,3}[.)])+|\([a-zA-Z0-9]{1,3}\)|[a-zA-Z][.)]|[-*+•·])[ \t]+"
)
_COMPACT_CHARACTER = re.compile(r"[\w\u3400-\u4dbf\u4e00-\u9fff]", re.UNICODE)
_CONVERTER_LOCK = threading.Lock()

logger = logging.getLogger(__name__)


@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 base64

    from app.core.cache_paths import tiktoken_cache_dir

    # 缓存目录已在 app.core.cache_paths.setup_local_caches 中统一指向 <repo>/.cache/tiktoken，
    # 此处直接读取；TIKTOKEN_CACHE_DIR 已在启动阶段写入。
    offline_cache = tiktoken_cache_dir()

    try:
        # 尝试标准方式加载（环境变量 TIKTOKEN_CACHE_DIR 已被统一设置）
        return tiktoken.get_encoding("cl100k_base")
    except Exception:
        # 如果失败，尝试手动从本地文件构造
        try:
            local_file = offline_cache / "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
            if not local_file.exists():
                # 尝试另一个可能的文件名
                local_file = 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={
                        "": 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)


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:
            result.append(_unlocated_chunk(content))
            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:
                    result.append(_unlocated_chunk(child.original_content))
                    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:
            chunks.append(_unlocated_chunk(content))
            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 _unlocated_chunk(content: str) -> DocumentChunk:
    """正文在源文本中定位失败时保底保留切片，只放弃行号信息。"""

    return DocumentChunk(
        original_content=content,
        contextualized_content=content,
        source_start=None,
        source_end=None,
        source_start_line=None,
        source_end_line=None,
        token_count=len(_tokenizer().encode(content)),
    )


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.datamodel.base_models import InputFormat
    from docling.datamodel.pipeline_options import PdfPipelineOptions
    from docling.document_converter import DocumentConverter, PdfFormatOption

    pipeline_options = PdfPipelineOptions()
    pipeline_options.do_ocr = False
    return DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
        }
    )


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 _expand_to_line_boundaries(source_text: str, start: int, end: int) -> tuple[int, int]:
    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


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]:
    for candidate in (content, _LIST_MARKER_PREFIX.sub("", content)):
        compact_content, _ = _compact_with_offsets(candidate)
        if len(compact_content) < 4:
            continue
        position = compact_source.find(compact_content, compact_start)
        if position < 0:
            position = compact_source.find(compact_content)
        if position < 0:
            continue
        start, end = _expand_to_line_boundaries(
            source_text,
            source_offsets[position],
            source_offsets[position + len(compact_content) - 1] + 1,
        )
        # 重复内容回退匹配可能命中已消费的更早位置，游标只进不退，
        # 避免后续切片跟着错位。
        return start, end, max(compact_start, position + len(compact_content))
    return _project_layout_span_by_anchors(
        source_text,
        content,
        compact_source=compact_source,
        source_offsets=source_offsets,
        compact_start=compact_start,
    )


def _project_layout_span_by_anchors(
    source_text: str,
    content: str,
    *,
    compact_source: str,
    source_offsets: list[int],
    compact_start: int,
) -> tuple[int | None, int | None, int]:
    """按行锚点顺序匹配，容忍切片里插入的重复表头等非连续内容。"""

    segments = [
        compact
        for compact in (
            _compact_with_offsets(line)[0]
            for line in _LIST_MARKER_PREFIX.sub("", content).split("\n")
        )
        if len(compact) >= 6
    ]
    if not segments:
        return None, None, compact_start
    total = sum(len(segment) for segment in segments)

    def match_from(cursor: int) -> tuple[list[tuple[int, int]], int]:
        matched: list[tuple[int, int]] = []
        position = cursor
        for segment in segments:
            found = compact_source.find(segment, position)
            if found < 0:
                continue
            matched.append((found, found + len(segment)))
            position = found + len(segment)
        return matched, sum(end - start for start, end in matched)

    matched, covered = match_from(compact_start)
    if covered * 2 < total:
        retried, retry_covered = match_from(0)
        if retry_covered > covered:
            matched, covered = retried, retry_covered
    # 覆盖不足一半时宁可不定位，也不能给出错误的行号。
    if not matched or covered * 2 < total:
        return None, None, compact_start
    start, end = _expand_to_line_boundaries(
        source_text,
        source_offsets[matched[0][0]],
        source_offsets[matched[-1][1] - 1] + 1,
    )
    return start, end, max(compact_start, matched[-1][1])


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

    from app.modules.data_process.algorithms import (
        detect_layout_repeated_blocks,
        remove_layout_repeated_blocks,
    )

    convert_started = time.perf_counter()
    try:
        with _CONVERTER_LOCK:
            conversion = _document_converter().convert(
                DocumentStream(name=filename, stream=BytesIO(raw))
            )
    except DoclingError as exc:
        raise ValueError(f"文档版面解析失败: {exc}") from exc
    logger.info(
        "layout chunking convert done file=%s elapsed=%.2fs",
        filename,
        time.perf_counter() - convert_started,
    )

    # 第二层启发式：扫描所有 docling item，识别跨页重复出现的短文本块
    # （docling layout 模型在中文企业 PDF 上把页眉页脚识别成普通 Table，
    # 因此 _MarkdownSerializerProvider 的标签排除规则收效甚微）。
    page_count = len(getattr(conversion.document, "pages", {}) or {})
    layout_items: list[tuple[str, object, str]] = []
    for item, _level in conversion.document.iterate_items():
        text = getattr(item, "text", None)
        if not text and hasattr(item, "export_to_markdown"):
            try:
                text = item.export_to_markdown(doc=conversion.document) or ""
            except TypeError:
                # 旧版 docling_core 无 doc 参数
                text = item.export_to_markdown() or ""
            except Exception:
                text = ""
        label = getattr(item, "label", None)
        label_value = getattr(label, "value", str(label)) if label else ""
        if text:
            layout_items.append((label_value, item, text))
    repeated_blocks = detect_layout_repeated_blocks(
        layout_items, page_count=page_count
    )

    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] = []
    covered_refs: set[str] = set()
    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
        if repeated_blocks:
            content = remove_layout_repeated_blocks(content, repeated_blocks)
            contextualized = remove_layout_repeated_blocks(
                contextualized, repeated_blocks
            )
        if not content:
            continue
        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))
            covered_refs.add(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),
            )
        )

    # HybridChunker(merge_peers=True) 会丢弃"末尾无正文的孤立标题"。
    # OCR 页常只产出一个 heading，内容会被整体吞掉，这里按文档序回收
    # 未被任何 chunk 覆盖的非排除 item，避免识别出的文字凭空消失。
    # 注意 heading 会进入 meta.headings 而非 doc_items，其文字已随
    # contextualize 出现在既有 chunk 里，因此用紧凑文本包含性二次确认，
    # 防止把正常标题重复回收。
    chunk_haystack = _compact_with_offsets(
        "\n".join(chunk.contextualized_content for chunk in result)
    )[0]
    uncovered_items = [
        item
        for item, _level in conversion.document.iterate_items()
        if item.label not in excluded
        and str(item.self_ref) not in covered_refs
        and (getattr(item, "text", None) or "").strip()
        and _compact_with_offsets(str(item.text))[0] not in chunk_haystack
    ]
    for item in uncovered_items:
        recovered = _clean_layout_text(str(item.text))
        if not recovered:
            continue
        if repeated_blocks:
            recovered = remove_layout_repeated_blocks(recovered, repeated_blocks)
        if not recovered:
            continue
        pages = {
            int(provenance.page_no) for provenance in item.prov or ()
        }
        bboxes = [
            {
                "page": int(provenance.page_no),
                "left": float(provenance.bbox.l),
                "top": float(provenance.bbox.t),
                "right": float(provenance.bbox.r),
                "bottom": float(provenance.bbox.b),
                "origin": str(provenance.bbox.coord_origin.value),
            }
            for provenance in item.prov or ()
        ]
        logger.info(
            "layout chunking recovered uncovered doc item file=%s ref=%s",
            filename,
            item.self_ref,
        )
        result.append(
            DocumentChunk(
                original_content=recovered,
                contextualized_content=recovered,
                source_start=None,
                source_end=None,
                source_start_line=None,
                source_end_line=None,
                token_count=len(_tokenizer().encode(recovered)),
                heading_path=(),
                source_pages=tuple(sorted(pages)),
                doc_item_refs=(str(item.self_ref),),
                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,
    )