diff --git a/src/core/ingest.py b/src/core/ingest.py new file mode 100644 index 0000000..0c6e02f --- /dev/null +++ b/src/core/ingest.py @@ -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 可能抛异常 diff --git a/tests/test_ingest.py b/tests/test_ingest.py new file mode 100644 index 0000000..d2c718a --- /dev/null +++ b/tests/test_ingest.py @@ -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