""" 自动分块策略 根据文件类型自动选择最佳分块器。 参考 Dify 的 auto 分块模式。 """ import logging from typing import Dict, Any, List, Optional from .base import BaseChunker, ChunkResult logger = logging.getLogger(__name__) # 文件类型 → 推荐分块策略 FILE_TYPE_STRATEGY_MAP = { # Markdown 文件使用 Markdown 分块器 'md': 'markdown', 'markdown': 'markdown', # 代码文件使用按句子分块(按行/语句边界) 'py': 'sentence', 'js': 'sentence', 'ts': 'sentence', 'java': 'sentence', 'go': 'sentence', 'rs': 'sentence', 'c': 'sentence', 'cpp': 'sentence', 'h': 'sentence', # 纯文本使用按句子分块 'txt': 'sentence', # CSV/Excel 使用固定大小(表格数据按行分割更合理) 'csv': 'fixed', 'xlsx': 'fixed', 'xls': 'fixed', # HTML 使用 Markdown 分块器(HTML 结构类似) 'html': 'markdown', 'htm': 'markdown', # 其他文档类型使用递归分块 'pdf': 'recursive', 'docx': 'recursive', 'doc': 'recursive', 'pptx': 'recursive', 'ppt': 'recursive', } class AutoChunker(BaseChunker): """ 自动分块器 根据文档的文件类型自动选择最佳分块策略。 metadata 中需要包含 'file_type' 字段。 """ def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]: """自动选择分块策略并执行""" if not text or not text.strip(): return [] metadata = metadata or {} file_type = metadata.get('file_type', '').lower().lstrip('.') # 根据文件类型选择策略 strategy = FILE_TYPE_STRATEGY_MAP.get(file_type, 'recursive') logger.info(f"AutoChunker: file_type={file_type} → strategy={strategy}") # 动态创建对应的分块器 chunker = self._get_chunker(strategy) return chunker.chunk(text, metadata) def _get_chunker(self, strategy: str) -> BaseChunker: """获取对应策略的分块器实例""" from .recursive import RecursiveChunker from .markdown import MarkdownChunker from .fixed import FixedChunker from .sentence import SentenceChunker chunkers = { 'recursive': RecursiveChunker, 'markdown': MarkdownChunker, 'fixed': FixedChunker, 'sentence': SentenceChunker, } cls = chunkers.get(strategy, RecursiveChunker) return cls( chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap, separator=self.separator, )