115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
|
"""Document chunking strategies used before Milvus ingestion."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import re
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
from rag.chunk_config import ChunkConfig, resolve_chunk_config
|
||
|
|
|
||
|
|
|
||
|
|
SUPPORTED_STRATEGIES = {"default", "qa_pair", "chapter_semantic"}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class Chunk:
|
||
|
|
text: str
|
||
|
|
section_title: str | None = None
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class ChunkingResult:
|
||
|
|
chunks: list[Chunk]
|
||
|
|
requested_strategy: str
|
||
|
|
actual_strategy: str
|
||
|
|
degraded: bool = False
|
||
|
|
warning: str | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def chunk_document(
|
||
|
|
text: str,
|
||
|
|
strategy: str,
|
||
|
|
*,
|
||
|
|
config: ChunkConfig | None = None,
|
||
|
|
) -> ChunkingResult:
|
||
|
|
if strategy not in SUPPORTED_STRATEGIES:
|
||
|
|
raise ValueError(f"Unsupported chunk strategy: {strategy}")
|
||
|
|
config = config or resolve_chunk_config()
|
||
|
|
if strategy == "default":
|
||
|
|
return ChunkingResult(_default_chunks(text, config), strategy, strategy)
|
||
|
|
if strategy == "qa_pair":
|
||
|
|
return ChunkingResult(_qa_pair_chunks(text, config), strategy, strategy)
|
||
|
|
return _chapter_chunks(text, config)
|
||
|
|
|
||
|
|
|
||
|
|
def _default_chunks(text: str, config: ChunkConfig) -> list[Chunk]:
|
||
|
|
paragraphs = [part.strip() for part in re.split(r"\n\s*\n", text) if part.strip()]
|
||
|
|
chunks: list[Chunk] = []
|
||
|
|
for paragraph in paragraphs:
|
||
|
|
chunks.extend(Chunk(part) for part in _split_long_text(paragraph, config))
|
||
|
|
return chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _split_long_text(text: str, config: ChunkConfig) -> list[str]:
|
||
|
|
if len(text) <= config.size:
|
||
|
|
return [text]
|
||
|
|
chunks = []
|
||
|
|
start = 0
|
||
|
|
while start < len(text):
|
||
|
|
end = min(start + config.size, len(text))
|
||
|
|
chunks.append(text[start:end])
|
||
|
|
if end == len(text):
|
||
|
|
break
|
||
|
|
start = end - config.overlap
|
||
|
|
return chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _qa_pair_chunks(text: str, config: ChunkConfig) -> list[Chunk]:
|
||
|
|
matches = list(
|
||
|
|
re.finditer(
|
||
|
|
r"^\s*Q:\s*(.*?)\r?\n\s*A:\s*(.*?)(?=^\s*Q:\s*|\Z)",
|
||
|
|
text,
|
||
|
|
flags=re.MULTILINE | re.DOTALL,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
if not matches or any(match.group(0).strip() == "" for match in matches):
|
||
|
|
raise ValueError("未识别FAQ问答格式，请确认文档包含Q:/A:标记")
|
||
|
|
if any(not match.group(1).strip() or not match.group(2).strip() for match in matches):
|
||
|
|
raise ValueError("FAQ问答对必须同时包含问题和答案")
|
||
|
|
chunks = [
|
||
|
|
Chunk(f"question: {match.group(1).strip()}\nanswer: {match.group(2).strip()}")
|
||
|
|
for match in matches
|
||
|
|
]
|
||
|
|
if any(len(chunk.text) > config.size for chunk in chunks):
|
||
|
|
raise ValueError("FAQ问答对超过chunk_size，整份文档拒绝上传")
|
||
|
|
return chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _chapter_chunks(text: str, config: ChunkConfig) -> ChunkingResult:
|
||
|
|
heading_pattern = re.compile(r"^(#{1,3})\s+(.+?)\s*$", re.MULTILINE)
|
||
|
|
headings = list(heading_pattern.finditer(text))
|
||
|
|
if not headings:
|
||
|
|
return ChunkingResult(
|
||
|
|
_default_chunks(text, config),
|
||
|
|
"chapter_semantic",
|
||
|
|
"default",
|
||
|
|
degraded=True,
|
||
|
|
warning="未发现Markdown标题，已降级为default策略",
|
||
|
|
)
|
||
|
|
|
||
|
|
chunks: list[Chunk] = []
|
||
|
|
title_stack: list[str] = []
|
||
|
|
for index, heading in enumerate(headings):
|
||
|
|
level = len(heading.group(1))
|
||
|
|
title = heading.group(2).strip()
|
||
|
|
title_stack = title_stack[: level - 1] + [title]
|
||
|
|
body_start = heading.end()
|
||
|
|
body_end = headings[index + 1].start() if index + 1 < len(headings) else len(text)
|
||
|
|
body = text[body_start:body_end].strip()
|
||
|
|
if not body:
|
||
|
|
continue
|
||
|
|
section_title = " > ".join(title_stack)
|
||
|
|
prefix = f"【章节：{section_title}】\n"
|
||
|
|
for part in _split_long_text(body, config):
|
||
|
|
chunks.append(Chunk(prefix + part, section_title=section_title))
|
||
|
|
return ChunkingResult(chunks, "chapter_semantic", "chapter_semantic")
|