"""Markdown 文档解析与入库模块.""" import logging import re from dataclasses import dataclass from pathlib import Path from src.core.db import VectorDB from src.core.embedder import Embedder, batch_embed logger = logging.getLogger("md-vector-db") 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 始终前进（避免分隔符距 start 小于 overlap 时 start 回退导致死循环） next_start = break_point - self.overlap if self.overlap > 0 else break_point start = max(start + 1, next_start) 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 数量. 使用文件路径的 SHA256 前 12 位 + 文件名作为唯一标识， 避免不同目录下同名文件冲突. """ import hashlib path = Path(file_path).resolve() content = path.read_text(encoding="utf-8") # 用路径 hash 保证同名文件在不同目录下不冲突 path_hash = hashlib.sha256(str(path).encode()).hexdigest()[:12] file_name = f"{path_hash}_{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 # 分批嵌入 (避免大文档 OOM) 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) 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: with self.db.write_lock: existing = self.collection.get( where={"source_file": file_name} ) if existing and existing["ids"]: self.collection.delete(ids=existing["ids"]) except ValueError: pass # collection 为空时 ChromaDB 抛 ValueError except Exception: logger.exception("去重检查失败: %s", file_name)