1015 lines
30 KiB
Markdown
1015 lines
30 KiB
Markdown
# 多格式文档支持 — 实现计划
|
||
|
||
> **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` | PDFSplitter(pymupdf 提取文字) |
|
||
| 新建 | `src/core/splitters/html.py` | HTMLSplitter(bs4 去标签) |
|
||
| 新建 | `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 自动选择 Splitter,ingest_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 # PDFSplitter(pymupdf 提取文字)
|
||
├── html.py # HTMLSplitter(bs4 去标签)
|
||
└── 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 可并行。
|