feat: add durable document indexing pipeline

2026-07-13 11:32:14 +08:00
parent eebcd5d6a7
commit 303116f950
12 changed files with 598 additions and 1 deletions
@@ -0,0 +1,125 @@
+import asyncio
+import logging
+from pathlib import Path
+
+from sqlalchemy.ext.asyncio import async_sessionmaker
+
+from kbqa.api.errors import AppError
+from kbqa.config import Settings
+from kbqa.documents.models import Chunk
+from kbqa.documents.repository import DocumentRepository
+from kbqa.indexing.loaders import load_document
+from kbqa.indexing.splitter import split_pages
+from kbqa.rag.embeddings import EmbeddingProvider
+from kbqa.rag.types import PreparedChunk
+from kbqa.rag.vector_store import MilvusVectorStore
+
+logger = logging.getLogger(__name__)
+
+
+class IndexWorker:
+    def __init__(
+        self,
+        settings: Settings,
+        session_factory: async_sessionmaker,
+        embeddings: EmbeddingProvider,
+        vector_store: MilvusVectorStore,
+    ) -> None:
+        self.settings = settings
+        self.session_factory = session_factory
+        self.embeddings = embeddings
+        self.vector_store = vector_store
+        self.queue: asyncio.Queue[str] = asyncio.Queue()
+        self._task: asyncio.Task[None] | None = None
+
+    async def start(self) -> None:
+        async with self.session_factory() as session:
+            pending_ids = await DocumentRepository(session).requeue_unfinished()
+        for job_id in pending_ids:
+            self.queue.put_nowait(job_id)
+        self._task = asyncio.create_task(self._run(), name="kbqa-index-worker")
+
+    async def stop(self) -> None:
+        if self._task is None:
+            return
+        self._task.cancel()
+        try:
+            await self._task
+        except asyncio.CancelledError:
+            pass
+        self._task = None
+
+    def enqueue(self, job_id: str) -> None:
+        self.queue.put_nowait(job_id)
+
+    async def _run(self) -> None:
+        while True:
+            job_id = await self.queue.get()
+            try:
+                await self._process(job_id)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.exception("Unexpected indexing worker failure job_id=%s", job_id)
+            finally:
+                self.queue.task_done()
+
+    async def _process(self, job_id: str) -> None:
+        async with self.session_factory() as session:
+            repository = DocumentRepository(session)
+            job = await repository.claim_job(job_id)
+            if job is None:
+                return
+            document = await repository.get(job.document_id)
+            if document is None:
+                return
+            document_id = document.id
+            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.session_factory() as session:
+                await DocumentRepository(session).replace_chunks(document_id, [])
+            pages = await asyncio.to_thread(load_document, Path(stored_path), file_type)
+            prepared = split_pages(
+                document_id,
+                pages,
+                self.settings.chunk_size,
+                self.settings.chunk_overlap,
+            )
+            if not prepared:
+                raise AppError("EMPTY_DOCUMENT", "文档分块结果为空", 415)
+            orm_chunks = [
+                Chunk(
+                    id=chunk.id,
+                    document_id=chunk.document_id,
+                    order_index=chunk.order_index,
+                    content=chunk.content,
+                    page_number=chunk.page_number,
+                    char_count=len(chunk.content),
+                )
+                for chunk in prepared
+            ]
+            async with self.session_factory() as session:
+                await DocumentRepository(session).replace_chunks(document_id, orm_chunks)
+            await self._embed_and_upsert(prepared)
+            async with self.session_factory() as session:
+                await DocumentRepository(session).mark_indexed(document_id, len(prepared))
+        except asyncio.CancelledError:
+            raise
+        except AppError as exc:
+            await self._mark_failed(document_id, exc.code, exc.message)
+        except Exception:
+            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:
+        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 def _mark_failed(self, document_id: str, code: str, message: str) -> None:
+        async with self.session_factory() as session:
+            await DocumentRepository(session).mark_failed(document_id, code, message)