""" md_chunker.py Markdown 语义分块器 支持标题层级、代码块、表格等元素的智能分块 """ import json import os import re from typing import List, Dict, Tuple from pathlib import Path from dataclasses import dataclass, field, asdict @dataclass class MDChunk: """Single Markdown chunk data structure""" chunk_id: int chunk_type: str human_description: str raw_content: str context: str metadata: Dict = field(default_factory=dict) class MarkdownSemanticChunker: """ Markdown 语义分块器 v1.0 分块策略： 1. 按标题层级（H1/H2/H3...）划分大段落 2. 代码块作为独立 chunk 3. 表格作为独立 chunk 4. 过长段落内部按句子/段落二次拆分 """ # Heading patterns HEADING_PATTERN = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE) # Code block pattern (fenced) CODE_BLOCK_PATTERN = re.compile(r'```(\w*)\n([\s\S]*?)```', re.MULTILINE) # Inline code pattern INLINE_CODE_PATTERN = re.compile(r'`([^`]+)`') # Table pattern TABLE_PATTERN = re.compile(r'\|.+\|\n\|[-| :]+\|\n((?:\|.+\|\n)*)', re.MULTILINE) # List pattern LIST_PATTERN = re.compile(r'^(\s*[-*+]\s+.+)+', re.MULTILINE) def __init__(self, max_chunk_size: int = 2000): self.max_chunk_size = max_chunk_size def chunk_file(self, file_path: str) -> List[Dict]: """处理单个 Markdown 文件""" if not os.path.exists(file_path): raise FileNotFoundError(f"File not found: {file_path}") with open(file_path, 'r', encoding='utf-8') as f: content = f.read() file_name = Path(file_path).stem chunks = [] chunk_id = 0 # 尝试提取文档标题（第一个 H1） title_match = re.search(r'^#\s+(.+)$', content, re.MULTILINE) doc_title = title_match.group(1).strip() if title_match else file_name # 按结构化元素分割 segments = self._split_by_structure(content) for segment in segments: seg_type = segment['type'] seg_content = segment['content'] if not seg_content.strip(): continue # 构建描述 description = self._build_description(seg_type, seg_content, doc_title) # 如果超过最大长度，尝试二次拆分 if len(seg_content) > self.max_chunk_size: sub_chunks = self._split_large_chunk( seg_content, seg_type, doc_title, chunk_id ) chunks.extend([asdict(c) for c in sub_chunks]) chunk_id += len(sub_chunks) else: chunks.append(asdict(MDChunk( chunk_id=chunk_id, chunk_type=seg_type, human_description=description, raw_content=seg_content.strip(), context=f"{doc_title}", metadata=segment.get('metadata', {}) ))) chunk_id += 1 return chunks def _split_by_structure(self, content: str) -> List[Dict]: """ 按 Markdown 结构分割内容 返回: [{'type': 'h1/h2/code/table/paragraph', 'content': '...', 'metadata': {...}}] """ segments = [] # 首先提取所有代码块（保留位置标记，稍后处理） code_blocks = [] code_pattern = re.compile(r'(```\w*\n[\s\S]*?```)', re.MULTILINE) last_end = 0 for match in code_pattern.finditer(content): # 处理代码块前的普通文本 before = content[last_end:match.start()] if before.strip(): segments.extend(self._process_text_section(before)) # 添加代码块 code_blocks.append(match.group(1)) lang_match = re.match(r'```(\w*)', match.group(1)) lang = lang_match.group(1) if lang_match else '' segments.append({ 'type': 'code', 'content': match.group(1), 'metadata': {'language': lang} }) last_end = match.end() # 处理剩余文本 remaining = content[last_end:] if remaining.strip(): segments.extend(self._process_text_section(remaining)) return segments def _process_text_section(self, text: str) -> List[Dict]: """处理普通文本区域，提取标题和段落""" segments = [] # 按标题分割 lines = text.split('\n') current_section = [] current_heading_level = 0 current_heading = '' for line in lines: heading_match = re.match(r'^(#{1,6})\s+(.+)', line) if heading_match: # 保存之前的段落 if current_section: section_text = '\n'.join(current_section).strip() if section_text: segments.append({ 'type': self._get_section_type(current_heading_level, current_heading), 'content': section_text, 'metadata': { 'heading': current_heading, 'heading_level': current_heading_level } }) current_section = [] # 开始新标题区域 current_heading_level = len(heading_match.group(1)) current_heading = heading_match.group(2).strip() else: current_section.append(line) # 保存最后一段 if current_section: section_text = '\n'.join(current_section).strip() if section_text: segments.append({ 'type': self._get_section_type(current_heading_level, current_heading), 'content': section_text, 'metadata': { 'heading': current_heading, 'heading_level': current_heading_level } }) return segments def _get_section_type(self, level: int, heading: str) -> str: """根据标题级别和内容确定段落类型""" heading_lower = heading.lower() if level == 1: return 'section_h1' elif level == 2: # 检测特殊章节类型 if any(kw in heading_lower for kw in ['install', '安装', 'setup', '部署']): return 'section_installation' elif any(kw in heading_lower for kw in ['config', '配置', 'setting']): return 'section_configuration' elif any(kw in heading_lower for kw in ['api', '接口']): return 'section_api' elif any(kw in heading_lower for kw in ['example', '示例', 'usage', '使用']): return 'section_example' elif any(kw in heading_lower for kw in ['faq', 'question', '问题', '常见']): return 'section_faq' elif any(kw in heading_lower for kw in ['changelog', '更新', 'release']): return 'section_changelog' return 'section_h2' elif level == 3: return 'section_h3' else: return 'section_other' def _build_description(self, chunk_type: str, content: str, doc_title: str) -> str: """为 chunk 生成人类可读描述""" lines = content.split('\n')[:5] preview = ' '.join(line.strip() for line in lines if line.strip())[:150] if chunk_type == 'code': lang = '' lang_match = re.match(r'```(\w*)', content) if lang_match: lang = lang_match.group(1) or 'text' return f"Code block (language: {lang}) in {doc_title}. Preview: {preview}" elif chunk_type.startswith('section_'): heading = content.split('\n')[0] if '\n' in content else content[:50] heading_clean = re.sub(r'^#+\s+', '', heading) type_hint = chunk_type.replace('section_', '') return f"[{type_hint.upper()}] {heading_clean}. Content: {preview}" else: return f"Document section in {doc_title}. Content: {preview}" def _split_large_chunk(self, content: str, chunk_type: str, doc_title: str, start_id: int) -> List[MDChunk]: """对过长的 chunk 进行二次拆分""" chunks = [] # 按段落分割（双换行符） paragraphs = re.split(r'\n\n+', content) current_chunk = [] current_size = 0 for para in paragraphs: para_size = len(para) if current_size + para_size > self.max_chunk_size and current_chunk: # 当前块已满，生成 chunk chunk_text = '\n\n'.join(current_chunk) chunks.append(MDChunk( chunk_id=start_id + len(chunks), chunk_type=f"{chunk_type}_part", human_description=f"Part of {doc_title} ({chunk_type}): {chunk_text[:100]}...", raw_content=chunk_text, context=f"{doc_title} (continued)", metadata={'part': len(chunks) + 1, 'original_type': chunk_type} )) current_chunk = [] current_size = 0 current_chunk.append(para) current_size += para_size + 2 # 处理剩余内容 if current_chunk: chunk_text = '\n\n'.join(current_chunk) chunks.append(MDChunk( chunk_id=start_id + len(chunks), chunk_type=f"{chunk_type}_part", human_description=f"Part of {doc_title} ({chunk_type}): {chunk_text[:100]}...", raw_content=chunk_text, context=f"{doc_title} (continued)", metadata={'part': len(chunks) + 1, 'original_type': chunk_type} )) return chunks if chunks else [MDChunk( chunk_id=start_id, chunk_type=chunk_type, human_description=f"{doc_title}: {content[:100]}...", raw_content=content[:self.max_chunk_size], context=doc_title, metadata={'truncated': True} )] def chunk_directory(self, dir_path: str, extensions: tuple = ('.md', '.markdown')) -> List[Dict]: """批量处理目录下所有 Markdown 文件""" all_chunks = [] file_count = 0 for root, _, files in os.walk(dir_path): for file in files: if file.lower().endswith(extensions): file_path = os.path.join(root, file) try: chunks = self.chunk_file(file_path) all_chunks.extend(chunks) file_count += 1 print(f"OK {file_path}: {len(chunks)} chunks") except Exception as e: print(f"FAIL {file_path}: {e}") print(f"\nTotal: {file_count} files, {len(all_chunks)} chunks") return all_chunks def save_chunks_to_json(chunks: List[Dict], output_path: str): """保存 chunks 到 JSON 文件""" with open(output_path, 'w', encoding='utf-8') as f: json.dump(chunks, f, ensure_ascii=False, indent=2) print(f"Saved {len(chunks)} chunks to {output_path}") def print_chunk_summary(chunks: List[Dict]): """打印 chunk 类型统计""" type_counts = {} for chunk in chunks: chunk_type = chunk["chunk_type"] type_counts[chunk_type] = type_counts.get(chunk_type, 0) + 1 print("\nChunk Type Summary:") for chunk_type, count in sorted(type_counts.items(), key=lambda x: -x[1]): print(f" {chunk_type}: {count}") if __name__ == "__main__": import sys chunker = MarkdownSemanticChunker(max_chunk_size=2000) if len(sys.argv) > 1: path = sys.argv[1] if os.path.isdir(path): all_chunks = chunker.chunk_directory(path) output_path = os.path.join(os.path.dirname(path.rstrip("/\\")) if os.path.dirname(path) else ".", os.path.basename(path.rstrip("/\\")) + "_md_chunks.json") save_chunks_to_json(all_chunks, output_path) print_chunk_summary(all_chunks) else: chunks = chunker.chunk_file(path) output_path = path.replace(".md", "_chunks.json").replace(".markdown", "_chunks.json") save_chunks_to_json(chunks, output_path) print(f"\n{'='*60}") print("Chunking Results Preview") print(f"{'='*60}") for chunk in chunks[:10]: print(f"\n[Chunk {chunk['chunk_id']}] Type: {chunk['chunk_type']}") print(f"Description: {chunk['human_description'][:120]}...") print(f"Content length: {len(chunk['raw_content'])} chars") if len(chunks) > 10: print(f"\n... and {len(chunks) - 10} more chunks") print_chunk_summary(chunks) else: print("=" * 60) print("Markdown Semantic Chunking v1.0") print("=" * 60) print("\nUsage: python md_chunker.py ")