gqt
136 lines
5.3 KiB
Python
136 lines
5.3 KiB
Python
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,
|
|
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
|
|
|
|
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:
|
|
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)
|
|
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)
|
|
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:
|
|
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, document_id: str, chunks: list[PreparedChunk]) -> bool:
|
|
batch_size = 10
|
|
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])
|
|
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:
|
|
await DocumentRepository(session).mark_failed(document_id, code, message)
|