From da4a5efc491601a058f626263ab4eb650210b32e Mon Sep 17 00:00:00 2001 From: binary-husky Date: Sun, 6 Oct 2024 16:26:26 +0000 Subject: [PATCH] lazy load llama-index lib --- crazy_functions/Rag_Interface.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/crazy_functions/Rag_Interface.py b/crazy_functions/Rag_Interface.py index d83d8ca5a2..1bc740adbf 100644 --- a/crazy_functions/Rag_Interface.py +++ b/crazy_functions/Rag_Interface.py @@ -2,20 +2,7 @@ from crazy_functions.crazy_utils import input_clipping from crazy_functions.crazy_utils import request_gpt_model_in_new_thread_with_ui_alive -VECTOR_STORE_TYPE = "Milvus" - -if VECTOR_STORE_TYPE == "Milvus": - try: - from crazy_functions.rag_fns.milvus_worker import MilvusRagWorker as LlamaIndexRagWorker - except: - VECTOR_STORE_TYPE = "Simple" - -if VECTOR_STORE_TYPE == "Simple": - from crazy_functions.rag_fns.llama_index_worker import LlamaIndexRagWorker - - RAG_WORKER_REGISTER = {} - MAX_HISTORY_ROUND = 5 MAX_CONTEXT_TOKEN_LIMIT = 4096 REMEMBER_PREVIEW = 1000 @@ -23,6 +10,16 @@ @CatchException def Rag问答(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request): + # import vector store lib + VECTOR_STORE_TYPE = "Milvus" + if VECTOR_STORE_TYPE == "Milvus": + try: + from crazy_functions.rag_fns.milvus_worker import MilvusRagWorker as LlamaIndexRagWorker + except: + VECTOR_STORE_TYPE = "Simple" + if VECTOR_STORE_TYPE == "Simple": + from crazy_functions.rag_fns.llama_index_worker import LlamaIndexRagWorker + # 1. we retrieve rag worker from global context user_name = chatbot.get_user() checkpoint_dir = get_log_folder(user_name, plugin_name='experimental_rag')