diff --git a/src/core/ingest.py b/src/core/ingest.py index 165564a..d870ebf 100644 --- a/src/core/ingest.py +++ b/src/core/ingest.py @@ -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: