feat: add durable document indexing pipeline
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user