Implemented the '--ai-settings' flag

huangweichun1997 · Apr 13, 2023 · 0f6fba7 · 0f6fba7
1 parent 428caa9
commit 0f6fba7
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/scripts/config.py b/scripts/config.py
@@ -39,6 +39,7 @@ def __init__(self):
         self.speak_mode = False
         self.skip_reprompt = False
 
+        self.ai_settings_file = os.getenv("AI_SETTINGS_FILE", "ai_settings.yaml")
         self.fast_llm_model = os.getenv("FAST_LLM_MODEL", "gpt-3.5-turbo")
         self.smart_llm_model = os.getenv("SMART_LLM_MODEL", "gpt-4")
         self.fast_token_limit = int(os.getenv("FAST_TOKEN_LIMIT", 4000))

diff --git a/scripts/main.py b/scripts/main.py
@@ -182,7 +182,7 @@ def load_variables(config_file="config.yaml"):
 
 def construct_prompt():
     """Construct the prompt for the AI to respond to"""
-    config = AIConfig.load()
+    config = AIConfig.load(cfg.ai_settings_file)
     if cfg.skip_reprompt and config.ai_name:
         logger.typewriter_log("Name :", Fore.GREEN, config.ai_name)
         logger.typewriter_log("Role :", Fore.GREEN, config.ai_role)
@@ -324,7 +324,21 @@ def parse_arguments():
     if args.skip_reprompt:
         logger.typewriter_log("Skip Re-prompt: ", Fore.GREEN, "ENABLED")
         cfg.skip_reprompt = True
+
+    if args.ai_settings_file:
+        file = args.ai_settings_file
+
+        # Validate file
+        (validated, message) = utils.validate_yaml_file(file)
+        if not validated:
+            logger.typewriter_log("FAILED FILE VALIDATION", Fore.RED, message)
+            exit(1)
 
+        logger.typewriter_log("Using AI Settings File:", Fore.GREEN, file)
+        cfg.ai_settings_file = file
+        cfg.skip_reprompt = True
+
+
 
 # TODO: fill in llm values here
 check_openai_api_key()

diff --git a/scripts/utils.py b/scripts/utils.py
@@ -1,3 +1,6 @@
+import yaml
+from colorama import Fore
+
 def clean_input(prompt: str=''):
     try:
         return input(prompt)
@@ -6,3 +9,14 @@ def clean_input(prompt: str=''):
         print("Quitting...")
         exit(0)
 
+
+def validate_yaml_file(file: str):
+    try:
+        with open(file) as file:
+            yaml.load(file, Loader=yaml.FullLoader)
+    except FileNotFoundError:
+        return (False, f"The file {Fore.CYAN}`{file}`{Fore.RESET} wasn't found")
+    except yaml.YAMLError as e:
+        return (False, f"There was an issue while trying to read with your AI Settings file: {e}")
+
+    return (True, f"Successfully validated {Fore.CYAN}`{file}`{Fore.RESET}!")