From 8c36079f743b687532f77bbbcd591749a61e6b92 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Mon, 13 Nov 2023 21:07:12 -0800 Subject: [PATCH] Add a first run experience to intialize the admin user if none exists and setup chat models --- src/khoj/main.py | 9 +++- src/khoj/routers/helpers.py | 4 +- src/khoj/utils/constants.py | 1 + src/khoj/utils/initialization.py | 82 ++++++++++++++++++++++++++++++++ 4 files changed, 92 insertions(+), 4 deletions(-) create mode 100644 src/khoj/utils/initialization.py diff --git a/src/khoj/main.py b/src/khoj/main.py index f92e8cbe..53dae1bc 100644 --- a/src/khoj/main.py +++ b/src/khoj/main.py @@ -56,6 +56,7 @@ locale.setlocale(locale.LC_ALL, "") from khoj.configure import configure_routes, initialize_server, configure_middleware from khoj.utils import state from khoj.utils.cli import cli +from khoj.utils.initialization import initialization # Setup Logger rich_handler = RichHandler(rich_tracebacks=True) @@ -74,8 +75,7 @@ def run(should_start_server=True): args = cli(state.cli_args) set_state(args) - # Create app directory, if it doesn't exist - state.config_file.parent.mkdir(parents=True, exist_ok=True) + logger.info(f"🚒 Initializing Khoj v{state.khoj_version}") # Set Logging Level if args.verbose == 0: @@ -83,6 +83,11 @@ def run(should_start_server=True): elif args.verbose >= 1: logger.setLevel(logging.DEBUG) + initialization() + + # Create app directory, if it doesn't exist + state.config_file.parent.mkdir(parents=True, exist_ok=True) + # Set Log File fh = logging.FileHandler(state.config_file.parent / "khoj.log", encoding="utf-8") fh.setLevel(logging.DEBUG) diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py index 2fc7ab79..a1af10f2 100644 --- a/src/khoj/routers/helpers.py +++ b/src/khoj/routers/helpers.py @@ -159,8 +159,8 @@ def generate_chat_response( conversation_config = ConversationAdapters.get_default_conversation_config() openai_chat_config = ConversationAdapters.get_openai_conversation_config() if offline_chat_config and offline_chat_config.enabled and conversation_config.model_type == "offline": - if state.gpt4all_processor_config.loaded_model is None: - state.gpt4all_processor_config = GPT4AllProcessorModel(offline_chat_config.chat_model) + if state.gpt4all_processor_config is None or state.gpt4all_processor_config.loaded_model is None: + state.gpt4all_processor_config = GPT4AllProcessorModel(conversation_config.chat_model) loaded_model = state.gpt4all_processor_config.loaded_model chat_response = converse_offline( diff --git a/src/khoj/utils/constants.py b/src/khoj/utils/constants.py index e9d431c6..8a106153 100644 --- a/src/khoj/utils/constants.py +++ b/src/khoj/utils/constants.py @@ -6,6 +6,7 @@ empty_escape_sequences = "\n|\r|\t| " app_env_filepath = "~/.khoj/env" telemetry_server = "https://khoj.beta.haletic.com/v1/telemetry" content_directory = "~/.khoj/content/" +default_offline_chat_model = "mistral-7b-instruct-v0.1.Q4_0.gguf" empty_config = { "search-type": { diff --git a/src/khoj/utils/initialization.py b/src/khoj/utils/initialization.py new file mode 100644 index 00000000..af9a0a87 --- /dev/null +++ b/src/khoj/utils/initialization.py @@ -0,0 +1,82 @@ +import logging + +from database.models import ( + KhojUser, + OfflineChatProcessorConversationConfig, + OpenAIProcessorConversationConfig, + ChatModelOptions, +) + +from khoj.utils.constants import default_offline_chat_model + +from database.adapters import ConversationAdapters + + +logger = logging.getLogger(__name__) + + +def initialization(): + def _create_admin_user(): + logger.info("👩‍✈️ Setting up admin user") + email_addr = input("Email Address: ") + password = input("Password: ") + admin_user = KhojUser.objects.create_superuser(email=email_addr, username=email_addr, password=password) + logger.info(f"👩‍✈️ Created admin user: {admin_user.email}") + + def _create_chat_configuration(): + logger.info( + "🗣️ Configure chat models available to your server. You can always update these at /django/admin using the credentials of your admin account" + ) + use_offline_model = input("Use offline chat model? (y/n): ") + if use_offline_model == "y": + logger.info("🗣️ Setting up offline chat model") + OfflineChatProcessorConversationConfig.objects.create(enabled=True) + + offline_chat_model = input( + f"Enter the name of the offline chat model you want to use, based on the models in HuggingFace (press enter to use the default: {default_offline_chat_model}): " + ) + if offline_chat_model == "": + ChatModelOptions.objects.create( + chat_model=default_offline_chat_model, model_type=ChatModelOptions.ModelType.OFFLINE + ) + else: + max_tokens = input("Enter the maximum number of tokens to use for the offline chat model:") + tokenizer = input("Enter the tokenizer to use for the offline chat model:") + ChatModelOptions.objects.create( + chat_model=offline_chat_model, + model_type=ChatModelOptions.ModelType.OFFLINE, + max_prompt_size=max_tokens, + tokenizer=tokenizer, + ) + + use_openai_model = input("Use OpenAI chat model? (y/n): ") + + if use_openai_model == "y": + logger.info("🗣️ Setting up OpenAI chat model") + api_key = input("Enter your OpenAI API key: ") + OpenAIProcessorConversationConfig.objects.create(api_key=api_key) + openai_chat_model = input("Enter the name of the OpenAI chat model you want to use: ") + max_tokens = input("Enter the maximum number of tokens to use for the OpenAI chat model:") + ChatModelOptions.objects.create( + chat_model=openai_chat_model, model_type=ChatModelOptions.ModelType.OPENAI, max_tokens=max_tokens + ) + + logger.info("🗣️ Chat model configuration complete") + + admin_user = KhojUser.objects.filter(is_staff=True).first() + if admin_user is None: + while True: + try: + _create_admin_user() + break + except Exception as e: + logger.error(f"🚨 Failed to create admin user: {e}", exc_info=True) + + chat_config = ConversationAdapters.get_default_conversation_config() + if admin_user is None and chat_config is None: + while True: + try: + _create_chat_configuration() + break + except Exception as e: + logger.error(f"🚨 Failed to create chat configuration: {e}", exc_info=True)