Serendipity
832201186d
H1: AppConfig 自定义 __init__ 改用 from_dict() 类方法 H2: list_collections_with_stats 改为从 ChromaDB 直接查询 H3: RateLimiter 过期 key 自动清理, 防止内存泄漏 M1: delete_by_source 区分 ValueError 与真实异常, 记日志 M2: VectorDB 新增 list_collections() 封装方法 M3: _remove_by_source 异常记日志, 不再静默吞掉 M4: CLI 集合回退支持 MD_VECTOR_DB_COLLECTION 环境变量 L1: DEFAULT_CONFIG_PATH 自动从项目根目录解析 L2: ingest 命令内重复 import 移至模块顶部 L3: 新增死循环回归测试 + 密集分隔符分块测试 L4: 统一 logger 名称为 md-vector-db 删除旧版审计文档 测试: 48 passed
252 lines
8.7 KiB
Python
252 lines
8.7 KiB
Python
"""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)
|