kplam
51dc8dc4f6
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索（dense + sparse RRF 融合） - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
122 lines
4.5 KiB
Python
122 lines
4.5 KiB
Python
"""文档 chunk 切分器
|
||
|
||
按文档原生标题树切分 chunk：
|
||
- 结构化文本：每个标题起点切分 section（标题行到下一个标题前），
|
||
超长 section 按空行段落二次切分，单段落仍超长则硬切
|
||
- 无结构文本：直接按空行段落累加切分
|
||
|
||
每个 chunk 记录 section_path（祖先标题链，" / " 连接），与 L2 大纲节点互相定位。
|
||
"""
|
||
|
||
import re
|
||
|
||
import structlog
|
||
|
||
from app.config import settings
|
||
from app.core.headings import Heading, parse_headings
|
||
from app.models.document import ChunkModel
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
# 段落分隔：一个或多个空行
|
||
_PARAGRAPH_SPLIT_PATTERN = re.compile(r"\n\s*\n")
|
||
|
||
|
||
class Chunker:
|
||
"""按标题树切分文档 chunk"""
|
||
|
||
def __init__(self, max_chars: int = settings.chunk_max_chars) -> None:
|
||
self.max_chars = max_chars
|
||
|
||
def chunk(self, text: str, doc_id: str) -> list[ChunkModel]:
|
||
"""将文档文本切分为 chunk 列表
|
||
|
||
Args:
|
||
text: 文档纯文本内容
|
||
doc_id: 文档 ID
|
||
|
||
Returns:
|
||
list[ChunkModel]: 切分结果，chunk_index 从 0 递增
|
||
"""
|
||
stripped = text.strip()
|
||
if not stripped:
|
||
return []
|
||
|
||
# 全文不超长：整篇单 chunk
|
||
if len(stripped) <= self.max_chars:
|
||
return [ChunkModel(doc_id=doc_id, chunk_index=0, text=stripped)]
|
||
|
||
headings = parse_headings(stripped)
|
||
chunks: list[ChunkModel] = []
|
||
|
||
if headings:
|
||
# 结构化：先按标题切分 section，再按长度二次切分
|
||
for section_text, section_path in self._split_sections(stripped, headings):
|
||
for piece in self._split_by_length(section_text):
|
||
chunks.append(
|
||
ChunkModel(doc_id=doc_id, chunk_index=len(chunks), text=piece, section_path=section_path)
|
||
)
|
||
else:
|
||
# 无结构：直接按段落累加切分
|
||
for piece in self._split_by_length(stripped):
|
||
chunks.append(ChunkModel(doc_id=doc_id, chunk_index=len(chunks), text=piece))
|
||
|
||
logger.info("文档切分完成", doc_id=doc_id, chunks_count=len(chunks), has_headings=bool(headings))
|
||
return chunks
|
||
|
||
def _split_sections(self, text: str, headings: list[Heading]) -> list[tuple[str, str]]:
|
||
"""按标题树切分 section，返回 (section 文本, section_path) 列表
|
||
|
||
每个标题起点切分一个 section，section 文本含标题行本身；
|
||
section_path 为祖先标题链（含自身标题），用 " / " 连接；
|
||
首个标题前的引导正文归入无前缀 section（section_path 为空）。
|
||
"""
|
||
lines = text.splitlines()
|
||
sections: list[tuple[str, str]] = []
|
||
|
||
# 首个标题前的引导内容
|
||
preamble = "\n".join(lines[: headings[0].line_index]).strip()
|
||
if preamble:
|
||
sections.append((preamble, ""))
|
||
|
||
# 维护祖先标题栈：遇到同级或更高级标题时弹栈
|
||
stack: list[Heading] = []
|
||
for i, heading in enumerate(headings):
|
||
while stack and stack[-1].level >= heading.level:
|
||
stack.pop()
|
||
stack.append(heading)
|
||
|
||
end = headings[i + 1].line_index if i + 1 < len(headings) else len(lines)
|
||
section_text = "\n".join(lines[heading.line_index : end]).strip()
|
||
section_path = " / ".join(h.title for h in stack)
|
||
sections.append((section_text, section_path))
|
||
|
||
return sections
|
||
|
||
def _split_by_length(self, text: str) -> list[str]:
|
||
"""按 max_chars 切分文本：先按空行段落累加，单段落超长则硬切"""
|
||
if len(text) <= self.max_chars:
|
||
return [text]
|
||
|
||
pieces: list[str] = []
|
||
current = ""
|
||
for paragraph in _PARAGRAPH_SPLIT_PATTERN.split(text):
|
||
paragraph = paragraph.strip()
|
||
if not paragraph:
|
||
continue
|
||
candidate = f"{current}\n\n{paragraph}" if current else paragraph
|
||
if len(candidate) <= self.max_chars:
|
||
current = candidate
|
||
continue
|
||
if current:
|
||
pieces.append(current)
|
||
current = ""
|
||
# 单段落仍超长：按 max_chars 硬切
|
||
if len(paragraph) > self.max_chars:
|
||
pieces.extend(paragraph[i : i + self.max_chars] for i in range(0, len(paragraph), self.max_chars))
|
||
else:
|
||
current = paragraph
|
||
if current:
|
||
pieces.append(current)
|
||
return pieces
|