iAOP/core/rag-kb/documents.py

# -*- coding: utf-8 -*-
"""RAG 文档模型与段落抽取 —— 抽取管线（Issue #46 / PRD 5.4 引用溯源）。

客户文档（工艺规范 / SOP / 国标，格式：文本或 Markdown）经抽取管线转成
**带来源的段落（chunk）**：每个 chunk 保留 `doc_id + 章节号 + 段落号`，
检索命中的段落可精确回指源文档章节（对齐 PRD 5.4「返回命中文档片段+来源」）。

分块规则（简单可靠，零外部依赖）：
- 按空行 / Markdown 标题切分段落；标题行单独记录为小节标题；
- 每段按 `max_chars` 阈值再切分（优先在 `。．.!?；;` 等句边界断开），
  避免超长段落撑爆向量块；
- chunk 携带 `source` 溯源串（如 `沸腾氯化炉异常处置SOP §2.1 ¶3`）。
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import List, Optional

# Markdown 标题行（# / ## / ### ...）—— 作为小节起点
_HEADING_RE = re.compile(r"^\s{0,3}#{1,6}\s+(.*)$")
# 句子边界（中文句号/省略号/英文句点/问号/感叹号/分号）
_SENT_BOUNDARY = re.compile(r"(?<=[。．.!?；;])")


@dataclass
class KbDocument:
    """一份知识库源文档（标题即唯一 ID，来源类型见 KnowledgeSourceKind）。"""

    doc_id: str                # 文档唯一 ID（清洗后用于对象键）
    title: str                 # 文档标题（显示与溯源用）
    text: str                  # 原始文本（Markdown 或纯文本）
    category: str = "process"  # 知识源类别（process/sop/standard）
    version: str = "1.0.0"     # 文档版本（来源追溯）

    def __post_init__(self) -> None:
        self.text = (self.text or "").strip()
        if not self.title:
            raise ValueError("KbDocument.title 不能为空")
        if not self.text:
            raise ValueError(f"KbDocument {self.title!r} 文本为空")


@dataclass
class Chunk:
    """抽取后的段落（检索最小单元），携带完整来源信息。"""

    doc_id: str
    title: str
    section: str          # 小节标题（无标题段落记 `§0 概述` 或空串）
    seq: int              # 段落序号（文档内 1 起）
    text: str
    category: str = "process"

    @property
    def source(self) -> str:
        """溯源串：`标题 §章节 ¶段落`（PRD 5.4 引用溯源返回给调用方）。"""
        base = f"{self.title}"
        if self.section:
            base += f" §{self.section}"
        return f"{base} ¶{self.seq}"


def _split_paragraphs(text: str) -> List[str]:
    """按空行 / 标题切分段落；顺带记录标题。返回 (section, para) 对交给调用方。

    实现：先按空行粗分块，再在每个块内识别标题行（标题行不入段落文本，
    而是成为随后段落的 section 名）。
    """
    blocks: List[List[str]] = []
    cur: List[str] = []
    for raw in text.split("\n"):
        line = raw.rstrip()
        if not line.strip():
            if cur:
                blocks.append(cur)
                cur = []
            continue
        cur.append(line)
    if cur:
        blocks.append(cur)
    return ["\n".join(b) for b in blocks]


def _split_by_chars(para: str, max_chars: int) -> List[str]:
    """超长段落按句边界切分，每片不超过 max_chars。"""
    if len(para) <= max_chars:
        return [para]
    pieces: List[str] = []
    for sentence in _SENT_BOUNDARY.split(para):
        if not sentence.strip():
            continue
        if pieces and len(pieces[-1]) + len(sentence) <= max_chars:
            pieces[-1] += sentence
        else:
            # 单句仍超长则硬切，避免无限循环
            while len(sentence) > max_chars:
                pieces.append(sentence[:max_chars])
                sentence = sentence[max_chars:]
            if sentence:
                pieces.append(sentence)
    return [p for p in pieces if p.strip()]


def chunk_document(doc: KbDocument, max_chars: int = 500) -> List[Chunk]:
    """文档 → 带来源的段落列表（抽取管线核心，纯函数便于测试）。"""
    chunks: List[Chunk] = []
    seq = 0
    section = ""
    for para in _split_paragraphs(doc.text):
        lines = para.split("\n")
        heading = None
        for ln in lines:
            m = _HEADING_RE.match(ln)
            if m:
                heading = m.group(1).strip()
        if heading is not None:
            section = heading
            body_lines = [ln for ln in lines if not _HEADING_RE.match(ln)]
            para = "\n".join(body_lines).strip()
            if not para:
                continue  # 纯标题行：仅更新小节名
        for piece in _split_by_chars(para, max_chars):
            seq += 1
            chunks.append(
                Chunk(
                    doc_id=doc.doc_id,
                    title=doc.title,
                    section=section,
                    seq=seq,
                    text=piece,
                    category=doc.category,
                )
            )
    return chunks


def build_document(
    title: str,
    text: str,
    category: str = "process",
    version: str = "1.0.0",
    doc_id: Optional[str] = None,
) -> KbDocument:
    """便捷构造：doc_id 缺省时由标题清洗生成（对齐 templating.sanitize）。"""
    from .templating import sanitize  # 局部导入避免循环依赖

    return KbDocument(
        doc_id=doc_id or sanitize(title),
        title=title,
        text=text,
        category=category,
        version=version,
    )