feat: ingest.py 使用 registry 自动选择 Splitter,ingest_directory 支持多格式
This commit is contained in:
+33
-17
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user