feat: add durable document indexing pipeline

backend/src/kbqa/indexing/__init__.py

@@ -0,0 +1 @@
+"""Document indexing worker."""

backend/src/kbqa/indexing/loaders.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+from pypdf import PdfReader
+
+from kbqa.api.errors import AppError
+
+
+@dataclass(frozen=True, slots=True)
+class LoadedPage:
+    content: str
+    page_number: int | None
+
+
+def load_document(path: Path, file_type: str) -> list[LoadedPage]:
+    if file_type == "pdf":
+        with path.open("rb") as handle:
+            if handle.read(5) != b"%PDF-":
+                raise AppError("INVALID_PDF", "PDF 文件格式无效", 415)
+        reader = PdfReader(str(path))
+        pages = [
+            LoadedPage(content=(page.extract_text() or "").strip(), page_number=index + 1)
+            for index, page in enumerate(reader.pages)
+        ]
+        pages = [page for page in pages if page.content]
+        if not pages:
+            raise AppError(
+                "PDF_TEXT_NOT_EXTRACTABLE",
+                "PDF 中没有可提取文本，暂不支持扫描件 OCR",
+                415,
+            )
+        return pages
+    try:
+        content = path.read_text(encoding="utf-8").strip()
+    except UnicodeDecodeError as exc:
+        raise AppError("TEXT_ENCODING_ERROR", "文本文件必须使用 UTF-8 编码", 415) from exc
+    if not content:
+        raise AppError("EMPTY_DOCUMENT", "文档内容为空", 415)
+    return [LoadedPage(content=content, page_number=None)]

backend/src/kbqa/indexing/splitter.py

@@ -0,0 +1,34 @@
+from uuid import UUID, uuid5
+
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+from kbqa.indexing.loaders import LoadedPage
+from kbqa.rag.types import PreparedChunk
+
+
+def split_pages(
+    document_id: str,
+    pages: list[LoadedPage],
+    chunk_size: int,
+    chunk_overlap: int,
+) -> list[PreparedChunk]:
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=len,
+    )
+    chunks: list[PreparedChunk] = []
+    namespace = UUID(document_id)
+    for page in pages:
+        for content in splitter.split_text(page.content):
+            order_index = len(chunks)
+            chunks.append(
+                PreparedChunk(
+                    id=str(uuid5(namespace, str(order_index))),
+                    document_id=document_id,
+                    order_index=order_index,
+                    page_number=page.page_number,
+                    content=content,
+                )
+            )
+    return chunks

backend/src/kbqa/indexing/worker.py

@@ -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)