Files
md-vector-db/docs/superpowers/plans/2026-07-10-multi-format-support.md
T

1015 lines
30 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# 多格式文档支持 — 实现计划
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** 将 md-vector-db 从"仅 Markdown"扩展为支持 `.txt``.pdf``.html` 的多格式文档向量数据库。
**Architecture:** 新建 `src/core/splitters/` 包,将 `MarkdownSplitter` + `Splitter(Protocol)``ingest.py` 移入。抽取 `BaseTextSplitter(ABC)` 作为通用文本切分基类,`TextSplitter`/`PDFSplitter`/`HTMLSplitter` 继承或组合它。`registry.py` 按文件扩展名自动选择 Splitter。`ingest.py` 保留兼容 import。
**Tech Stack:** pymupdf (PDF), beautifulsoup4 (HTML), 均为可选依赖。
---
## 文件清单
| 操作 | 文件 | 职责 |
|------|------|------|
| 新建 | `src/core/splitters/__init__.py` | 导出所有公共符号 |
| 新建 | `src/core/splitters/base.py` | Splitter(Protocol) + BaseTextSplitter(ABC) |
| 新建 | `src/core/splitters/markdown.py` | MarkdownSplitter(从 ingest.py 移入) |
| 新建 | `src/core/splitters/text.py` | TextSplitter(纯文本段落切分) |
| 新建 | `src/core/splitters/pdf.py` | PDFSplitterpymupdf 提取文字) |
| 新建 | `src/core/splitters/html.py` | HTMLSplitterbs4 去标签) |
| 新建 | `src/core/splitters/registry.py` | 扩展名→Splitter 映射 + get_splitter() |
| 修改 | `src/core/ingest.py` | 精简,用 registry 自动选择 splitter |
| 修改 | `pyproject.toml` | 添加可选依赖组 |
| 新建 | `tests/test_splitters.py` | 注册表 + TextSplitter 测试 |
| 新建 | `tests/test_splitters_pdf.py` | PDFSplitter 测试 |
| 新建 | `tests/test_splitters_html.py` | HTMLSplitter 测试 |
---
### Task 1: 创建 `splitters/base.py` — Protocol + 基类
**Files:**
- Create: `src/core/splitters/__init__.py`
- Create: `src/core/splitters/base.py`
- [ ] **Step 1: 创建包目录和 `__init__.py`**
```bash
mkdir -p src/core/splitters
```
- [ ] **Step 2: 写入 `splitters/__init__.py`**
```python
"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML."""
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.registry import get_splitter, register_splitter, SUPPORTED_SUFFIXES
__all__ = [
"Splitter",
"BaseTextSplitter",
"MarkdownSplitter",
"TextSplitter",
"get_splitter",
"register_splitter",
"SUPPORTED_SUFFIXES",
]
```
- [ ] **Step 3: 写入 `splitters/base.py`**
```python
"""Splitter Protocol 和文本切分基类."""
import re
from abc import ABC, abstractmethod
from typing import Protocol
class Splitter(Protocol):
"""文档分块器接口 — 将文本拆分为带元数据的 chunk 列表.
每个 chunk 为 dict: {"content": str, "section_title": str, "heading_level": int, ...}
"""
def split(self, text: str, source_file: str = "") -> list[dict]: ...
class BaseTextSplitter(ABC):
"""文本切分基类 — 提供段落切分和硬切逻辑,子类实现 split()."""
def __init__(self, max_size: int = 1000, overlap: int = 100):
self.max_size = max_size
self.overlap = overlap
@abstractmethod
def split(self, text: str, source_file: str = "") -> list[dict]: ...
def _split_by_paragraphs(
self, text: str, section_title: str = "", heading_level: int = 0
) -> list[dict]:
"""按段落边界拆分超长文本,若单段仍超长则硬切."""
paragraphs = re.split(r"\n\n+", text)
chunks = []
current = ""
for para in paragraphs:
if len(para) > self.max_size:
if current.strip():
chunks.append(self._make_chunk(current, section_title, heading_level))
current = ""
for sub in self._split_single_paragraph(para):
chunks.append(self._make_chunk(sub, section_title, heading_level))
continue
if len(current) + len(para) > self.max_size and current:
chunks.append(self._make_chunk(current, section_title, heading_level))
if self.overlap > 0 and len(current) > self.overlap:
current = current[-self.overlap:] + "\n\n" + para
else:
current = para
else:
current = f"{current}\n\n{para}" if current else para
if current.strip():
chunks.append(self._make_chunk(current, section_title, heading_level))
return chunks
def _split_single_paragraph(self, text: str) -> list[str]:
"""按字符边界硬切单个超长段落(带 overlap."""
parts = []
start = 0
while start < len(text):
end = start + self.max_size
if end >= len(text):
parts.append(text[start:].strip())
break
break_point = end
for sep in ("。", "", "", "\n", ". ", " "):
pos = text.rfind(sep, start, end)
if pos > start:
break_point = pos + len(sep)
break
part = text[start:break_point].strip()
if part:
parts.append(part)
next_start = break_point - self.overlap if self.overlap > 0 else break_point
start = max(start + 1, next_start)
return parts
@staticmethod
def _make_chunk(content: str, section_title: str, heading_level: int) -> dict:
return {
"content": content.strip(),
"section_title": section_title,
"heading_level": heading_level,
}
```
- [ ] **Step 4: 验证导入**
```bash
uv run python -c "from src.core.splitters.base import Splitter, BaseTextSplitter; print('OK')"
```
Expected: `OK`
- [ ] **Step 5: Commit**
```bash
git add src/core/splitters/__init__.py src/core/splitters/base.py
git commit -m "feat: 创建 splitters 包骨架 — Splitter Protocol + BaseTextSplitter 基类"
```
---
### Task 2: 将 MarkdownSplitter 移入 `splitters/markdown.py`
**Files:**
- Create: `src/core/splitters/markdown.py`
- Modify: `src/core/ingest.py` — 删除 MarkdownSplitter 类,添加兼容 import
- Modify: `tests/test_ingest.py` — 更新 import 路径
- [ ] **Step 1: 写入 `splitters/markdown.py`**
```python
"""Markdown 文档分块器."""
import re
from src.core.splitters.base import BaseTextSplitter
class MarkdownSplitter(BaseTextSplitter):
"""Markdown 混合分块器:先按标题拆,超长再按段落拆."""
def split(self, text: str, source_file: str = "") -> list[dict]:
if not text.strip():
return []
sections = self._split_by_headings(text)
chunks = []
for section in sections:
if len(section["content"]) <= self.max_size:
chunks.append(section)
else:
sub_chunks = self._split_by_paragraphs(
section["content"],
section["section_title"],
section["heading_level"],
)
chunks.extend(sub_chunks)
for i, chunk in enumerate(chunks):
chunk["source_file"] = source_file or chunk.get("source_file", "")
chunk["chunk_index"] = i
return chunks
def _split_by_headings(self, text: str) -> list[dict]:
heading_pattern = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
matches = list(heading_pattern.finditer(text))
if not matches:
return [{
"content": text.strip(),
"section_title": "",
"heading_level": 0,
}]
sections = []
for i, match in enumerate(matches):
level = len(match.group(1))
title = match.group(2).strip()
start = match.end()
end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
content = text[start:end].strip()
if content:
sections.append({
"content": f"{match.group(0)}\n{content}",
"section_title": title,
"heading_level": level,
})
if matches and matches[0].start() > 0:
preamble = text[:matches[0].start()].strip()
if preamble:
sections.insert(0, {
"content": preamble,
"section_title": "",
"heading_level": 0,
})
return sections
```
- [ ] **Step 2: 从 `ingest.py` 中删除 `Splitter` Protocol 和 `MarkdownSplitter` 类,替换为兼容 import**
`ingest.py` 中删除第 14-174 行(Splitter Protocol + MarkdownSplitter 全部代码),替换为:
```python
import logging
import re
from pathlib import Path
from src.core.db import VectorDB
from src.core.embedder import Embedder, batch_embed
from src.core.splitters.registry import get_splitter
from src.core.splitters.markdown import MarkdownSplitter # 兼容旧 import
from src.core.splitters.base import Splitter # 兼容旧 import
logger = logging.getLogger("md-vector-db")
```
- [ ] **Step 3: 更新 `tests/test_ingest.py` 的 import**
```python
# 将原有 from src.core.ingest import MarkdownSplitter
# 改为:
from src.core.splitters import MarkdownSplitter
```
- [ ] **Step 4: 运行测试验证**
```bash
uv run pytest tests/test_ingest.py tests/test_search.py -v
```
Expected: all 12+ tests PASS
- [ ] **Step 5: Commit**
```bash
git add src/core/splitters/markdown.py src/core/ingest.py tests/test_ingest.py
git commit -m "refactor: 将 MarkdownSplitter 移入 splitters 包,ingest.py 保留兼容 import"
```
---
### Task 3: 创建 `TextSplitter` 和 `registry.py`
**Files:**
- Create: `src/core/splitters/text.py`
- Create: `src/core/splitters/registry.py`
- [ ] **Step 1: 写入 `splitters/text.py`**
```python
"""纯文本分块器 — 按段落双换行切分."""
from src.core.splitters.base import BaseTextSplitter
class TextSplitter(BaseTextSplitter):
"""纯文本分块器:按 \n\n 切段落,超长按标点硬切."""
def split(self, text: str, source_file: str = "") -> list[dict]:
if not text.strip():
return []
chunks = self._split_by_paragraphs(text)
for i, chunk in enumerate(chunks):
chunk["source_file"] = source_file
chunk["chunk_index"] = i
return chunks
```
- [ ] **Step 2: 写入 `splitters/registry.py`**
```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",
}
# 所有支持的扩展名集合(供外部遍历文件使用)
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)
# 回退
from src.core.splitters.text import TextSplitter
return TextSplitter(max_size=max_size, overlap=overlap)
```
- [ ] **Step 3: 更新 `splitters/__init__.py`**(追加 TextSplitter 和 SUPPORTED_SUFFIXES 导出)
```python
"""文档分块器包 — 支持 Markdown / 纯文本 / PDF / HTML."""
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.registry import get_splitter, register_splitter, SUPPORTED_SUFFIXES
__all__ = [
"Splitter",
"BaseTextSplitter",
"MarkdownSplitter",
"TextSplitter",
"get_splitter",
"register_splitter",
"SUPPORTED_SUFFIXES",
]
```
- [ ] **Step 4: 验证注册表**
```bash
uv run python -c "
from src.core.splitters.registry import get_splitter, SUPPORTED_SUFFIXES
print('Suffixes:', sorted(SUPPORTED_SUFFIXES))
s = get_splitter('test.txt')
print('TextSplitter:', type(s).__name__)
s = get_splitter('doc.md')
print('MarkdownSplitter:', type(s).__name__)
s = get_splitter('unknown.xyz')
print('Fallback:', type(s).__name__)
"
```
Expected: `TextSplitter`, `MarkdownSplitter`, `Fallback: TextSplitter`
- [ ] **Step 5: Commit**
```bash
git add src/core/splitters/text.py src/core/splitters/registry.py src/core/splitters/__init__.py
git commit -m "feat: 添加 TextSplitter 和 registry 自动选择机制"
```
---
### Task 4: 更新 `ingest.py` — 使用 registry 自动选择 Splitter
**Files:**
- Modify: `src/core/ingest.py`
- [ ] **Step 1: 更新 `DocumentIngestor.__init__`**
将默认 splitter 从 `MarkdownSplitter()` 改为 `None`None 时由 `ingest_file` 自动选择):
```python
def __init__(
self,
db: VectorDB,
embedder: Embedder,
collection_name: str,
splitter: Splitter | None = None,
):
self.db = db
self.embedder = embedder
self.collection_name = collection_name
self.splitter = splitter # None = 按扩展名自动选择
```
- [ ] **Step 2: 更新 `ingest_file` 使用 registry**
```python
def ingest_file(self, file_path: str) -> int:
"""入库单个文件, 返回 chunk 数量.
根据文件扩展名自动选择 Splitter.md→MarkdownSplitter, .pdf→PDFSplitter 等)。
使用文件路径的 SHA256 前 12 位 + 文件名作为唯一标识。
"""
import hashlib
path = Path(file_path).resolve()
path_hash = hashlib.sha256(str(path).encode()).hexdigest()[:12]
file_name = f"{path_hash}_{path.name}"
splitter = self.splitter or get_splitter(file_path)
content = path.read_text(encoding="utf-8")
return self._ingest_with_splitter(content, file_name, splitter)
def _ingest_with_splitter(self, content: str, file_name: str, splitter: Splitter) -> int:
"""使用指定 splitter 分块并入库."""
self._remove_by_source(file_name)
chunks = splitter.split(content, source_file=file_name)
if not chunks:
return 0
texts = [c["content"] for c in chunks]
embeddings = batch_embed(self.embedder, texts)
ids = [f"{file_name}_{i}" for i in range(len(chunks))]
metadatas = [
{
"source_file": c.get("source_file", file_name),
"section_title": c.get("section_title", ""),
"heading_level": c.get("heading_level", 0),
"chunk_index": i,
}
for i, c in enumerate(chunks)
]
with self.db.write_lock:
self.collection.add(
ids=ids, embeddings=embeddings, documents=texts, metadatas=metadatas,
)
return len(chunks)
```
- [ ] **Step 3: 更新 `ingest_content` 保持向后兼容**
```python
def ingest_content(self, content: str, file_name: str) -> int:
"""入库内容(无需实际文件)。若未指定 splitter,默认用 MarkdownSplitter."""
splitter = self.splitter or MarkdownSplitter()
return self._ingest_with_splitter(content, file_name, splitter)
```
- [ ] **Step 4: 更新 `ingest_directory` 支持多格式**
```python
def ingest_directory(self, dir_path: str) -> dict[str, int]:
"""入库目录下所有支持的文档格式."""
from src.core.splitters.registry import SUPPORTED_SUFFIXES
results = {}
for f in Path(dir_path).rglob("*"):
if f.suffix.lower() in SUPPORTED_SUFFIXES:
count = self.ingest_file(str(f))
results[f.name] = count
return results
```
- [ ] **Step 5: 运行现有测试确保无回归**
```bash
uv run pytest tests/ -v
```
Expected: all 70 tests PASS
- [ ] **Step 6: Commit**
```bash
git add src/core/ingest.py
git commit -m "feat: ingest.py 使用 registry 自动选择 Splitteringest_directory 支持多格式"
```
---
### Task 5: 创建 `PDFSplitter`
**Files:**
- Create: `src/core/splitters/pdf.py`
- [ ] **Step 1: 写入 `splitters/pdf.py`**
```python
"""PDF 文档分块器 — 使用 pymupdf 提取文字后委托 TextSplitter."""
import logging
from src.core.splitters.base import Splitter, BaseTextSplitter
from src.core.splitters.text import TextSplitter
logger = logging.getLogger("md-vector-db")
class PDFSplitter:
"""PDF 分块器:pymupdf 提取文字 → TextSplitter 分块.
实现 Splitter Protocol,内部组合 TextSplitter 实例。
"""
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 文件路径(而非文本内容)
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)
```
- [ ] **Step 2: 修改 `ingest_file` 处理 PDF 二进制文件**
PDF 文件不能像文本文件那样 `read_text()`。需要在 `ingest_file` 中特殊处理:
```python
def ingest_file(self, file_path: str) -> int:
import hashlib
path = Path(file_path).resolve()
path_hash = hashlib.sha256(str(path).encode()).hexdigest()[:12]
file_name = f"{path_hash}_{path.name}"
splitter = self.splitter or get_splitter(file_path)
# PDF 二进制文件特殊处理:传入路径给 Splitter
suffix = path.suffix.lower()
if suffix in (".pdf",):
# PDFSplitter.split() 接收文件路径而非文本内容
chunks = splitter.split(str(path), source_file=file_name)
return self._ingest_chunks(chunks, file_name)
content = path.read_text(encoding="utf-8")
return self._ingest_content_with_splitter(content, file_name, splitter)
```
并在类中添加 `_ingest_chunks` 辅助方法(与 `_ingest_with_splitter` 的后半段相同)。
- [ ] **Step 3: 安装 pymupdf 并验证导入**
```bash
uv sync --extra pdf
uv run python -c "from src.core.splitters.pdf import PDFSplitter; print('OK')"
```
Expected: `OK`
- [ ] **Step 4: Commit**
```bash
git add src/core/splitters/pdf.py src/core/ingest.py
git commit -m "feat: 添加 PDFSplitter — pymupdf 提取文字后分块"
```
---
### Task 6: 创建 `HTMLSplitter`
**Files:**
- Create: `src/core/splitters/html.py`
- [ ] **Step 1: 写入 `splitters/html.py`**
```python
"""HTML 文档分块器 — 使用 BeautifulSoup 去标签后委托 TextSplitter."""
import logging
from src.core.splitters.text import TextSplitter
logger = logging.getLogger("md-vector-db")
class HTMLSplitter:
"""HTML 分块器:bs4 去标签提取文字 → TextSplitter 分块.
实现 Splitter Protocol,内部组合 TextSplitter 实例。
"""
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]:
"""从 HTML 文本去标签并分块.
Args:
text: HTML 文本内容
source_file: 来源文件名
"""
try:
from bs4 import BeautifulSoup
except ImportError:
raise ImportError(
"HTML 支持需要 beautifulsoup4 库. 请执行: uv sync --extra html"
)
try:
soup = BeautifulSoup(text, "html.parser")
# 移除 script/style 标签
for tag in soup(["script", "style"]):
tag.decompose()
plain_text = soup.get_text(separator="\n")
except Exception as e:
logger.error("HTML 解析失败: %s%s", source_file, e)
raise ValueError(f"HTML 解析失败: {e}") from e
if not plain_text.strip():
return []
return self._text_splitter.split(plain_text, source_file=source_file)
```
- [ ] **Step 2: 安装 beautifulsoup4 并验证导入**
```bash
uv sync --extra html
uv run python -c "from src.core.splitters.html import HTMLSplitter; print('OK')"
```
Expected: `OK`
- [ ] **Step 3: Commit**
```bash
git add src/core/splitters/html.py
git commit -m "feat: 添加 HTMLSplitter — bs4 去标签后分块"
```
---
### Task 7: 更新 `pyproject.toml` 可选依赖
**Files:**
- Modify: `pyproject.toml`
- [ ] **Step 1: 添加可选依赖组**
`pyproject.toml``[project.optional-dependencies]` 段改为:
```toml
[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"]
```
- [ ] **Step 2: 验证依赖安装**
```bash
uv sync --extra all
uv run python -c "import fitz; from bs4 import BeautifulSoup; print('OK')"
```
Expected: `OK`
- [ ] **Step 3: Commit**
```bash
git add pyproject.toml
git commit -m "feat: 添加 pdf/html/all 可选依赖组"
```
---
### Task 8: 编写测试 — 注册表 + TextSplitter
**Files:**
- Create: `tests/test_splitters.py`
- [ ] **Step 1: 写入 `tests/test_splitters.py`**
```python
"""Splitter 注册表和 TextSplitter 测试."""
import pytest
from src.core.splitters import TextSplitter, MarkdownSplitter, get_splitter, register_splitter, SUPPORTED_SUFFIXES
class TestTextSplitter:
"""TextSplitter 纯文本分块测试."""
def test_empty_text(self):
s = TextSplitter()
assert s.split("") == []
assert s.split(" \n\n ") == []
def test_short_text_single_chunk(self):
s = TextSplitter(max_size=1000)
chunks = s.split("这是一段短文本。", source_file="test.txt")
assert len(chunks) == 1
assert chunks[0]["source_file"] == "test.txt"
assert chunks[0]["content"] == "这是一段短文本。"
def test_long_paragraph_split(self):
s = TextSplitter(max_size=50, overlap=10)
long_text = "这是第一句。" * 20
chunks = s.split(long_text, source_file="long.txt")
assert len(chunks) > 1
for c in chunks:
assert len(c["content"]) <= 60 # max_size + 少许容差
def test_paragraph_boundary_split(self):
s = TextSplitter(max_size=100)
text = "短段落A。\n\n短段落B。\n\n短段落C。"
chunks = s.split(text)
assert len(chunks) >= 1
assert all("content" in c for c in chunks)
def test_chunk_metadata(self):
s = TextSplitter()
chunks = s.split("测试内容。", source_file="doc.txt")
assert chunks[0]["source_file"] == "doc.txt"
assert chunks[0]["section_title"] == ""
assert chunks[0]["heading_level"] == 0
assert chunks[0]["chunk_index"] == 0
class TestRegistry:
"""注册表测试."""
def test_get_splitter_for_md(self):
s = get_splitter("doc.md")
assert isinstance(s, MarkdownSplitter)
def test_get_splitter_for_txt(self):
s = get_splitter("notes.txt")
assert isinstance(s, TextSplitter)
def test_get_splitter_fallback(self):
s = get_splitter("data.xyz")
assert isinstance(s, TextSplitter)
def test_supported_suffixes(self):
assert ".md" in SUPPORTED_SUFFIXES
assert ".txt" in SUPPORTED_SUFFIXES
assert ".pdf" in SUPPORTED_SUFFIXES
assert ".html" in SUPPORTED_SUFFIXES
def test_custom_register(self):
class FakeSplitter:
def __init__(self, max_size=1000, overlap=100): pass
def split(self, text, source_file=""): return []
register_splitter(".fake", FakeSplitter)
s = get_splitter("test.fake")
assert isinstance(s, FakeSplitter)
```
- [ ] **Step 2: 运行测试**
```bash
uv run pytest tests/test_splitters.py -v
```
Expected: 8 tests PASS
- [ ] **Step 3: Commit**
```bash
git add tests/test_splitters.py
git commit -m "test: 添加 Splitter 注册表和 TextSplitter 测试"
```
---
### Task 9: 编写 PDF/HTML Splitter 测试
**Files:**
- Create: `tests/test_splitters_pdf.py`
- Create: `tests/test_splitters_html.py`
- [ ] **Step 1: 写入 `tests/test_splitters_pdf.py`**
```python
"""PDFSplitter 测试."""
import pytest
from pathlib import Path
pymupdf = pytest.importorskip("fitz", reason="pymupdf 未安装")
class TestPDFSplitter:
"""PDFSplitter 测试(需 pymupdf."""
def test_split_simple_pdf(self, tmp_path):
"""用 pymupdf 创建一个简单 PDF 并测试分块."""
from src.core.splitters.pdf import PDFSplitter
import fitz
pdf_path = tmp_path / "test.pdf"
doc = fitz.open()
doc.new_page().insert_text((72, 72), "这是PDF文档内容。\n\n第二段文字。")
doc.save(str(pdf_path))
doc.close()
s = PDFSplitter(max_size=500)
chunks = s.split(str(pdf_path), source_file="test.pdf")
assert len(chunks) >= 1
assert "PDF文档内容" in chunks[0]["content"]
def test_pdf_missing_lib_error(self, monkeypatch):
"""未安装 pymupdf 时的错误提示."""
# 此测试在已安装 pymupdf 时跳过语义检查,仅验证 split 方法存在
from src.core.splitters.pdf import PDFSplitter
s = PDFSplitter()
assert hasattr(s, "split")
```
- [ ] **Step 2: 写入 `tests/test_splitters_html.py`**
```python
"""HTMLSplitter 测试."""
import pytest
bs4 = pytest.importorskip("bs4", reason="beautifulsoup4 未安装")
class TestHTMLSplitter:
"""HTMLSplitter 测试(需 beautifulsoup4."""
def test_split_simple_html(self):
from src.core.splitters.html import HTMLSplitter
html = "<html><body><h1>标题</h1><p>这是段落内容。</p><p>第二段。</p></body></html>"
s = HTMLSplitter(max_size=500)
chunks = s.split(html, source_file="test.html")
assert len(chunks) >= 1
# 验证去标签后的内容
all_text = "".join(c["content"] for c in chunks)
assert "标题" in all_text
assert "段落内容" in all_text
assert "第二段" in all_text
def test_strips_script_and_style(self):
from src.core.splitters.html import HTMLSplitter
html = """
<html><head><style>.a{color:red}</style><script>alert('xss')</script></head>
<body><p>可见内容。</p></body></html>
"""
s = HTMLSplitter()
chunks = s.split(html, source_file="test.html")
all_text = "".join(c["content"] for c in chunks)
assert "可见内容" in all_text
assert "alert" not in all_text
assert ".a{color:red}" not in all_text
def test_empty_html(self):
from src.core.splitters.html import HTMLSplitter
s = HTMLSplitter()
assert s.split("<html></html>") == []
assert s.split("") == []
```
- [ ] **Step 3: 运行测试**
```bash
uv run pytest tests/test_splitters_pdf.py tests/test_splitters_html.py -v
```
Expected: all tests PASS (pymupdf + bs4 已安装)
- [ ] **Step 4: 运行全量测试确保无回归**
```bash
uv run pytest tests/ -v
```
Expected: all tests PASS
- [ ] **Step 5: Commit**
```bash
git add tests/test_splitters_pdf.py tests/test_splitters_html.py
git commit -m "test: 添加 PDFSplitter 和 HTMLSplitter 测试"
```
---
### Task 10: 更新文档
**Files:**
- Modify: `CLAUDE.md`
- Modify: `README.md`
- [ ] **Step 1: 更新 `CLAUDE.md` 的架构描述**
在 "架构" 段落后增加 splitters 说明:
```
src/core/splitters/ # 文档分块器包(新增)
├── base.py # Splitter Protocol + BaseTextSplitter 基类
├── markdown.py # MarkdownSplitter(标题+段落混合分块)
├── text.py # TextSplitter(纯文本段落切分)
├── pdf.py # PDFSplitterpymupdf 提取文字)
├── html.py # HTMLSplitterbs4 去标签)
└── registry.py # 扩展名→Splitter 自动选择
```
- [ ] **Step 2: 更新 `README.md`**
在 "功能特性" 列表中添加 "多格式文档:支持 .md / .txt / .pdf / .html,按扩展名自动选择分块器",并提供安装命令:
```bash
uv sync --extra all # 安装 PDF + HTML 支持
```
- [ ] **Step 3: Commit**
```bash
git add CLAUDE.md README.md
git commit -m "docs: 更新文档记录多格式支持特性"
```
---
## 执行顺序
```
Task 1 (base.py 骨架)
→ Task 2 (MarkdownSplitter 迁移)
→ Task 3 (TextSplitter + registry)
→ Task 4 (ingest.py 更新)
→ Task 5 (PDFSplitter)
→ Task 6 (HTMLSplitter)
→ Task 7 (pyproject.toml)
→ Task 8 (测试 — 注册表+TextSplitter)
→ Task 9 (测试 — PDF+HTML)
→ Task 10 (文档)
```
Task 5/6 可并行,Task 8/9 可并行。