diff --git a/scripts/config.py b/scripts/config.py index fd370a72..ad968fb2 100644 --- a/scripts/config.py +++ b/scripts/config.py @@ -39,6 +39,7 @@ class Config(metaclass=Singleton): 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 index f81b09a7..07d2bbd2 100644 --- 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 index 5039796f..bca8d4a8 100644 --- 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}!") \ No newline at end of file