feat: DocumentIngestor 支持增量入库(FileTracker)
This commit is contained in:
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+24
-4
@@ -6,6 +6,7 @@ from pathlib import Path
|
||||
from src.core.config import ChunkConfig
|
||||
from src.core.db import VectorDB
|
||||
from src.core.embedder import Embedder, batch_embed
|
||||
from src.core.file_tracker import FileTracker
|
||||
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
|
||||
@@ -23,28 +24,41 @@ class DocumentIngestor:
|
||||
collection_name: str,
|
||||
splitter: Splitter | None = None,
|
||||
chunk_config: ChunkConfig | None = None,
|
||||
file_tracker: FileTracker | None = None,
|
||||
):
|
||||
self.db = db
|
||||
self.embedder = embedder
|
||||
self.collection_name = collection_name
|
||||
self.splitter = splitter or MarkdownSplitter()
|
||||
self.chunk_config = chunk_config or ChunkConfig()
|
||||
self.file_tracker = file_tracker # None = 不使用增量功能
|
||||
|
||||
@property
|
||||
def collection(self):
|
||||
return self.db.get_or_create_collection(self.collection_name)
|
||||
|
||||
def ingest_file(self, file_path: str) -> int:
|
||||
def ingest_file(self, file_path: str, incremental: bool = False, force: bool = False) -> int:
|
||||
"""入库单个文件, 返回 chunk 数量.
|
||||
|
||||
根据文件扩展名自动选择 Splitter(.md→MarkdownSplitter, .txt→TextSplitter, .pdf→PDFSplitter 等)。
|
||||
使用文件路径的 SHA256 前 12 位 + 文件名作为唯一标识。
|
||||
|
||||
Args:
|
||||
file_path: 文件路径
|
||||
incremental: 启用增量模式(需 file_tracker 已注入)
|
||||
force: 强制重新入库(忽略增量检查)
|
||||
"""
|
||||
import hashlib
|
||||
path = Path(file_path).resolve()
|
||||
path_hash = hashlib.sha256(str(path).encode()).hexdigest()[:12]
|
||||
file_name = f"{path_hash}_{path.name}"
|
||||
|
||||
# 增量模式:检查是否需要重新入库
|
||||
if incremental and not force and self.file_tracker is not None:
|
||||
if not self.file_tracker.is_stale(str(path)):
|
||||
logger.debug("跳过未变更文件: %s", path.name)
|
||||
return 0
|
||||
|
||||
splitter = self.splitter or get_splitter(
|
||||
file_path,
|
||||
max_size=self.chunk_config.max_size,
|
||||
@@ -55,10 +69,16 @@ class DocumentIngestor:
|
||||
suffix = path.suffix.lower()
|
||||
if suffix in (".pdf", ".epub"):
|
||||
chunks = splitter.split(str(path), source_file=file_name)
|
||||
return self._add_chunks(chunks, file_name)
|
||||
result = self._add_chunks(chunks, file_name)
|
||||
else:
|
||||
content = path.read_text(encoding="utf-8")
|
||||
result = self._ingest_with_splitter(content, file_name, splitter)
|
||||
|
||||
content = path.read_text(encoding="utf-8")
|
||||
return self._ingest_with_splitter(content, file_name, splitter)
|
||||
# 入库成功后更新 tracker
|
||||
if self.file_tracker is not None:
|
||||
self.file_tracker.mark_ingested(str(path))
|
||||
|
||||
return result
|
||||
|
||||
def _ingest_with_splitter(self, content: str, file_name: str, splitter) -> int:
|
||||
"""分块 + 嵌入 + 入库(文本文件通用路径)."""
|
||||
|
||||
@@ -221,3 +221,77 @@ class TestIngestDirectory:
|
||||
|
||||
results = ingestor.ingest_directory(str(tmp_path))
|
||||
assert results == {}
|
||||
|
||||
|
||||
class TestIncrementalIngest:
|
||||
"""增量入库测试."""
|
||||
|
||||
def test_ingest_file_incremental_skips_unchanged(self, tmp_path):
|
||||
"""增量模式: 未修改的文件跳过入库."""
|
||||
from src.core.config import EmbedConfig, ChunkConfig
|
||||
from src.core.db import VectorDB
|
||||
from src.core.embedder import create_embedder
|
||||
from src.core.ingest import DocumentIngestor
|
||||
from src.core.file_tracker import FileTracker
|
||||
|
||||
file = tmp_path / "stable.md"
|
||||
file.write_text("# 稳定文档\n\n内容不变。", encoding="utf-8")
|
||||
tracker_path = str(tmp_path / "tracker.json")
|
||||
db = VectorDB(persist_dir=str(tmp_path / "db"))
|
||||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||
ingestor = DocumentIngestor(
|
||||
db, embedder, "test_incr",
|
||||
chunk_config=ChunkConfig(max_size=1000, overlap=100),
|
||||
file_tracker=FileTracker(tracker_path),
|
||||
)
|
||||
count1 = ingestor.ingest_file(str(file), incremental=True)
|
||||
assert count1 > 0
|
||||
count2 = ingestor.ingest_file(str(file), incremental=True)
|
||||
assert count2 == 0 # 跳过
|
||||
|
||||
def test_ingest_file_incremental_reingests_modified(self, tmp_path):
|
||||
"""增量模式: 修改后的文件重新入库."""
|
||||
from src.core.config import EmbedConfig, ChunkConfig
|
||||
from src.core.db import VectorDB
|
||||
from src.core.embedder import create_embedder
|
||||
from src.core.ingest import DocumentIngestor
|
||||
from src.core.file_tracker import FileTracker
|
||||
|
||||
file = tmp_path / "changing.md"
|
||||
file.write_text("# v1\n\n初始版本的内容段落。", encoding="utf-8")
|
||||
tracker_path = str(tmp_path / "tracker2.json")
|
||||
db = VectorDB(persist_dir=str(tmp_path / "db2"))
|
||||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||
ingestor = DocumentIngestor(
|
||||
db, embedder, "test_incr2",
|
||||
chunk_config=ChunkConfig(max_size=1000, overlap=100),
|
||||
file_tracker=FileTracker(tracker_path),
|
||||
)
|
||||
count1 = ingestor.ingest_file(str(file), incremental=True)
|
||||
assert count1 > 0
|
||||
file.write_text("# v2\n\n新增段落,内容完全不同了。", encoding="utf-8")
|
||||
count2 = ingestor.ingest_file(str(file), incremental=True)
|
||||
assert count2 > 0
|
||||
|
||||
def test_ingest_file_force_mode_always_reingests(self, tmp_path):
|
||||
"""force=True 时始终重新入库(忽略 tracker)."""
|
||||
from src.core.config import EmbedConfig, ChunkConfig
|
||||
from src.core.db import VectorDB
|
||||
from src.core.embedder import create_embedder
|
||||
from src.core.ingest import DocumentIngestor
|
||||
from src.core.file_tracker import FileTracker
|
||||
|
||||
file = tmp_path / "force.md"
|
||||
file.write_text("# force test\n\n这是强制入库测试的内容。", encoding="utf-8")
|
||||
tracker_path = str(tmp_path / "tracker3.json")
|
||||
db = VectorDB(persist_dir=str(tmp_path / "db3"))
|
||||
embedder = create_embedder(EmbedConfig(mode="local"))
|
||||
ingestor = DocumentIngestor(
|
||||
db, embedder, "test_force",
|
||||
chunk_config=ChunkConfig(max_size=1000, overlap=100),
|
||||
file_tracker=FileTracker(tracker_path),
|
||||
)
|
||||
count1 = ingestor.ingest_file(str(file), incremental=True)
|
||||
count2 = ingestor.ingest_file(str(file), incremental=True, force=True)
|
||||
assert count1 > 0
|
||||
assert count2 > 0 # force 模式重新入库
|
||||
|
||||
Reference in New Issue
Block a user