diff --git a/backend/src/kbqa/documents/repository.py b/backend/src/kbqa/documents/repository.py index c1add13..fc3e610 100644 --- a/backend/src/kbqa/documents/repository.py +++ b/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) diff --git a/backend/src/kbqa/documents/service.py b/backend/src/kbqa/documents/service.py index 5271e24..8edb626 100644 --- a/backend/src/kbqa/documents/service.py +++ b/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) diff --git a/backend/src/kbqa/indexing/worker.py b/backend/src/kbqa/indexing/worker.py index 05364ef..55e9774 100644 --- a/backend/src/kbqa/indexing/worker.py +++ b/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: diff --git a/backend/src/kbqa/lifecycle.py b/backend/src/kbqa/lifecycle.py index b162220..2247f2c 100644 --- a/backend/src/kbqa/lifecycle.py +++ b/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