feat: add Markdown splitter and document ingestor
- MarkdownSplitter: hybrid splitting by headings then paragraphs - DocumentIngestor: read MD -> split -> embed -> store with dedup - All 6 tests passing
This commit is contained in:
@@ -0,0 +1,235 @@
|
||||
"""Markdown 文档解析与入库模块."""
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from src.core.db import VectorDB
|
||||
from src.core.embedder import Embedder
|
||||
|
||||
|
||||
class MarkdownSplitter:
|
||||
"""Markdown 混合分块器：先按标题拆，超长再按段落拆."""
|
||||
|
||||
def __init__(self, max_size: int = 1000, overlap: int = 100):
|
||||
self.max_size = max_size
|
||||
self.overlap = overlap
|
||||
|
||||
def split(self, text: str, source_file: str = "") -> list[dict]:
|
||||
"""将 Markdown 文本拆分为带元数据的 chunk 列表."""
|
||||
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)
|
||||
|
||||
# 为所有 chunk 补充 source_file 和 chunk_index
|
||||
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]:
|
||||
"""按 Markdown 标题拆分."""
|
||||
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
|
||||
|
||||
def _split_by_paragraphs(
|
||||
self, text: str, section_title: str, heading_level: int
|
||||
) -> list[dict]:
|
||||
"""按段落边界拆分超长章节.
|
||||
|
||||
优先在段落边界拆分，若单个段落仍超长则按字符硬切。
|
||||
"""
|
||||
paragraphs = re.split(r"\n\n+", text)
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for para in paragraphs:
|
||||
# 单一段落超出 max_size 时直接硬切
|
||||
if len(para) > self.max_size:
|
||||
# 先 flush 当前累积
|
||||
if current.strip():
|
||||
chunks.append({
|
||||
"content": current.strip(),
|
||||
"section_title": section_title,
|
||||
"heading_level": heading_level,
|
||||
})
|
||||
current = ""
|
||||
# 硬切该段落
|
||||
for sub in self._split_single_paragraph(para):
|
||||
chunks.append({
|
||||
"content": sub,
|
||||
"section_title": section_title,
|
||||
"heading_level": heading_level,
|
||||
})
|
||||
continue
|
||||
|
||||
if len(current) + len(para) > self.max_size and current:
|
||||
chunks.append({
|
||||
"content": current.strip(),
|
||||
"section_title": section_title,
|
||||
"heading_level": heading_level,
|
||||
})
|
||||
# overlap: 保留上一块的末尾部分
|
||||
if self.overlap > 0 and len(current) > self.overlap:
|
||||
current = current[-self.overlap:] + "\n\n" + para
|
||||
else:
|
||||
current = para
|
||||
else:
|
||||
if current:
|
||||
current += "\n\n" + para
|
||||
else:
|
||||
current = para
|
||||
|
||||
if current.strip():
|
||||
chunks.append({
|
||||
"content": current.strip(),
|
||||
"section_title": section_title,
|
||||
"heading_level": 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)
|
||||
start = break_point - self.overlap if self.overlap > 0 else break_point
|
||||
return parts
|
||||
|
||||
|
||||
class DocumentIngestor:
|
||||
"""文档入库器: 读取 MD 文件 → 分块 → 嵌入 → 入库."""
|
||||
|
||||
def __init__(self, db: VectorDB, embedder: Embedder, collection_name: str):
|
||||
self.db = db
|
||||
self.embedder = embedder
|
||||
self.collection_name = collection_name
|
||||
self.splitter = MarkdownSplitter()
|
||||
|
||||
@property
|
||||
def collection(self):
|
||||
return self.db.get_or_create_collection(self.collection_name)
|
||||
|
||||
def ingest_file(self, file_path: str) -> int:
|
||||
"""入库单个 Markdown 文件, 返回 chunk 数量."""
|
||||
path = Path(file_path)
|
||||
content = path.read_text(encoding="utf-8")
|
||||
file_name = path.name
|
||||
|
||||
return self.ingest_content(content, file_name)
|
||||
|
||||
def ingest_content(self, content: str, file_name: str) -> int:
|
||||
"""入库 Markdown 内容（无需实际文件）."""
|
||||
# 去重：先删旧 chunks
|
||||
self._remove_by_source(file_name)
|
||||
|
||||
# 分块
|
||||
chunks = self.splitter.split(content, source_file=file_name)
|
||||
if not chunks:
|
||||
return 0
|
||||
|
||||
# 嵌入
|
||||
texts = [c["content"] for c in chunks]
|
||||
embeddings = self.embedder.embed(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)
|
||||
]
|
||||
|
||||
self.collection.add(
|
||||
ids=ids,
|
||||
embeddings=embeddings,
|
||||
documents=texts,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
|
||||
return len(chunks)
|
||||
|
||||
def ingest_directory(self, dir_path: str) -> dict[str, int]:
|
||||
"""入库目录下所有 Markdown 文件."""
|
||||
results = {}
|
||||
for md_file in Path(dir_path).rglob("*.md"):
|
||||
count = self.ingest_file(str(md_file))
|
||||
results[md_file.name] = count
|
||||
return results
|
||||
|
||||
def _remove_by_source(self, file_name: str) -> None:
|
||||
"""按 source_file 删除已有 chunks."""
|
||||
try:
|
||||
existing = self.collection.get(
|
||||
where={"source_file": file_name}
|
||||
)
|
||||
if existing and existing["ids"]:
|
||||
self.collection.delete(ids=existing["ids"])
|
||||
except Exception:
|
||||
pass # collection 为空时 get 可能抛异常
|
||||
@@ -0,0 +1,87 @@
|
||||
"""文档入库测试."""
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from src.core.ingest import MarkdownSplitter, DocumentIngestor
|
||||
|
||||
|
||||
class TestMarkdownSplitter:
|
||||
"""Markdown 分块器测试."""
|
||||
|
||||
@pytest.fixture
|
||||
def splitter(self):
|
||||
return MarkdownSplitter(max_size=1000, overlap=100)
|
||||
|
||||
def test_split_simple_document(self, splitter):
|
||||
"""简单文档按标题拆分."""
|
||||
md = """# 标题一
|
||||
这是第一段内容。
|
||||
|
||||
## 标题二
|
||||
这是第二段内容。
|
||||
|
||||
# 标题三
|
||||
这是第三段内容。"""
|
||||
chunks = splitter.split(md, source_file="test.md")
|
||||
assert len(chunks) >= 3
|
||||
# 每个 chunk 有元数据
|
||||
for chunk in chunks:
|
||||
assert "content" in chunk
|
||||
assert chunk["source_file"] == "test.md"
|
||||
|
||||
def test_chunk_has_heading_metadata(self, splitter):
|
||||
"""chunk 附带标题元数据."""
|
||||
md = "# 配置指南\n这里是配置说明。"
|
||||
chunks = splitter.split(md, source_file="config.md")
|
||||
assert len(chunks) >= 1
|
||||
title = chunks[0]["section_title"]
|
||||
assert "配置指南" in title or title == ""
|
||||
|
||||
def test_long_section_is_split(self, splitter):
|
||||
"""超长章节被进一步拆分."""
|
||||
# 创建一个超过 max_size 的段落
|
||||
long_text = "这是很长的文本。" * 300 # ~3000 字符
|
||||
md = f"# 长章节\n{long_text}"
|
||||
small_splitter = MarkdownSplitter(max_size=500, overlap=50)
|
||||
chunks = small_splitter.split(md, source_file="long.md")
|
||||
assert len(chunks) > 1
|
||||
|
||||
def test_empty_document(self, splitter):
|
||||
"""空文档返回空列表."""
|
||||
chunks = splitter.split("", source_file="empty.md")
|
||||
assert chunks == []
|
||||
|
||||
def test_code_blocks_preserved(self, splitter):
|
||||
"""代码块不被拆分."""
|
||||
md = """# 代码示例
|
||||
```python
|
||||
def hello():
|
||||
print("world")
|
||||
```
|
||||
|
||||
"""
|
||||
chunks = splitter.split(md, source_file="code.md")
|
||||
assert len(chunks) >= 1
|
||||
# 代码块内容应在某个 chunk 中
|
||||
all_content = " ".join(c["content"] for c in chunks)
|
||||
assert "def hello()" in all_content
|
||||
|
||||
|
||||
class TestDocumentIngestor:
|
||||
"""文档入库器测试."""
|
||||
|
||||
@pytest.fixture
|
||||
def temp_md_dir(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
# 创建测试 Markdown 文件
|
||||
md_path = Path(d) / "test.md"
|
||||
md_path.write_text("# 测试\n这是测试内容。", encoding="utf-8")
|
||||
yield d
|
||||
|
||||
def test_read_markdown_file(self, temp_md_dir):
|
||||
"""读取 Markdown 文件."""
|
||||
content = Path(temp_md_dir + "/test.md").read_text(encoding="utf-8")
|
||||
assert "测试" in content
|
||||
assert "这是测试内容" in content
|
||||
Reference in New Issue
Block a user