155 lines
5.4 KiB
Python
155 lines
5.4 KiB
Python
|
# -*- 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,
|
|||
|
|
)
|