79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
"""Splitter 注册表 — 按文件扩展名自动选择分块器."""
|
|
from pathlib import Path
|
|
from src.core.splitters.base import Splitter
|
|
|
|
# 扩展名 → Splitter 类名映射
|
|
_DEFAULT_MAP: dict[str, str] = {
|
|
".md": "markdown",
|
|
".markdown": "markdown",
|
|
".txt": "text",
|
|
".pdf": "pdf",
|
|
".html": "html",
|
|
".htm": "html",
|
|
".epub": "epub",
|
|
}
|
|
|
|
# 所有支持的扩展名集合（供外部遍历文件使用）
|
|
SUPPORTED_SUFFIXES = frozenset(_DEFAULT_MAP.keys())
|
|
|
|
# 用户可注册自定义 Splitter
|
|
_custom_registry: dict[str, type[Splitter]] = {}
|
|
|
|
|
|
def register_splitter(ext: str, splitter_cls: type[Splitter]) -> None:
|
|
"""注册自定义 Splitter 类.
|
|
|
|
注意: 此函数非线程安全，请在程序启动时调用（单线程阶段）。
|
|
运行时动态注册需自行加锁。
|
|
"""
|
|
ext = ext.lower() if ext.startswith(".") else f".{ext}"
|
|
_custom_registry[ext] = splitter_cls
|
|
|
|
|
|
def get_splitter(
|
|
file_path: str,
|
|
max_size: int = 1000,
|
|
overlap: int = 100,
|
|
) -> Splitter:
|
|
"""根据文件扩展名自动选择 Splitter，未匹配回退到 TextSplitter.
|
|
|
|
Args:
|
|
file_path: 文件路径（用于提取扩展名）
|
|
max_size: 分块最大字符数
|
|
overlap: 相邻块重叠字符数
|
|
|
|
Returns:
|
|
对应格式的 Splitter 实例
|
|
"""
|
|
ext = Path(file_path).suffix.lower()
|
|
|
|
# 优先查用户自定义注册
|
|
if ext in _custom_registry:
|
|
return _custom_registry[ext](max_size=max_size, overlap=overlap)
|
|
|
|
kind = _DEFAULT_MAP.get(ext, "text")
|
|
|
|
if kind == "markdown":
|
|
from src.core.splitters.markdown import MarkdownSplitter
|
|
return MarkdownSplitter(max_size=max_size, overlap=overlap)
|
|
|
|
if kind == "text":
|
|
from src.core.splitters.text import TextSplitter
|
|
return TextSplitter(max_size=max_size, overlap=overlap)
|
|
|
|
if kind == "pdf":
|
|
from src.core.splitters.pdf import PDFSplitter
|
|
return PDFSplitter(max_size=max_size, overlap=overlap)
|
|
|
|
if kind == "html":
|
|
from src.core.splitters.html import HTMLSplitter
|
|
return HTMLSplitter(max_size=max_size, overlap=overlap)
|
|
|
|
if kind == "epub":
|
|
from src.core.splitters.epub import EPUBSplitter
|
|
return EPUBSplitter(max_size=max_size, overlap=overlap)
|
|
|
|
# 回退
|
|
from src.core.splitters.text import TextSplitter
|
|
return TextSplitter(max_size=max_size, overlap=overlap)
|