Add create collection to app startup

This commit is contained in:
Weves
2023-05-17 15:36:10 -07:00
committed by Chris Weaver
parent bdebb9d441
commit 8685beceb2
2 changed files with 19 additions and 1 deletions

View File

@@ -27,6 +27,18 @@ logger = setup_logger()
DEFAULT_BATCH_SIZE = 30
def create_collection(
collection_name: str, embedding_dim: int = DOC_EMBEDDING_DIM
) -> None:
logger.info(f"Attempting to create collection {collection_name}")
result = get_qdrant_client().create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=embedding_dim, distance=Distance.COSINE),
)
if not result:
raise RuntimeError("Could not create Qdrant collection")
def recreate_collection(
collection_name: str, embedding_dim: int = DOC_EMBEDDING_DIM
) -> None:

View File

@@ -93,9 +93,15 @@ def get_application() -> FastAPI:
@application.on_event("startup")
def startup_event() -> None:
# To avoid circular imports
from danswer.semantic_search.semantic_search import (
warm_up_models,
) # To avoid circular imports
)
from danswer.datastores.qdrant.indexing import create_collection
from danswer.configs.app_configs import QDRANT_DEFAULT_COLLECTION
create_collection(collection_name=QDRANT_DEFAULT_COLLECTION)
logger.info("Collection ready")
warm_up_models()
logger.info("Semantic Search models are ready.")