Improved coding style and comment for read_yaml_config and set_openai_api_key function

This commit is contained in:
Heiko Joerg Schick
2023-09-05 14:48:54 +02:00
parent 34eb29b926
commit ed65b2ab4b
+52 -29
View File
@@ -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")