Codex
173 lines
5.4 KiB
Python
173 lines
5.4 KiB
Python
"""
|
|
Markdown 结构化分块策略
|
|
|
|
按 Markdown 标题层级分割文档，保留文档结构信息
|
|
"""
|
|
import re
|
|
import logging
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
from .base import BaseChunker, ChunkResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MarkdownChunker(BaseChunker):
|
|
"""
|
|
Markdown 分块器
|
|
|
|
按标题层级分割 Markdown 文档，每个标题下的内容作为一个分块
|
|
如果单个标题下的内容超过 chunk_size，则使用递归分割
|
|
"""
|
|
|
|
def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
|
|
"""按 Markdown 标题分块"""
|
|
if not text or not text.strip():
|
|
return []
|
|
|
|
text = self._clean_text(text)
|
|
metadata = metadata or {}
|
|
|
|
# 按标题分割
|
|
sections = self._split_by_headers(text)
|
|
|
|
# 处理每个 section
|
|
raw_chunks = []
|
|
for section in sections:
|
|
header = section.get('header', '')
|
|
content = section.get('content', '')
|
|
level = section.get('level', 0)
|
|
|
|
if not content.strip():
|
|
continue
|
|
|
|
# 组合标题和内容
|
|
full_text = f"{header}\n{content}" if header else content
|
|
|
|
if len(full_text) <= self.chunk_size:
|
|
raw_chunks.append({
|
|
'content': full_text.strip(),
|
|
'metadata': {
|
|
**metadata,
|
|
'header': header,
|
|
'header_level': level,
|
|
}
|
|
})
|
|
else:
|
|
# 内容超长，递归分割
|
|
sub_chunks = self._split_long_section(content, header)
|
|
for i, sub in enumerate(sub_chunks):
|
|
raw_chunks.append({
|
|
'content': sub.strip(),
|
|
'metadata': {
|
|
**metadata,
|
|
'header': header,
|
|
'header_level': level,
|
|
'sub_chunk': i,
|
|
}
|
|
})
|
|
|
|
# 合并过小的分块
|
|
merged = self._merge_small_section_chunks(raw_chunks)
|
|
|
|
# 构建结果
|
|
results = []
|
|
for i, item in enumerate(merged):
|
|
if item['content'].strip():
|
|
results.append(ChunkResult(
|
|
content=item['content'],
|
|
position=i,
|
|
metadata=item.get('metadata', {}),
|
|
))
|
|
|
|
return results
|
|
|
|
def _split_by_headers(self, text: str) -> List[Dict[str, Any]]:
|
|
"""按 Markdown 标题分割"""
|
|
# 匹配 Markdown 标题: # Title, ## Title, ### Title 等
|
|
header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
|
|
|
|
sections = []
|
|
last_end = 0
|
|
last_header = ''
|
|
last_level = 0
|
|
|
|
for match in header_pattern.finditer(text):
|
|
# 保存前一个 section 的内容
|
|
if last_end > 0 or match.start() > 0:
|
|
content = text[last_end:match.start()]
|
|
if content.strip() or last_header:
|
|
sections.append({
|
|
'header': last_header,
|
|
'content': content.strip(),
|
|
'level': last_level,
|
|
})
|
|
|
|
last_header = match.group(0)
|
|
last_level = len(match.group(1))
|
|
last_end = match.end()
|
|
|
|
# 最后一个 section
|
|
remaining = text[last_end:]
|
|
if remaining.strip() or last_header:
|
|
sections.append({
|
|
'header': last_header,
|
|
'content': remaining.strip(),
|
|
'level': last_level,
|
|
})
|
|
|
|
# 如果没有找到任何标题，整个文本作为一个 section
|
|
if not sections:
|
|
sections.append({
|
|
'header': '',
|
|
'content': text.strip(),
|
|
'level': 0,
|
|
})
|
|
|
|
return sections
|
|
|
|
def _split_long_section(self, content: str, header: str = '') -> List[str]:
|
|
"""分割超长的 section 内容"""
|
|
from .recursive import RecursiveChunker
|
|
|
|
chunker = RecursiveChunker(
|
|
chunk_size=self.chunk_size,
|
|
chunk_overlap=self.chunk_overlap,
|
|
)
|
|
results = chunker.chunk(content)
|
|
|
|
chunks = []
|
|
for i, result in enumerate(results):
|
|
# 第一个分块带上标题
|
|
if i == 0 and header:
|
|
chunks.append(f"{header}\n{result.content}")
|
|
else:
|
|
chunks.append(result.content)
|
|
|
|
return chunks if chunks else [content]
|
|
|
|
def _merge_small_section_chunks(self, chunks: List[Dict], min_size: int = 80) -> List[Dict]:
|
|
"""合并过小的 section 分块"""
|
|
if not chunks:
|
|
return []
|
|
|
|
merged = []
|
|
buffer = None
|
|
|
|
for chunk in chunks:
|
|
if buffer is None:
|
|
buffer = chunk
|
|
elif len(buffer['content']) < min_size and len(buffer['content']) + len(chunk['content']) <= self.chunk_size:
|
|
buffer['content'] = buffer['content'] + "\n\n" + chunk['content']
|
|
else:
|
|
merged.append(buffer)
|
|
buffer = chunk
|
|
|
|
if buffer:
|
|
if merged and len(buffer['content']) < min_size:
|
|
merged[-1]['content'] = merged[-1]['content'] + "\n\n" + buffer['content']
|
|
else:
|
|
merged.append(buffer)
|
|
|
|
return merged
|