From 3b8b585f31e118d61efd84eca826071e9c49dbb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com> Date: Fri, 10 Jul 2026 14:32:41 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20EPUBSplitter=20?= =?UTF-8?q?=E2=80=94=20ebooklib=20=E6=8F=90=E5=8F=96=E7=AB=A0=E8=8A=82?= =?UTF-8?q?=E6=96=87=E5=AD=97=E5=90=8E=E5=88=86=E5=9D=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 splitters/epub.py: ebooklib 读取 EPUB → BeautifulSoup 去标签 → TextSplitter 分块 - epub 作为可选依赖: uv sync --extra epub - registry 新增 .epub 映射, ingest_file 处理 EPUB 二进制文件 - 测试: 3 个 EPUBSplitter 测试 (含空 EPUB/Protocol 合规) Co-Authored-By: Claude --- pyproject.toml | 20 +++------- src/core/ingest.py | 4 +- src/core/splitters/__init__.py | 4 +- src/core/splitters/epub.py | 68 ++++++++++++++++++++++++++++++++++ src/core/splitters/registry.py | 5 +++ tests/test_splitters_epub.py | 56 ++++++++++++++++++++++++++++ 6 files changed, 139 insertions(+), 18 deletions(-) create mode 100644 src/core/splitters/epub.py create mode 100644 tests/test_splitters_epub.py diff --git a/pyproject.toml b/pyproject.toml index 13d29c9..b1cb6fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,21 +18,11 @@ dependencies = [ md-vector-db = "src.cli.main:app" [project.optional-dependencies] -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", -] +dev = ["pytest>=8.0", "httpx>=0.27.0"] +pdf = ["pymupdf>=1.24.0"] +html = ["beautifulsoup4>=4.12.0"] +epub = ["ebooklib>=0.18"] +all = ["md-vector-db[pdf,html,epub]", "requests>=2.31.0", "openai>=1.0.0"] [build-system] requires = ["hatchling"] diff --git a/src/core/ingest.py b/src/core/ingest.py index d870ebf..520c409 100644 --- a/src/core/ingest.py +++ b/src/core/ingest.py @@ -46,9 +46,9 @@ class DocumentIngestor: file_path, max_size=1000, overlap=100 ) - # PDF 文件特殊处理：splitter 内部读取二进制内容 + # PDF/EPUB 二进制文件特殊处理：splitter 内部读取文件 suffix = path.suffix.lower() - if suffix in (".pdf",): + if suffix in (".pdf", ".epub"): chunks = splitter.split(str(path), source_file=file_name) return self._add_chunks(chunks, file_name) diff --git a/src/core/splitters/__init__.py b/src/core/splitters/__init__.py index 69764ee..038b148 100644 --- a/src/core/splitters/__init__.py +++ b/src/core/splitters/__init__.py @@ -1,9 +1,10 @@ -"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML.""" +"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML / EPUB.""" 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.epub import EPUBSplitter from src.core.splitters.registry import get_splitter, register_splitter, SUPPORTED_SUFFIXES __all__ = [ @@ -12,6 +13,7 @@ __all__ = [ "MarkdownSplitter", "TextSplitter", "PDFSplitter", + "EPUBSplitter", "get_splitter", "register_splitter", "SUPPORTED_SUFFIXES", diff --git a/src/core/splitters/epub.py b/src/core/splitters/epub.py new file mode 100644 index 0000000..0c4b247 --- /dev/null +++ b/src/core/splitters/epub.py @@ -0,0 +1,68 @@ +"""EPUB 电子书分块器 — 使用 ebooklib 提取文字后委托 TextSplitter.""" +import logging +from src.core.splitters.text import TextSplitter + +logger = logging.getLogger("md-vector-db") + + +class EPUBSplitter: + """EPUB 分块器：ebooklib 提取各章节文字 → TextSplitter 分块. + + 实现 Splitter Protocol，内部组合 TextSplitter 实例。 + split() 的 text 参数实际接收 EPUB 文件路径（非文本内容）。 + """ + + 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]: + """从 EPUB 文件提取各章节文字并分块. + + Args: + text: EPUB 文件路径（非文本内容，由 ingest_file 传入） + source_file: 来源文件名 + """ + try: + import ebooklib + from ebooklib import epub + except ImportError: + raise ImportError( + "EPUB 支持需要 ebooklib 库. 请执行: uv sync --extra epub" + ) + + epub_path = text + try: + book = epub.read_epub(epub_path) + except Exception as e: + logger.error("EPUB 解析失败: %s — %s", epub_path, e) + raise ValueError(f"EPUB 解析失败: {e}") from e + + extracted_chapters = [] + for item in book.get_items_of_type(ebooklib.ITEM_DOCUMENT): + try: + # ebooklib 的 get_content() 返回 bytes + content = item.get_content().decode("utf-8") + except UnicodeDecodeError: + logger.warning("EPUB 跳过一个无法解码的章节: %s", item.get_name()) + continue + + # 用 BeautifulSoup 去标签（如果可用），否则保留原样 + try: + from bs4 import BeautifulSoup + soup = BeautifulSoup(content, "html.parser") + for tag in soup(["script", "style"]): + tag.decompose() + text_content = soup.get_text(separator="\n") + except ImportError: + # 无 bs4 时手动去除简单标签 + import re + text_content = re.sub(r"<[^>]+>", "", content) + + if text_content.strip(): + extracted_chapters.append(text_content) + + if not extracted_chapters: + return [] + + full_text = "\n\n".join(extracted_chapters) + return self._text_splitter.split(full_text, source_file=source_file) diff --git a/src/core/splitters/registry.py b/src/core/splitters/registry.py index 70a1006..538ef4a 100644 --- a/src/core/splitters/registry.py +++ b/src/core/splitters/registry.py @@ -10,6 +10,7 @@ _DEFAULT_MAP: dict[str, str] = { ".pdf": "pdf", ".html": "html", ".htm": "html", + ".epub": "epub", } # 所有支持的扩展名集合（供外部遍历文件使用） @@ -64,6 +65,10 @@ def get_splitter( 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) diff --git a/tests/test_splitters_epub.py b/tests/test_splitters_epub.py new file mode 100644 index 0000000..c695f45 --- /dev/null +++ b/tests/test_splitters_epub.py @@ -0,0 +1,56 @@ +"""EPUBSplitter 测试.""" +import pytest + +ebooklib = pytest.importorskip("ebooklib", reason="ebooklib 未安装") + + +class TestEPUBSplitter: + """EPUBSplitter 测试（需 ebooklib）.""" + + def test_split_simple_epub(self, tmp_path): + """用 ebooklib 创建一个简单 EPUB 并测试分块.""" + from src.core.splitters.epub import EPUBSplitter + from ebooklib import epub + + epub_path = tmp_path / "test.epub" + + book = epub.EpubBook() + book.set_identifier("test123") + book.set_title("测试书名") + book.set_language("zh") + + chapter = epub.EpubHtml( + title="第一章", + file_name="chap01.xhtml", + lang="zh", + ) + chapter.content = "

第一章

这是EPUB电子书的内容。

第二段文字在这里。

" + book.add_item(chapter) + + book.toc = [epub.Link("chap01.xhtml", "第一章", "chap01")] + book.add_item(epub.EpubNcx()) + book.add_item(epub.EpubNav()) + book.spine = ["nav", chapter] + + epub.write_epub(str(epub_path), book) + + s = EPUBSplitter(max_size=500) + chunks = s.split(str(epub_path), source_file="test.epub") + assert len(chunks) >= 1 + all_text = "".join(c["content"] for c in chunks) + assert "第一章" in all_text + assert "EPUB电子书" in all_text + assert "第二段文字" in all_text + + def test_epub_has_split_method(self): + """EPUBSplitter 遵循 Splitter Protocol.""" + from src.core.splitters.epub import EPUBSplitter + s = EPUBSplitter() + assert hasattr(s, "split") + + def test_epub_missing_dependency_message(self): + """EPUBSplitter.split() 在未安装 ebooklib 时应给出明确提示.""" + from src.core.splitters.epub import EPUBSplitter + s = EPUBSplitter() + # 验证 split 方法存在且可调用（有明确的 import error 信息） + assert callable(s.split)