From 303116f9509c428d79009fcb1798ab4adcaba115 Mon Sep 17 00:00:00 2001 From: gqt <3217233537@qq.com> Date: Mon, 13 Jul 2026 11:32:14 +0800 Subject: [PATCH] feat: add durable document indexing pipeline --- backend/src/kbqa/documents/service.py | 130 ++++++++++++++++++++++++++ backend/src/kbqa/health/routes.py | 2 +- backend/src/kbqa/indexing/__init__.py | 1 + backend/src/kbqa/indexing/loaders.py | 39 ++++++++ backend/src/kbqa/indexing/splitter.py | 34 +++++++ backend/src/kbqa/indexing/worker.py | 125 +++++++++++++++++++++++++ backend/src/kbqa/lifecycle.py | 25 +++++ backend/src/kbqa/rag/__init__.py | 1 + backend/src/kbqa/rag/embeddings.py | 63 +++++++++++++ backend/src/kbqa/rag/retriever.py | 56 +++++++++++ backend/src/kbqa/rag/types.py | 29 ++++++ backend/src/kbqa/rag/vector_store.py | 94 +++++++++++++++++++ 12 files changed, 598 insertions(+), 1 deletion(-) create mode 100644 backend/src/kbqa/documents/service.py create mode 100644 backend/src/kbqa/indexing/__init__.py create mode 100644 backend/src/kbqa/indexing/loaders.py create mode 100644 backend/src/kbqa/indexing/splitter.py create mode 100644 backend/src/kbqa/indexing/worker.py create mode 100644 backend/src/kbqa/rag/__init__.py create mode 100644 backend/src/kbqa/rag/embeddings.py create mode 100644 backend/src/kbqa/rag/retriever.py create mode 100644 backend/src/kbqa/rag/types.py create mode 100644 backend/src/kbqa/rag/vector_store.py diff --git a/backend/src/kbqa/documents/service.py b/backend/src/kbqa/documents/service.py new file mode 100644 index 0000000..90d5f8e --- /dev/null +++ b/backend/src/kbqa/documents/service.py @@ -0,0 +1,130 @@ +import asyncio +import hashlib +import os +from pathlib import Path +from uuid import uuid4 + +from fastapi import UploadFile +from sqlalchemy.ext.asyncio import async_sessionmaker + +from kbqa.api.errors import AppError +from kbqa.config import Settings +from kbqa.documents.models import Document +from kbqa.documents.repository import DocumentRepository, StoredUpload +from kbqa.indexing.worker import IndexWorker +from kbqa.rag.vector_store import MilvusVectorStore + +ALLOWED_EXTENSIONS = { + ".pdf": ("pdf", "application/pdf"), + ".md": ("md", "text/markdown"), + ".markdown": ("md", "text/markdown"), + ".txt": ("txt", "text/plain"), +} + + +class DocumentService: + def __init__( + self, + settings: Settings, + session_factory: async_sessionmaker, + worker: IndexWorker, + vector_store: MilvusVectorStore, + ) -> None: + self.settings = settings + self.session_factory = session_factory + self.worker = worker + self.vector_store = vector_store + + async def upload(self, file: UploadFile) -> Document: + original_name = Path(file.filename or "").name + suffix = Path(original_name).suffix.lower() + if suffix not in ALLOWED_EXTENSIONS: + raise AppError("UNSUPPORTED_FILE_TYPE", "仅支持 PDF、Markdown 和 TXT 文件", 415) + file_type, expected_mime = ALLOWED_EXTENSIONS[suffix] + document_id = str(uuid4()) + target = self.settings.raw_data_dir / f"{document_id}{suffix}" + temporary = target.with_suffix(f"{suffix}.part") + maximum = self.settings.max_upload_mib * 1024 * 1024 + digest = hashlib.sha256() + size = 0 + prefix = b"" + try: + with temporary.open("wb") as handle: + while chunk := await file.read(1024 * 1024): + size += len(chunk) + if size > maximum: + raise AppError("FILE_TOO_LARGE", "上传文件超过大小限制", 413) + if len(prefix) < 5: + prefix = (prefix + chunk)[:5] + digest.update(chunk) + handle.write(chunk) + if size == 0: + raise AppError("EMPTY_DOCUMENT", "上传文件为空", 415) + self._validate_content(file_type, prefix, temporary) + os.replace(temporary, target) + upload = StoredUpload( + document_id=document_id, + filename=original_name, + stored_path=str(target), + file_type=file_type, + mime_type=expected_mime, + file_size=size, + sha256=digest.hexdigest(), + ) + async with self.session_factory() as session: + document, job = await DocumentRepository(session).create_pending(upload) + self.worker.enqueue(job.id) + return document + except Exception: + temporary.unlink(missing_ok=True) + if target.exists(): + target.unlink(missing_ok=True) + raise + finally: + await file.close() + + @staticmethod + def _validate_content(file_type: str, prefix: bytes, path: Path) -> None: + if file_type == "pdf": + if prefix != b"%PDF-": + raise AppError("INVALID_PDF", "PDF 文件格式无效", 415) + return + try: + path.read_text(encoding="utf-8") + except UnicodeDecodeError as exc: + raise AppError("TEXT_ENCODING_ERROR", "文本文件必须使用 UTF-8 编码", 415) from exc + + async def retry(self, document_id: str) -> Document: + async with self.session_factory() as session: + repository = DocumentRepository(session) + job = await repository.requeue_failed(document_id) + document = await repository.get(document_id) + if document is None: + raise AppError("DOCUMENT_NOT_FOUND", "文档不存在", 404) + if job is None: + raise AppError("DOCUMENT_NOT_RETRYABLE", "当前文档状态不能重试", 409) + self.worker.enqueue(job.id) + return document + + async def delete(self, document_id: str) -> None: + async with self.session_factory() as session: + repository = DocumentRepository(session) + document = await repository.mark_deleting(document_id) + if document is None: + raise AppError("DOCUMENT_NOT_FOUND", "文档不存在", 404) + await self._finish_delete(document) + + async def recover_deletions(self) -> None: + async with self.session_factory() as session: + documents = await DocumentRepository(session).deleting_documents() + for document in documents: + try: + await self._finish_delete(document) + except Exception: + 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) diff --git a/backend/src/kbqa/health/routes.py b/backend/src/kbqa/health/routes.py index 4a7d434..19ceff9 100644 --- a/backend/src/kbqa/health/routes.py +++ b/backend/src/kbqa/health/routes.py @@ -21,7 +21,7 @@ async def health(request: Request) -> dict[str, object]: components["sqlite"] = {"status": "error"} try: - await asyncio.to_thread(request.app.state.milvus_client.list_collections) + await asyncio.to_thread(request.app.state.vector_store.ping) except Exception: components["milvus"] = {"status": "error"} diff --git a/backend/src/kbqa/indexing/__init__.py b/backend/src/kbqa/indexing/__init__.py new file mode 100644 index 0000000..6a060e0 --- /dev/null +++ b/backend/src/kbqa/indexing/__init__.py @@ -0,0 +1 @@ +"""Document indexing worker.""" diff --git a/backend/src/kbqa/indexing/loaders.py b/backend/src/kbqa/indexing/loaders.py new file mode 100644 index 0000000..7476bc7 --- /dev/null +++ b/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)] diff --git a/backend/src/kbqa/indexing/splitter.py b/backend/src/kbqa/indexing/splitter.py new file mode 100644 index 0000000..03ad0a0 --- /dev/null +++ b/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 diff --git a/backend/src/kbqa/indexing/worker.py b/backend/src/kbqa/indexing/worker.py new file mode 100644 index 0000000..05364ef --- /dev/null +++ b/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) diff --git a/backend/src/kbqa/lifecycle.py b/backend/src/kbqa/lifecycle.py index 1a9c2da..b162220 100644 --- a/backend/src/kbqa/lifecycle.py +++ b/backend/src/kbqa/lifecycle.py @@ -1,12 +1,18 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from pathlib import Path +import asyncio from fastapi import FastAPI from pymilvus import MilvusClient from kbqa.config import get_settings from kbqa.database import Database +from kbqa.documents.service import DocumentService +from kbqa.indexing.worker import IndexWorker +from kbqa.rag.embeddings import EmbeddingProvider +from kbqa.rag.retriever import KnowledgeRetriever +from kbqa.rag.vector_store import MilvusVectorStore def _ensure_runtime_directories(database_url: str, milvus_uri: str, raw_data_dir: Path) -> None: @@ -28,11 +34,30 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: ) database = Database(settings.database_url) milvus_client = MilvusClient(uri=settings.milvus_uri) + vector_store = MilvusVectorStore(milvus_client, settings.embedding_dim) + 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) + document_service = DocumentService( + settings, + database.session_factory, + worker, + vector_store, + ) app.state.settings = settings app.state.database = database app.state.milvus_client = milvus_client + app.state.vector_store = vector_store + app.state.embeddings = embeddings + app.state.retriever = retriever + app.state.index_worker = worker + app.state.document_service = document_service + await worker.start() + await document_service.recover_deletions() try: yield finally: + await worker.stop() milvus_client.close() await database.dispose() diff --git a/backend/src/kbqa/rag/__init__.py b/backend/src/kbqa/rag/__init__.py new file mode 100644 index 0000000..bca0d5f --- /dev/null +++ b/backend/src/kbqa/rag/__init__.py @@ -0,0 +1 @@ +"""Retrieval-augmented generation infrastructure.""" diff --git a/backend/src/kbqa/rag/embeddings.py b/backend/src/kbqa/rag/embeddings.py new file mode 100644 index 0000000..7534c72 --- /dev/null +++ b/backend/src/kbqa/rag/embeddings.py @@ -0,0 +1,63 @@ +from langchain_openai import OpenAIEmbeddings + +from kbqa.api.errors import AppError +from kbqa.config import Settings, validate_live_settings + + +class EmbeddingProvider: + def __init__(self, settings: Settings) -> None: + self.settings = settings + self._client: OpenAIEmbeddings | None = None + + def _get_client(self) -> OpenAIEmbeddings: + if self._client is None: + validate_live_settings(self.settings) + self._client = OpenAIEmbeddings( + api_key=self.settings.dashscope_api_key, + base_url=self.settings.dashscope_base_url, + model=self.settings.embedding_model, + dimensions=self.settings.embedding_dim, + chunk_size=20, + max_retries=2, + timeout=60.0, + check_embedding_ctx_length=False, + ) + return self._client + + def _validate_vectors(self, vectors: list[list[float]]) -> list[list[float]]: + if any(len(vector) != self.settings.embedding_dim for vector in vectors): + raise AppError( + "EMBEDDING_DIMENSION_MISMATCH", + "Embedding 返回的向量维度不正确", + 502, + retryable=True, + ) + return vectors + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + try: + vectors = await self._get_client().aembed_documents(texts) + except AppError: + raise + except Exception as exc: + raise AppError( + "EMBEDDING_UPSTREAM_ERROR", + "文档向量化失败，请稍后重试", + 502, + retryable=True, + ) from exc + return self._validate_vectors(vectors) + + async def embed_query(self, text: str) -> list[float]: + try: + vector = await self._get_client().aembed_query(text) + except AppError: + raise + except Exception as exc: + raise AppError( + "EMBEDDING_UPSTREAM_ERROR", + "查询向量化失败，请稍后重试", + 502, + retryable=True, + ) from exc + return self._validate_vectors([vector])[0] diff --git a/backend/src/kbqa/rag/retriever.py b/backend/src/kbqa/rag/retriever.py new file mode 100644 index 0000000..a30470f --- /dev/null +++ b/backend/src/kbqa/rag/retriever.py @@ -0,0 +1,56 @@ +import asyncio + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import async_sessionmaker + +from kbqa.documents.models import Chunk, Document +from kbqa.rag.embeddings import EmbeddingProvider +from kbqa.rag.types import SourceCandidate +from kbqa.rag.vector_store import MilvusVectorStore + + +class KnowledgeRetriever: + def __init__( + self, + embeddings: EmbeddingProvider, + vector_store: MilvusVectorStore, + session_factory: async_sessionmaker, + ) -> None: + self.embeddings = embeddings + self.vector_store = vector_store + self.session_factory = session_factory + + async def retrieve(self, query: str, top_k: int) -> list[SourceCandidate]: + vector = await self.embeddings.embed_query(query) + candidate_limit = max(top_k * 3, top_k) + hits = await asyncio.to_thread(self.vector_store.search, vector, candidate_limit) + if not hits: + return [] + ids = [hit.chunk_id for hit in hits] + async with self.session_factory() as session: + rows = await session.execute( + select(Chunk, Document) + .join(Document, Document.id == Chunk.document_id) + .where(Chunk.id.in_(ids), Document.status == "indexed") + ) + lookup = {chunk.id: (chunk, document) for chunk, document in rows} + sources: list[SourceCandidate] = [] + for hit in hits: + record = lookup.get(hit.chunk_id) + if record is None: + continue + chunk, document = record + sources.append( + SourceCandidate( + chunk_id=chunk.id, + document_id=document.id, + filename=document.filename, + order_index=chunk.order_index, + page_number=chunk.page_number, + content=chunk.content, + score=hit.score, + ) + ) + if len(sources) == top_k: + break + return sources diff --git a/backend/src/kbqa/rag/types.py b/backend/src/kbqa/rag/types.py new file mode 100644 index 0000000..746a028 --- /dev/null +++ b/backend/src/kbqa/rag/types.py @@ -0,0 +1,29 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class SourceCandidate: + chunk_id: str + document_id: str + filename: str + order_index: int + page_number: int | None + content: str + score: float + + +@dataclass(frozen=True, slots=True) +class PreparedChunk: + id: str + document_id: str + order_index: int + page_number: int | None + content: str + + +@dataclass(frozen=True, slots=True) +class VectorHit: + chunk_id: str + document_id: str + order_index: int + score: float diff --git a/backend/src/kbqa/rag/vector_store.py b/backend/src/kbqa/rag/vector_store.py new file mode 100644 index 0000000..e497648 --- /dev/null +++ b/backend/src/kbqa/rag/vector_store.py @@ -0,0 +1,94 @@ +import json + +from pymilvus import DataType, MilvusClient + +from kbqa.rag.types import PreparedChunk, VectorHit + + +class MilvusVectorStore: + collection_name = "kbqa_chunks" + + def __init__(self, client: MilvusClient, embedding_dim: int) -> None: + self.client = client + self.embedding_dim = embedding_dim + + def ensure_collection(self) -> None: + if self.client.has_collection(self.collection_name): + return + schema = self.client.create_schema(auto_id=False, enable_dynamic_field=False) + schema.add_field( + field_name="chunk_id", + datatype=DataType.VARCHAR, + is_primary=True, + max_length=36, + ) + schema.add_field( + field_name="document_id", + datatype=DataType.VARCHAR, + max_length=36, + ) + schema.add_field(field_name="order_index", datatype=DataType.INT64) + schema.add_field( + field_name="vector", + datatype=DataType.FLOAT_VECTOR, + dim=self.embedding_dim, + ) + index_params = self.client.prepare_index_params() + index_params.add_index( + field_name="vector", + index_type="AUTOINDEX", + metric_type="COSINE", + ) + self.client.create_collection( + collection_name=self.collection_name, + schema=schema, + index_params=index_params, + ) + + def upsert(self, chunks: list[PreparedChunk], vectors: list[list[float]]) -> None: + if len(chunks) != len(vectors): + raise ValueError("chunks and vectors must have equal lengths") + if not chunks: + return + data = [ + { + "chunk_id": chunk.id, + "document_id": chunk.document_id, + "order_index": chunk.order_index, + "vector": vector, + } + for chunk, vector in zip(chunks, vectors, strict=True) + ] + self.client.upsert(collection_name=self.collection_name, data=data) + + def search(self, vector: list[float], limit: int) -> list[VectorHit]: + results = self.client.search( + collection_name=self.collection_name, + data=[vector], + anns_field="vector", + limit=limit, + output_fields=["document_id", "order_index"], + search_params={"metric_type": "COSINE"}, + ) + hits: list[VectorHit] = [] + for item in results[0] if results else []: + entity = item.get("entity", {}) + hits.append( + VectorHit( + chunk_id=str(item["id"]), + document_id=str(entity["document_id"]), + order_index=int(entity["order_index"]), + score=float(item["distance"]), + ) + ) + return hits + + def delete_document(self, document_id: str) -> None: + encoded = json.dumps(document_id) + self.client.delete( + collection_name=self.collection_name, + filter=f"document_id == {encoded}", + ) + + def ping(self) -> None: + self.client.list_collections()