Embedding Model Swap Changes (#35)

This commit is contained in:
Yuhong Sun 2024-01-29 00:05:05 -08:00 committed by Chris Weaver
parent 15934ee268
commit bac34a47b2
2 changed files with 18 additions and 3 deletions

View File

@ -5,6 +5,7 @@ from sqlalchemy.orm import Session
from danswer.auth.users import current_user
from danswer.configs.chat_configs import DISABLE_LLM_CHUNK_FILTER
from danswer.configs.chat_configs import NUM_RETURNED_HITS
from danswer.db.embedding_model import get_current_db_embedding_model
from danswer.db.engine import get_session
from danswer.db.models import User
from danswer.document_index.factory import get_default_document_index
@ -66,9 +67,13 @@ def handle_search_request(
skip_llm_chunk_filter=disable_llm_chunk_filter,
)
db_embedding_model = get_current_db_embedding_model(db_session)
document_index = get_default_document_index(
primary_index_name=db_embedding_model.index_name, secondary_index_name=None
)
top_chunks, llm_selection = full_chunk_search(
query=search_query,
document_index=get_default_document_index(),
query=search_query, document_index=document_index, db_session=db_session
)
top_docs = chunks_to_search_docs(top_chunks)

View File

@ -2,6 +2,8 @@ from sqlalchemy.orm import Session
from danswer.access.access import get_access_for_documents
from danswer.db.document import prepare_to_modify_documents
from danswer.db.embedding_model import get_current_db_embedding_model
from danswer.db.embedding_model import get_secondary_db_embedding_model
from danswer.db.engine import get_sqlalchemy_engine
from danswer.document_index.factory import get_default_document_index
from danswer.document_index.interfaces import DocumentIndex
@ -46,7 +48,15 @@ def _sync_user_group_batch(
def sync_user_groups(user_group_id: int, db_session: Session) -> None:
"""Sync the status of Postgres for the specified user group"""
document_index = get_default_document_index()
db_embedding_model = get_current_db_embedding_model(db_session)
secondary_db_embedding_model = get_secondary_db_embedding_model(db_session)
document_index = get_default_document_index(
primary_index_name=db_embedding_model.index_name,
secondary_index_name=secondary_db_embedding_model.index_name
if secondary_db_embedding_model
else None,
)
user_group = fetch_user_group(db_session=db_session, user_group_id=user_group_id)
if user_group is None: