feat: ingest.py 使用 registry 自动选择 Splitter,ingest_directory 支持多格式

This commit is contained in:
2026-07-10 14:19:03 +08:00
parent 13d63ba6ff
commit cde4ae7ddc
+33 -17
View File
@@ -7,6 +7,7 @@ from src.core.db import VectorDB
from src.core.embedder import Embedder, batch_embed
from src.core.splitters.markdown import MarkdownSplitter # 兼容旧 import 路径
from src.core.splitters.base import Splitter # 兼容旧 import 路径
from src.core.splitters.registry import SUPPORTED_SUFFIXES, get_splitter
logger = logging.getLogger("md-vector-db")
@@ -31,35 +32,44 @@ class DocumentIngestor:
return self.db.get_or_create_collection(self.collection_name)
def ingest_file(self, file_path: str) -> int:
"""入库单个 Markdown 文件, 返回 chunk 数量.
"""入库单个文件, 返回 chunk 数量.
使用文件路径的 SHA256 前 12 位 + 文件名作为唯一标识,
避免不同目录下同名文件冲突.
根据文件扩展名自动选择 Splitter.md→MarkdownSplitter, .txt→TextSplitter, .pdf→PDFSplitter 等)。
使用文件路径的 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)
splitter = self.splitter or get_splitter(
file_path, max_size=1000, overlap=100
)
def ingest_content(self, content: str, file_name: str) -> int:
"""入库 Markdown 内容(无需实际文件)."""
# 去重:先删旧 chunks
# PDF 文件特殊处理:splitter 内部读取二进制内容
suffix = path.suffix.lower()
if suffix in (".pdf",):
chunks = splitter.split(str(path), source_file=file_name)
return self._add_chunks(chunks, file_name)
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) -> int:
"""分块 + 嵌入 + 入库(文本文件通用路径)."""
self._remove_by_source(file_name)
# 分块
chunks = self.splitter.split(content, source_file=file_name)
chunks = splitter.split(content, source_file=file_name)
return self._add_chunks(chunks, file_name)
def _add_chunks(self, chunks: list[dict], file_name: str) -> int:
"""嵌入 + 写入 ChromaDB(分块已完成)."""
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 = [
{
@@ -81,12 +91,18 @@ class DocumentIngestor:
return len(chunks)
def ingest_content(self, content: str, file_name: str) -> int:
"""入库 Markdown 内容(无需实际文件)。若未指定 splitter,默认用 MarkdownSplitter."""
splitter = self.splitter or MarkdownSplitter()
return self._ingest_with_splitter(content, file_name, splitter)
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
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
def _remove_by_source(self, file_name: str) -> None: