perf: enrich knowledge base embeddings with document context

[lxfight]

Short knowledge base chunks often lose the document and section context needed to match user queries. This PR enriches embedding inputs with the document title and preserves heading paths for formats that are parsed into Markdown, while keeping vector storage and returned chunk behavior compatible.

Modifications / 改动点

• Allow vector insertion to use enriched embedding text separately from stored chunk text.
• Prefix chunk embedding inputs with the source document title.
• Apply the existing structure-aware Markdown chunker to DOCX, XLS/XLSX, RST, AsciiDoc, and EPUB parser output.
• Preserve parent heading paths such as Handbook > Installation for nested sections.
• Keep plain-text fallback behavior when parsed output contains no headings.
• Do not add adjacent passage text: the held-out experiment showed a substantial regression.
• This is NOT a breaking change. / 这不是一个破坏性变更。

Screenshots or Test Results / 运行截图或测试结果

Document title experiment

• Public dataset: MIRACL Chinese development split
• Corpus revision: d921ec7e349ce0d28daf30b2da9da5ee698bef0d
• Fixed split seed: astrbot-kb-context-v1
• Corpus: 50,000 passages
• Queries: 196 development / 197 held-out test
• Embedding model: baai/bge-m3

Held-out fusion metrics:

Embedding input	Recall@5	HitRate@5	nDCG@5
Passage body	0.9419	0.9898	0.9284
Document title + passage body	0.9594	1.0000	0.9600

Heading path and adjacent context experiment

To recover real heading paths, MIRACL passages were joined with public Chinese Wikipedia revisions. The candidate set contained all 22,230 passages from the 783 relevant articles, with 393 queries and 994 positive labels. This experiment is a separate corpus and its absolute scores are not directly comparable to the title experiment.

Held-out fusion metrics:

Embedding input	Recall@5	HitRate@5	nDCG@5
Document title + passage	0.7320	0.9086	0.7445
Document title + heading path + passage	0.7489	0.9239	0.7623
Heading path + previous/next 128 characters	0.4548	0.7107	0.4569

The heading-only variant was selected on the development split and improved every held-out top-five metric. Adjacent context diluted the target passage and was intentionally excluded. Benchmark datasets and generated results remain local and are not included in this PR.

Verification steps:

uv run pytest -q tests/unit/test_faiss_vec_db.py tests/unit/test_kb_manager_resilience.py tests/unit/test_kb_upload_atomicity.py tests/unit/test_kb_document_cleanup.py tests/unit/test_knowledge_base_service_contract.py tests/test_epub_parser.py
uv run ruff format --check .
uv run ruff check .

Results:

• 53 tests passed.
• All 483 Python files passed Ruff formatting checks.
• Ruff lint checks passed.

---

Checklist / 检查清单

• 😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
  / 如果 PR 中有新加入的功能，已经通过 Issue / 邮件等方式和作者讨论过。
• 👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
  / 我的更改经过了良好的测试，并已在上方提供了“验证步骤”和“运行截图”。
• 🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in requirements.txt and pyproject.toml.
  / 我确保没有引入新依赖库，或者引入了新依赖库的同时将其添加到 requirements.txt 和 pyproject.toml 文件相应位置。
• 😮 My changes do not introduce malicious code.
  / 我的更改没有引入恶意代码。

Summary by Sourcery

Enrich knowledge base embeddings with document and section context while keeping stored chunks unchanged and maintaining upload atomicity guarantees.

New Features:

• Support passing separate enriched embedding texts alongside stored chunk contents when inserting batches into the vector database.
• Prefix chunk embedding inputs with the source document title to improve retrieval quality.
• Apply the structure-aware Markdown chunker to additional structured formats including DOCX, XLS/XLSX, RST, AsciiDoc, and EPUB to preserve heading hierarchies in chunks.

Bug Fixes:

• Ensure heading paths from structured parser output are preserved in uploaded document chunks and their corresponding embeddings.
• Validate that the number of embedding texts matches the number of content chunks, raising a user-friendly error on mismatch to avoid inconsistent storage.

Tests:

• Add unit tests verifying that upload_document uses MarkdownChunker for structured formats and preserves heading paths and document titles in embedding contents.
• Add a unit test ensuring FaissVecDB embeds using enriched texts but stores the original chunk contents unchanged.

[dosubot]

📄 Knowledge review

Dosu skipped reviewing this PR because your organization has used its 200 included credits for the month. Your usage will reset on 2026-08-01. To have Dosu review this PR before then, ask your organization admin to upgrade to a pro account.

---

[sourcery-ai]

Hey - I've found 2 issues, and left some high level feedback:

• In FaissVecDB.insert_batch, the debug log still reports len(contents) while embeddings are now generated from embedding_contents; consider logging the actual embedding input count to avoid confusion when these lists differ.
• The new embedding_contents parameter was added to the VecDB base interface; double-check that all concrete VecDB implementations accept and correctly handle this argument to prevent runtime errors or silently ignored enrichment.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In `FaissVecDB.insert_batch`, the debug log still reports `len(contents)` while embeddings are now generated from `embedding_contents`; consider logging the actual embedding input count to avoid confusion when these lists differ.
- The new `embedding_contents` parameter was added to the `VecDB` base interface; double-check that all concrete `VecDB` implementations accept and correctly handle this argument to prevent runtime errors or silently ignored enrichment.

