From ed65b2ab4bd433ca40dec1dc59e3bb74bf5295d6 Mon Sep 17 00:00:00 2001
From: Heiko Joerg Schick <heiko.schick@huawei.com>
Date: Tue, 5 Sep 2023 14:48:54 +0200
Subject: [PATCH] Improved coding style and comment for read_yaml_config and
 set_openai_api_key function

---
 yolo.py | 81 ++++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 52 insertions(+), 29 deletions(-)

diff --git a/yolo.py b/yolo.py
index c1ad174..92551a4 100755
--- a/yolo.py
+++ b/yolo.py
@@ -19,17 +19,62 @@ import pyperclip
 from termcolor import colored
 from colorama import init
 
-def read_config() -> any:
+CONFIG_FILE = "yolo.yaml"
+
+def read_yaml_config() -> any:
     """
-    Find the executing directory (e.g. in case an alias is set), so we can find the config file.
+    Read the configuration file from the executing directory.
+
+    This function determines the execution folder (which may vary if an alias is set) in order to
+    find the configuration file. It reads the file and returns its content in a Python data
+    structure.
+
+    Returns:
+        The content of the configuration file. Could be dictionary, list, etc. depending on
+        the YAML file structure.
     """
     yolo_path = os.path.abspath(__file__)
     prompt_path = os.path.dirname(yolo_path)
 
-    config_file = os.path.join(prompt_path, "yolo.yaml")
+    config_file = os.path.join(prompt_path, CONFIG_FILE)
     with open(config_file, 'r') as file:
         return yaml.safe_load(file)
 
+def set_openai_api_key(config):
+    """
+    Set the OpenAI API key by attempting several methods.
+
+    This function first tries to grab the OpenAI API key from environment variables,
+    if not found, it then looks for the key in the `.openai.apikey` in the home directory,
+    and lastly, it will look in the provided config dictionary. It sets the `openai.api_key`
+    with the retrieved key.
+
+    Parameters:
+        config (dict): A dictionary containing configuration values.
+                       It may contain `openai_api_key` as one of the keys.
+    """
+    dotenv.load_dotenv()
+
+    # Method 1: Read API key from environment variable
+    # The user can set their OpenAI API key by creating a ".env" file in the same
+    # directory as this script or by exporting it to their environment variables.
+    # The file or environment variable should contain the line `OPENAI_API_KEY="<yourkey>"`.
+    config["openai_api_key"] = os.getenv("OPENAI_API_KEY")
+
+    # Method 2: Read API key from a file in the home directory
+    # The user can also place a file named ".openai.apikey" in their home directory,
+    # which includes the API key in raw format. This method might be deprecated in future versions.
+    if not openai.api_key:  # Check this to avoid potential "invalid filepath" error.
+        home_path = os.path.expanduser("~")
+        openai.api_key_path = os.path.join(home_path, ".openai.apikey")
+
+    # Method 3: Read API key from the provided config dictionary
+    # The final method to set the API key is by providing it in the 'config' dictionary under the
+    # key 'openai_api_key'. For instance, in a `yolo.yaml` config file, it would appear as
+    # `openai_apikey: <yourkey>`.
+    if not openai.api_key:
+        openai.api_key = config["openai_api_key"]
+
 # Construct the prompt
 def get_full_prompt(user_prompt, shell):
     """
@@ -83,30 +128,6 @@ def get_os_friendly_name():
     else:
         return os_name
 
-def set_api_key(config):
-    """
-    Set the OpenAI API key.
-    """
-
-    # Two options for the user to specify they openai api key.
-    #1. Place a ".env" file in same directory as this with the line:
-    #   OPENAI_API_KEY="<yourkey>"
-    #   or do `export OPENAI_API_KEY=<yourkey>` before use
-    dotenv.load_dotenv()
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-
-    #2. Place a ".openai.apikey" in the home directory that holds the line:
-    #   <yourkey>
-    #   Note: This options will likely be removed in the future
-    if not openai.api_key:  #If statement to avoid "invalid filepath" error
-        home_path = os.path.expanduser("~")
-        openai.api_key_path = os.path.join(home_path,".openai.apikey")
-
-    #3. Final option is the key might be in the yolo.yaml config file
-    #   openai_apikey: <yourkey>
-    if not openai.api_key:
-        openai.api_key = config["openai_api_key"]
-
 def call_open_ai(config, shell, query):
     """
     Do we have a prompt from the user?
@@ -209,8 +230,10 @@ def main():
     """
     Defined starting point of source code.
     """
-    config = read_config()
-    set_api_key(config)
+    config = read_yaml_config()
+    set_openai_api_key(config)
+
+    # set_api_key(config)
 
     # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
     shell = os.environ.get("SHELL", "powershell.exe")