fix: serialize vector indexing and deletion

backend/src/kbqa/documents/repository.py

@@ -60,6 +60,12 @@ class DocumentRepository:
     async def get(self, document_id: str) -> Document | None:
         return await self.session.get(Document, document_id)
 
+    async def is_indexing(self, document_id: str) -> bool:
+        status = await self.session.scalar(
+            select(Document.status).where(Document.id == document_id)
+        )
+        return status == "indexing"
+
     async def get_job_for_document(self, document_id: str) -> IndexJob | None:
         return await self.session.scalar(
             select(IndexJob).where(IndexJob.document_id == document_id)

backend/src/kbqa/documents/service.py

@@ -35,11 +35,13 @@ class DocumentService:
         session_factory: async_sessionmaker,
         worker: IndexWorker,
         vector_store: MilvusVectorStore,
+        vector_mutation_lock: asyncio.Lock,
     ) -> None:
         self.settings = settings
         self.session_factory = session_factory
         self.worker = worker
         self.vector_store = vector_store
+        self.vector_mutation_lock = vector_mutation_lock
 
     async def upload(self, file: UploadFile) -> Document:
         original_name = Path(file.filename or "").name
@@ -132,7 +134,8 @@ class DocumentService:
                 continue
 
     async def _finish_delete(self, document: Document) -> None:
-        await asyncio.to_thread(self.vector_store.delete_document, document.id)
-        await asyncio.to_thread(Path(document.stored_path).unlink, True)
-        async with self.session_factory() as session:
-            await DocumentRepository(session).delete_document(document.id)
+        async with self.vector_mutation_lock:
+            await asyncio.to_thread(self.vector_store.delete_document, document.id)
+            await asyncio.to_thread(Path(document.stored_path).unlink, True)
+            async with self.session_factory() as session:
+                await DocumentRepository(session).delete_document(document.id)

backend/src/kbqa/indexing/worker.py

@@ -24,11 +24,13 @@ class IndexWorker:
         session_factory: async_sessionmaker,
         embeddings: EmbeddingProvider,
         vector_store: MilvusVectorStore,
+        vector_mutation_lock: asyncio.Lock,
     ) -> None:
         self.settings = settings
         self.session_factory = session_factory
         self.embeddings = embeddings
         self.vector_store = vector_store
+        self.vector_mutation_lock = vector_mutation_lock
         self.queue: asyncio.Queue[str] = asyncio.Queue()
         self._task: asyncio.Task[None] | None = None
 
@@ -77,7 +79,8 @@ class IndexWorker:
             stored_path = document.stored_path
             file_type = document.file_type
         try:
-            await asyncio.to_thread(self.vector_store.delete_document, document_id)
+            async with self.vector_mutation_lock:
+                await asyncio.to_thread(self.vector_store.delete_document, document_id)
             async with self.session_factory() as session:
                 await DocumentRepository(session).replace_chunks(document_id, [])
             pages = await asyncio.to_thread(load_document, Path(stored_path), file_type)
@@ -102,7 +105,9 @@ class IndexWorker:
             ]
             async with self.session_factory() as session:
                 await DocumentRepository(session).replace_chunks(document_id, orm_chunks)
-            await self._embed_and_upsert(prepared)
+            completed = await self._embed_and_upsert(document_id, prepared)
+            if not completed:
+                return
             async with self.session_factory() as session:
                 await DocumentRepository(session).mark_indexed(document_id, len(prepared))
         except asyncio.CancelledError:
@@ -113,12 +118,17 @@ class IndexWorker:
             logger.exception("Document indexing failed document_id=%s", document_id)
             await self._mark_failed(document_id, "INDEXING_FAILED", "文档索引失败，请重试")
 
-    async def _embed_and_upsert(self, chunks: list[PreparedChunk]) -> None:
+    async def _embed_and_upsert(self, document_id: str, chunks: list[PreparedChunk]) -> bool:
         batch_size = 20
         for start in range(0, len(chunks), batch_size):
             batch = chunks[start : start + batch_size]
             vectors = await self.embeddings.embed_documents([chunk.content for chunk in batch])
-            await asyncio.to_thread(self.vector_store.upsert, batch, vectors)
+            async with self.vector_mutation_lock:
+                async with self.session_factory() as session:
+                    if not await DocumentRepository(session).is_indexing(document_id):
+                        return False
+                await asyncio.to_thread(self.vector_store.upsert, batch, vectors)
+        return True
 
     async def _mark_failed(self, document_id: str, code: str, message: str) -> None:
         async with self.session_factory() as session:

backend/src/kbqa/lifecycle.py

@@ -38,12 +38,20 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     await asyncio.to_thread(vector_store.ensure_collection)
     embeddings = EmbeddingProvider(settings)
     retriever = KnowledgeRetriever(embeddings, vector_store, database.session_factory)
-    worker = IndexWorker(settings, database.session_factory, embeddings, vector_store)
+    vector_mutation_lock = asyncio.Lock()
+    worker = IndexWorker(
+        settings,
+        database.session_factory,
+        embeddings,
+        vector_store,
+        vector_mutation_lock,
+    )
     document_service = DocumentService(
         settings,
         database.session_factory,
         worker,
         vector_store,
+        vector_mutation_lock,
     )
     app.state.settings = settings
     app.state.database = database