## Individual Comments

### Comment 1
<location path="astrbot/core/knowledge_base/kb_helper.py" line_range="395" />
<code_context>
                         "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]
</code_context>
<issue_to_address>
**issue (bug_risk):** Guard against `file_name` being `None` before constructing a `Path`.

Some call paths may pass `file_name=None` or omit it, and `Path(None)` will raise a `TypeError`. Please add a guard (e.g. `document_title = Path(file_name).stem.strip() if file_name else ""`) or ensure the title is derived earlier when `file_name` is guaranteed to be set.
</issue_to_address>

### Comment 2
<location path="tests/unit/test_faiss_vec_db.py" line_range="68-77" />
<code_context>
     vec_db.embedding_storage.insert_batch.assert_not_awaited()


+@pytest.mark.asyncio
+async def test_insert_batch_embeds_context_but_stores_original_contents() -> 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=["guide\n\nchunk one", "guide\n\nchunk two"],
+    )
+
+    vec_db.embedding_provider.get_embeddings_batch.assert_awaited_once_with(
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding tests for the default `embedding_contents=None` path and the new mismatch error condition.

To fully cover the new `insert_batch` behavior, please add:

1. A test where `embedding_contents` is omitted or `None`, verifying that `get_embeddings_batch` is called with `contents` and `insert_documents_batch` receives the original `contents`.
2. A test where `embedding_contents` and `contents` differ in length, verifying that a `KnowledgeBaseUploadError` is raised with the expected `stage` and `details`.

These will explicitly exercise both the default path and the new safety check and help prevent subtle caller bugs.

Suggested implementation:

```python
        ["doc-1", "doc-2"],
        ["chunk one", "chunk two"],
        [{}, {}],
    )

    @pytest.mark.asyncio
    async def test_insert_batch_uses_contents_when_embedding_contents_not_provided() -> None:
        vec_db = FaissVecDB.__new__(FaissVecDB)
        vec_db.embedding_provider = AsyncMock()
        vec_db.embedding_provider.get_embeddings_batch.return_value = [
            [0.5, 0.6],
            [0.7, 0.8],
        ]
        vec_db.document_storage = AsyncMock()
        vec_db.document_storage.insert_documents_batch.return_value = [21, 22]
        vec_db.embedding_storage = AsyncMock()
        vec_db.embedding_storage.dimension = 2

        await FaissVecDB.insert_batch(
            vec_db,
            contents=["default one", "default two"],
            metadatas=[{"m": 1}, {"m": 2}],
            ids=["doc-3", "doc-4"],
            # omit embedding_contents to take the default path
        )

        vec_db.embedding_provider.get_embeddings_batch.assert_awaited_once_with(
            ["default one", "default two"],
            batch_size=32,
            tasks_limit=3,
            max_retries=3,
            progress_callback=None,
        )
        vec_db.document_storage.insert_documents_batch.assert_awaited_once_with(
            ["doc-3", "doc-4"],
            ["default one", "default two"],
            [{"m": 1}, {"m": 2}],
        )

    @pytest.mark.asyncio
    async def test_insert_batch_raises_on_embedding_contents_length_mismatch() -> None:
        vec_db = FaissVecDB.__new__(FaissVecDB)
        vec_db.embedding_provider = AsyncMock()
        vec_db.document_storage = AsyncMock()
        vec_db.embedding_storage = AsyncMock()
        vec_db.embedding_storage.dimension = 2

        with pytest.raises(KnowledgeBaseUploadError) as excinfo:
            await FaissVecDB.insert_batch(
                vec_db,
                contents=["chunk one", "chunk two"],
                metadatas=[{}, {}],
                ids=["doc-1", "doc-2"],
                embedding_contents=["guide\n\nchunk one"],
            )

        # Ensure we fail fast with a clear stage and details
        assert excinfo.value.stage == "embedding"
        assert "length mismatch" in excinfo.value.details

```

1. Ensure `KnowledgeBaseUploadError` is imported at the top of `tests/unit/test_faiss_vec_db.py` (e.g., `from <module> import KnowledgeBaseUploadError`) consistent with the rest of the codebase.
2. Confirm that `FaissVecDB.insert_batch`:
   - Defaults `embedding_contents` to `contents` when `embedding_contents` is `None` or omitted.
   - Raises `KnowledgeBaseUploadError` on length mismatch with `stage="embedding"` and a `details` string that includes "length mismatch" (or adjust the assertion to match the actual message).
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[lxfight]

Reviewed the remaining high-level suggestions:

• The debug count cannot diverge from the embedding input count: embedding_contents length is validated against contents immediately before logging and embedding generation.
• FaissVecDB is the only concrete BaseVecDB implementation in the repository, and its signature handles the new optional argument.

The actionable test-coverage suggestion is addressed in 97b026b.