diff --git a/pyproject.toml b/pyproject.toml index 0291985..13d29c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,17 @@ dev = [ "pytest>=8.0", "httpx>=0.27.0", ] +pdf = [ + "pymupdf>=1.24.0", +] +html = [ + "beautifulsoup4>=4.12.0", +] +all = [ + "md-vector-db[pdf,html]", + "requests>=2.31.0", + "openai>=1.0.0", +] [build-system] requires = ["hatchling"] diff --git a/src/core/splitters/__init__.py b/src/core/splitters/__init__.py index 4fc8eed..69764ee 100644 --- a/src/core/splitters/__init__.py +++ b/src/core/splitters/__init__.py @@ -3,6 +3,7 @@ from src.core.splitters.base import Splitter, BaseTextSplitter from src.core.splitters.markdown import MarkdownSplitter from src.core.splitters.text import TextSplitter +from src.core.splitters.pdf import PDFSplitter from src.core.splitters.registry import get_splitter, register_splitter, SUPPORTED_SUFFIXES __all__ = [ @@ -10,6 +11,7 @@ __all__ = [ "BaseTextSplitter", "MarkdownSplitter", "TextSplitter", + "PDFSplitter", "get_splitter", "register_splitter", "SUPPORTED_SUFFIXES", diff --git a/src/core/splitters/pdf.py b/src/core/splitters/pdf.py new file mode 100644 index 0000000..255ac93 --- /dev/null +++ b/src/core/splitters/pdf.py @@ -0,0 +1,49 @@ +"""PDF 文档分块器 — 使用 pymupdf 提取文字后委托 TextSplitter.""" +import logging +from src.core.splitters.text import TextSplitter + +logger = logging.getLogger("md-vector-db") + + +class PDFSplitter: + """PDF 分块器:pymupdf 提取文字 → TextSplitter 分块. + + 实现 Splitter Protocol,内部组合 TextSplitter 实例。 + 注意: split() 的 text 参数实际接收 PDF 文件路径(非文本内容)。 + """ + + def __init__(self, max_size: int = 1000, overlap: int = 100): + self._text_splitter = TextSplitter(max_size=max_size, overlap=overlap) + + def split(self, text: str, source_file: str = "") -> list[dict]: + """从 PDF 文件提取文字并分块. + + Args: + text: PDF 文件路径(非文本内容!由 ingest_file 传入) + source_file: 来源文件名 + """ + try: + import fitz # pymupdf + except ImportError: + raise ImportError( + "PDF 支持需要 pymupdf 库. 请执行: uv sync --extra pdf" + ) + + pdf_path = text # text 参数实际是文件路径 + extracted_pages = [] + try: + doc = fitz.open(pdf_path) + for page in doc: + page_text = page.get_text() + if page_text.strip(): + extracted_pages.append(page_text) + doc.close() + except Exception as e: + logger.error("PDF 解析失败: %s — %s", pdf_path, e) + raise ValueError(f"PDF 解析失败: {e}") from e + + if not extracted_pages: + return [] + + full_text = "\n\n".join(extracted_pages) + return self._text_splitter.split(full_text, source_file=source_file)