perf: enrich knowledge base embeddings with document context

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.
 
         """
         ...

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,

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

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"))

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()