Codex
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
"""
|
|
自动分块策略
|
|
|
|
根据文件类型自动选择最佳分块器。
|
|
参考 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,
|
|
)
|