Initial changes to support multiple search model configurations

- All search models are loaded into memory, and stored in a dictionary indexed by name
- Still need to add database migrations and create a UI for user to select their choice. Presently, it uses the default option
This commit is contained in:
sabaimran
2023-12-05 00:35:40 -05:00
parent d2ddbef08f
commit ef21d78c99
9 changed files with 46 additions and 19 deletions

View File

@@ -45,8 +45,10 @@ def enable_db_access_for_all_tests(db):
@pytest.fixture(scope="session")
def search_config() -> SearchConfig:
state.embeddings_model = EmbeddingsModel()
state.cross_encoder_model = CrossEncoderModel()
state.embeddings_model = dict()
state.embeddings_model["default"] = EmbeddingsModel()
state.cross_encoder_model = dict()
state.cross_encoder_model["default"] = CrossEncoderModel()
model_dir = resolve_absolute_path("~/.khoj/search")
model_dir.mkdir(parents=True, exist_ok=True)
@@ -317,8 +319,10 @@ def client(
state.config.content_type = content_config
state.config.search_type = search_config
state.SearchType = configure_search_types()
state.embeddings_model = EmbeddingsModel()
state.cross_encoder_model = CrossEncoderModel()
state.embeddings_model = dict()
state.embeddings_model["default"] = EmbeddingsModel()
state.cross_encoder_model = dict()
state.cross_encoder_model["default"] = CrossEncoderModel()
# These lines help us Mock the Search models for these search types
state.search_models.image_search = image_search.initialize_model(search_config.image)