Improved coding style and comment for read_yaml_config and set_openai_api_key function
This commit is contained in:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user