diff --git a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py index c7684506d7..ad4da13fd2 100644 --- a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py @@ -130,6 +130,11 @@ def _write_index(index: faiss.Index, path: str) -> None: async def insert(self, vector: np.ndarray, id: int) -> None: """插入向量""" assert self.index is not None, "FAISS index is not initialized." + if not np.all(np.isfinite(vector)): + nan_count = int(np.sum(~np.isfinite(vector))) + raise RuntimeError( + f"向量包含 {nan_count} 个非有限值 (NaN/Inf)，无法写入 FAISS 索引。" + ) if vector.shape[0] != self.dimension: raise ValueError( f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}", @@ -144,6 +149,13 @@ async def insert_batch(self, vectors: np.ndarray, ids: list[int]) -> None: raise ValueError( f"向量必须是二维数组, 当前维度: {len(vectors.shape)}", ) + # Validate vector values before FAISS write; non-finite values cause + # C++ segfaults or silent index corruption. + if not np.all(np.isfinite(vectors)): + nan_count = int(np.sum(~np.isfinite(vectors))) + raise RuntimeError( + f"向量包含 {nan_count} 个非有限值 (NaN/Inf)，无法写入 FAISS 索引。请检查嵌入模型配置。" + ) if vectors.shape[1] != self.dimension: raise ValueError( f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}", @@ -158,6 +170,11 @@ async def search(self, vector: np.ndarray, k: int) -> tuple: """ assert self.index is not None, "FAISS index is not initialized." vector = np.asarray(vector, dtype=np.float32).ravel() + if not np.all(np.isfinite(vector)): + nan_count = int(np.sum(~np.isfinite(vector))) + raise RuntimeError( + f"查询向量包含 {nan_count} 个非有限值 (NaN/Inf)。" + ) if vector.shape[0] != self.dimension: raise ValueError( f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}", 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 13fde6c0ef..5a1c006812 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -1,4 +1,5 @@ import time +import traceback import uuid import numpy as np @@ -164,6 +165,19 @@ async def insert_batch( ), details={"vector_count": len(vectors)}, ) from exc + # Reject non-finite values before they can reach FAISS. Replacing them + # with zero would silently persist a corrupted embedding and degrade + # retrieval quality. + if not np.all(np.isfinite(vectors_array)): + nan_count = int(np.sum(~np.isfinite(vectors_array))) + raise KnowledgeBaseUploadError( + stage="embedding", + user_message=( + f"向量化失败：嵌入模型返回的向量包含 {nan_count} 个非有限值" + "（NaN/Inf），无法写入知识库。请检查嵌入模型配置。" + ), + details={"non_finite_values": nan_count}, + ) if vectors_array.ndim != 2: raise KnowledgeBaseUploadError( stage="embedding", @@ -198,7 +212,7 @@ async def insert_batch( raise KnowledgeBaseUploadError( stage="storage", user_message=( - f"存储失败：写入文档索引后返回的内部 ID 数量与文本分块数量不一致" + f"存储失败：写入文档索引后返回的内部 ID 数量不一致" f"（期望 {content_count}，实际 {len(int_ids)}）。" ), details={ @@ -207,9 +221,29 @@ async def insert_batch( }, ) await self.embedding_storage.insert_batch(vectors_array, int_ids) - except Exception: + except KnowledgeBaseUploadError: + # Roll back partial inserts before re-raising to prevent orphaned + # document rows when DocumentStorage returns mismatched ID count. await self._rollback_partial_insert(ids=ids, int_ids=int_ids) raise + except Exception as _faiss_err: + # Low-level FAISS errors (index corruption, resource exhaustion, + # I/O failure, etc.) bubble up here. Roll back partial writes and + # surface structured diagnostics so the upload caller can report + # the exact cause rather than the opaque "write index error". + await self._rollback_partial_insert(ids=ids, int_ids=int_ids) + raise KnowledgeBaseUploadError( + stage="storage", + user_message=( + "存储失败：嵌入向量已成功生成，但在写入 FAISS 向量索引时发生错误。" + "这可能是索引文件损坏、磁盘空间不足或嵌入维度变更导致的。" + "请查看日志中的详细原因。" + ), + details={ + "cause": str(_faiss_err), + "traceback": traceback.format_exc(), + }, + ) from _faiss_err return int_ids async def retrieve( diff --git a/tests/unit/test_faiss_vec_db.py b/tests/unit/test_faiss_vec_db.py index 8b30d2ca01..1683f526c6 100644 --- a/tests/unit/test_faiss_vec_db.py +++ b/tests/unit/test_faiss_vec_db.py @@ -1,6 +1,7 @@ import asyncio from unittest.mock import AsyncMock +import numpy as np import pytest from astrbot.core.db.vec_db.faiss_impl.embedding_storage import EmbeddingStorage @@ -138,6 +139,50 @@ async def test_insert_batch_rejects_embedding_content_count_mismatch() -> None: vec_db.document_storage.insert_documents_batch.assert_not_awaited() +@pytest.mark.asyncio +async def test_insert_batch_rejects_non_finite_embeddings_before_storage() -> None: + vec_db = FaissVecDB.__new__(FaissVecDB) + vec_db.embedding_provider = AsyncMock() + vec_db.embedding_provider.get_embeddings_batch.return_value = [ + [0.1, float("nan")], + [0.3, 0.4], + ] + vec_db.document_storage = AsyncMock() + vec_db.embedding_storage = AsyncMock() + vec_db.embedding_storage.dimension = 2 + + with pytest.raises(KnowledgeBaseUploadError, match="非有限值"): + await FaissVecDB.insert_batch( + vec_db, + contents=["chunk-1", "chunk-2"], + metadatas=[{}, {}], + ids=["doc-1", "doc-2"], + ) + + vec_db.document_storage.insert_documents_batch.assert_not_awaited() + vec_db.embedding_storage.insert_batch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_embedding_storage_rejects_non_finite_single_and_batch_vectors() -> None: + storage = EmbeddingStorage(2) + storage.save_index = AsyncMock() + + with pytest.raises(RuntimeError, match="非有限值"): + await storage.insert(np.array([float("inf"), 0.2], dtype=np.float32), 1) + + with pytest.raises(RuntimeError, match="非有限值"): + await storage.insert_batch( + np.array([[0.1, 0.2], [float("nan"), 0.4]], dtype=np.float32), + [1, 2], + ) + + with pytest.raises(RuntimeError, match="非有限值"): + await storage.search(np.array([0.1, float("nan")], dtype=np.float32), 1) + + storage.save_index.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"))