Add a first run experience to intialize the admin user if none exists and setup chat models

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)

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(

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": {

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)