feat: 添加 PDFSplitter — pymupdf 提取文字后分块

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"]

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",

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)