diff --git a/astrbot/core/db/vec_db/base.py b/astrbot/core/db/vec_db/base.py index 04f8903b15..c926f613b4 100644 --- a/astrbot/core/db/vec_db/base.py +++ b/astrbot/core/db/vec_db/base.py @@ -32,11 +32,13 @@ async def insert_batch( tasks_limit: int = 3, max_retries: int = 3, progress_callback=None, + embedding_contents: list[str] | None = None, ) -> int: """批量插入文本和其对应向量，自动生成 ID 并保持一致性。 Args: progress_callback: 进度回调函数，接收参数 (current, total) + embedding_contents: Optional enriched texts used only for embeddings. """ ... diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index c641dd56a4..13fde6c0ef 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -65,15 +65,19 @@ async def insert_batch( tasks_limit: int = 3, max_retries: int = 3, progress_callback=None, + embedding_contents: list[str] | None = None, ) -> list[int]: """批量插入文本和其对应向量，自动生成 ID 并保持一致性。 Args: progress_callback: 进度回调函数，接收参数 (current, total) + embedding_contents: Optional enriched texts used only for embeddings. """ metadatas = metadatas or [{} for _ in contents] ids = ids or [str(uuid.uuid4()) for _ in contents] + if embedding_contents is None: + embedding_contents = contents if not contents: logger.debug( @@ -106,11 +110,23 @@ async def insert_batch( "actual_ids": len(ids), }, ) + if len(embedding_contents) != content_count: + raise KnowledgeBaseUploadError( + stage="storage", + user_message=( + f"存储失败：文本分块数量与向量化文本数量不一致（期望 {content_count}，" + f"实际 {len(embedding_contents)}）。" + ), + details={ + "expected_contents": content_count, + "actual_embedding_contents": len(embedding_contents), + }, + ) start = time.time() logger.debug(f"Generating embeddings for {len(contents)} contents...") vectors = await self.embedding_provider.get_embeddings_batch( - contents, + embedding_contents, batch_size=batch_size, tasks_limit=tasks_limit, max_retries=max_retries, diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py index d3f43b085e..75234840cb 100644 --- a/astrbot/core/knowledge_base/kb_helper.py +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -323,16 +323,28 @@ async def upload_document( await progress_callback("chunking", 0, 100) try: - # 根据文件类型选择分块器：Markdown 文件使用结构感知分块 + # These parsers return Markdown, so retain their heading hierarchy. effective_chunker = self.chunker file_ext = Path(file_name).suffix.lower() if file_name else "" - if file_ext in (".md", ".markdown", ".mkd", ".mdx"): + if file_ext in { + ".adoc", + ".docx", + ".epub", + ".markdown", + ".md", + ".mdx", + ".mkd", + ".rst", + ".xls", + ".xlsx", + }: effective_chunker = MarkdownChunker( chunk_size=chunk_size, chunk_overlap=chunk_overlap, ) logger.info( - f"检测到 Markdown 文件 '{file_name}'，使用 MarkdownChunker 进行结构化分块" + f"Using MarkdownChunker for structured document " + f"'{file_name}'." ) chunks_text = await effective_chunker.chunk( @@ -380,6 +392,12 @@ async def upload_document( "chunk_index": idx, }, ) + document_title = Path(file_name).stem.strip() + embedding_contents = ( + [f"{document_title}\n\n{chunk_text}" for chunk_text in chunks_text] + if document_title + else contents + ) if progress_callback: await progress_callback("chunking", 100, 100) @@ -397,6 +415,7 @@ async def embedding_progress_callback(current, total) -> None: tasks_limit=tasks_limit, max_retries=max_retries, progress_callback=embedding_progress_callback, + embedding_contents=embedding_contents, ) except KnowledgeBaseUploadError: raise diff --git a/tests/unit/test_faiss_vec_db.py b/tests/unit/test_faiss_vec_db.py index 2d0d7f9dc6..8b30d2ca01 100644 --- a/tests/unit/test_faiss_vec_db.py +++ b/tests/unit/test_faiss_vec_db.py @@ -65,6 +65,79 @@ async def test_insert_batch_raises_friendly_error_for_embedding_count_mismatch() vec_db.embedding_storage.insert_batch.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("embedding_contents", "expected_embedding_contents"), + [ + (None, ["chunk one", "chunk two"]), + ( + ["guide\n\nchunk one", "guide\n\nchunk two"], + ["guide\n\nchunk one", "guide\n\nchunk two"], + ), + ], +) +async def test_insert_batch_uses_embedding_contents_without_changing_storage( + embedding_contents: list[str] | None, + expected_embedding_contents: list[str], +) -> None: + vec_db = FaissVecDB.__new__(FaissVecDB) + vec_db.embedding_provider = AsyncMock() + vec_db.embedding_provider.get_embeddings_batch.return_value = [ + [0.1, 0.2], + [0.3, 0.4], + ] + vec_db.document_storage = AsyncMock() + vec_db.document_storage.insert_documents_batch.return_value = [11, 12] + vec_db.embedding_storage = AsyncMock() + vec_db.embedding_storage.dimension = 2 + + await FaissVecDB.insert_batch( + vec_db, + contents=["chunk one", "chunk two"], + metadatas=[{}, {}], + ids=["doc-1", "doc-2"], + embedding_contents=embedding_contents, + ) + + vec_db.embedding_provider.get_embeddings_batch.assert_awaited_once_with( + expected_embedding_contents, + batch_size=32, + tasks_limit=3, + max_retries=3, + progress_callback=None, + ) + vec_db.document_storage.insert_documents_batch.assert_awaited_once_with( + ["doc-1", "doc-2"], + ["chunk one", "chunk two"], + [{}, {}], + ) + + +@pytest.mark.asyncio +async def test_insert_batch_rejects_embedding_content_count_mismatch() -> None: + vec_db = FaissVecDB.__new__(FaissVecDB) + vec_db.embedding_provider = AsyncMock() + vec_db.document_storage = AsyncMock() + vec_db.embedding_storage = AsyncMock() + + with pytest.raises(KnowledgeBaseUploadError) as exc_info: + await FaissVecDB.insert_batch( + vec_db, + contents=["chunk one", "chunk two"], + metadatas=[{}, {}], + ids=["doc-1", "doc-2"], + embedding_contents=["guide\n\nchunk one"], + ) + + assert exc_info.value.stage == "storage" + assert exc_info.value.details == { + "expected_contents": 2, + "actual_embedding_contents": 1, + } + vec_db.embedding_provider.get_embeddings_batch.assert_not_awaited() + vec_db.document_storage.insert_documents_batch.assert_not_awaited() + + def test_embedding_storage_rejects_zero_dimension_for_a_fresh_index(tmp_path) -> None: with pytest.raises(ValueError, match="无效的嵌入向量维度"): EmbeddingStorage(0, str(tmp_path / "index.faiss")) diff --git a/tests/unit/test_kb_upload_atomicity.py b/tests/unit/test_kb_upload_atomicity.py index e9570d1969..c7ad2d750b 100644 --- a/tests/unit/test_kb_upload_atomicity.py +++ b/tests/unit/test_kb_upload_atomicity.py @@ -349,6 +349,76 @@ async def fake_save_media(**kwargs): assert not media_file.exists() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("file_name", "file_type"), + [ + ("guide.docx", "docx"), + ("guide.xlsx", "xlsx"), + ("guide.xls", "xls"), + ("guide.rst", "rst"), + ("guide.adoc", "adoc"), + ("guide.epub", "epub"), + ], +) +async def test_upload_document_preserves_markdown_heading_paths( + file_name: str, + file_type: str, + tmp_path: Path, + stub_provider_manager_module, +) -> None: + """Structured parser output should retain parent headings in embeddings.""" + KBHelper = _import_kb_helper() + + helper = KBHelper.__new__(KBHelper) + helper.kb = KnowledgeBase( + kb_name="Test KB", + description="", + embedding_provider_id="emb", + ) + helper.kb_db = MagicMock() + helper.kb_db.get_db = _failing_get_db() + helper.vec_db = AsyncMock() + helper.vec_db.insert_batch.side_effect = RuntimeError("stop after chunking") + helper.vec_db.delete_documents = AsyncMock() + helper.kb_medias_dir = tmp_path / "medias" + helper.chunker = AsyncMock() + + parse_result = MagicMock( + text=( + "# Handbook\n\nOverview.\n\n" + "## Installation\n\nInstall the package.\n\n" + "### Linux\n\nRun the command." + ), + media=[], + ) + + with ( + patch( + "astrbot.core.knowledge_base.kb_helper.select_parser", + new=AsyncMock( + return_value=MagicMock(parse=AsyncMock(return_value=parse_result)), + ), + ), + patch.object(helper, "_ensure_vec_db", new=AsyncMock()), + pytest.raises(KnowledgeBaseUploadError) as exc_info, + ): + await helper.upload_document( + file_name=file_name, + file_content=b"structured document", + file_type=file_type, + ) + + assert exc_info.value.stage == "storage" + helper.chunker.chunk.assert_not_awaited() + contents = helper.vec_db.insert_batch.await_args.kwargs["contents"] + embedding_contents = helper.vec_db.insert_batch.await_args.kwargs[ + "embedding_contents" + ] + assert "Handbook > Installation\n\n### Linux" in contents[2] + assert embedding_contents[2] == f"guide\n\n{contents[2]}" + + @pytest.mark.asyncio async def test_upload_document_cleans_up_on_metadata_failure( stub_provider_manager_module, @@ -384,6 +454,14 @@ async def test_upload_document_cleans_up_on_metadata_failure( assert exc_info.value.stage == "metadata" helper.vec_db.insert_batch.assert_awaited_once() + assert helper.vec_db.insert_batch.await_args.kwargs["contents"] == [ + "chunk a", + "chunk b", + ] + assert helper.vec_db.insert_batch.await_args.kwargs["embedding_contents"] == [ + "demo\n\nchunk a", + "demo\n\nchunk b", + ] helper.vec_db.delete_documents.assert_awaited()