md-vector-db/src/core/splitters/epub.py

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