Refactor configuration handling, add setup command. Resolves #3

This commit introduces a significant overhaul of the configuration handling in the AICodeBot. The OpenAI API key is no longer fetched directly from the environment variables in the cli.py file. Instead, a new helper function, get_config_file, has been added to helpers.py. This function checks if the AICODEBOT_CONFIG_FILE environment variable is set and uses its value as the path to the configuration file. If the environment variable is not set, it defaults to using a file named .aicodebot.yaml in the user's home directory.

The setup command in cli.py has been updated to use this new function. It now also accepts a new option, --gpt-4-supported, which is a flag indicating whether the user has access to GPT-4. This information is stored in the configuration file along with the OpenAI API key.

The setup_config function in cli.py has been updated to reflect these changes. It now accepts two arguments: the OpenAI API key and a boolean indicating whether GPT-4 is supported. If the OpenAI API key is not provided, the function will attempt to fetch it from the environment variables or prompt the user to enter it. The function then validates the API key and checks if GPT-4 is supported (unless the --gpt-4-supported flag was used). The configuration data is then written to the configuration file.

The read_config function in helpers.py has also been updated to use the new get_config_file function.

Finally, tests have been added to test_cli.py to test the new setup command and the changes to the configuration handling.

Author: TechNickAI

aicodebot/cli.py

@@ -1,6 +1,7 @@
 from aicodebot import version as aicodebot_version
 from aicodebot.helpers import (
     exec_and_get_output,
+    get_config_file,
     get_llm_model,
     get_token_length,
     git_diff_context,
@@ -26,7 +27,6 @@
 PRECISE_TEMPERATURE = 0
 CREATIVE_TEMPERATURE = 0.7
 DEFAULT_SPINNER = "point"
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
 # ----------------------- Setup for rich console output ---------------------- #
 
@@ -295,23 +295,24 @@ def review(commit, verbose):
 
 @cli.command()
 @click.option("--openai-api-key", "-k", help="Your OpenAI API key")
-def setup(openai_api_key):
+@click.option("--gpt-4-supported", "-4", help="Whether you have access to GPT-4", is_flag=True)
+def setup(openai_api_key, gpt_4_supported):
     """Set up the configuration file with your OpenAI API key
     If the config file already exists, it will ask you if you want to remove it and recreate it.
     """
-    config_file = Path(read_config.CONFIG_FILE)
+    config_file = get_config_file()
     if config_file.exists():
         if not click.confirm(
             f"The config file already exists at {config_file}. Do you want to remove it and recreate it?"
         ):
             console.print("Setup cancelled. 🚫")
-            return
+            sys.exit(1)
 
         # Remove the existing config file
         config_file.unlink()
 
     # Call the setup_config function with the provided arguments
-    setup_config(openai_api_key)
+    setup_config(openai_api_key, gpt_4_supported)
 
 
 @cli.command
@@ -392,18 +393,21 @@ def sidekick(request, verbose, files):
 # ---------------------------------------------------------------------------- #
 
 
-def setup_config(openai_api_key=OPENAI_API_KEY):
+def setup_config(openai_api_key=None, gpt_4_supported=None):
     config = read_config()
+    openai.api_key = openai_api_key
     if config:
         openai.api_key = config["OPENAI_API_KEY"]
+        logger.success(f"Using OpenAI API key from {get_config_file()}")
         return config
-    elif openai_api_key:
-        openai.api_key = openai_api_key
+    elif os.getenv("OPENAI_API_KEY"):
+        logger.info("Using OPENAI_API_KEY environment variable")
+        openai.api_key = os.getenv("OPENAI_API_KEY")
 
-    console.print(
-        f"[bold red]The config file does not exist.[/bold red]\n"
-        f"Let's set that up for you at {read_config.CONFIG_FILE}\n"
-    )
+    logger.info(f"openai.api_key {openai.api_key}")
+
+    config_file = get_config_file()
+    console.print(f"[bold red]The config file does not exist.[/bold red]\nLet's set that up for you at {config_file}\n")
 
     if not openai.api_key:
         openai_api_key_url = "https://platform.openai.com/account/api-keys"
@@ -417,30 +421,33 @@ def setup_config(openai_api_key=OPENAI_API_KEY):
 
         openai.api_key = click.prompt("Please enter your OpenAI API key")
 
+    logger.info(f"openai.api_key 2 {openai.api_key}")
     # Validate the API key and check if it supports GPT-4
-    try:
-        click.echo("Validating the API key, and checking if GPT-4 is supported...")
-        engines = engine.Engine.list()
-        logger.trace(f"Engines: {engines}")
-        gpt_4_supported = "gpt-4" in [engine.id for engine in engines.data]
-        if gpt_4_supported:
-            click.echo("✅ The API key is valid and supports GPT-4.")
-        else:
-            click.echo("✅ The API key is valid, but does not support GPT-4. GPT-3.5 will be used instead.")
-    except Exception as e:
-        raise click.ClickException(f"Failed to validate the API key: {str(e)}") from e
+    if gpt_4_supported is None:
+        try:
+            click.echo("Validating the API key, and checking if GPT-4 is supported...")
+            engines = engine.Engine.list()
+            logger.trace(f"Engines: {engines}")
+            gpt_4_supported = "gpt-4" in [engine.id for engine in engines.data]
+            if gpt_4_supported:
+                click.echo("✅ The API key is valid and supports GPT-4.")
+            else:
+                click.echo("✅ The API key is valid, but does not support GPT-4. GPT-3.5 will be used instead.")
+        except Exception as e:
+            raise click.ClickException(f"Failed to validate the API key: {str(e)}") from e
 
     config_data = {"config_version": 1, "OPENAI_API_KEY": openai.api_key, "gpt_4_supported": gpt_4_supported}
 
-    with Path.open(read_config.CONFIG_FILE, "w") as f:
+    logger.info(f"openai.api_key 3 {openai.api_key}")
+
+    with Path.open(config_file, "w") as f:
         yaml.dump(config_data, f)
 
     console.print(
-        f"[bold green]Created {read_config.CONFIG_FILE} with your OpenAI API key.[/bold green] "
+        f"[bold green]Created {config_file} with your OpenAI API key.[/bold green] "
         "Now, please re-run aicodebot and let's get started!"
     )
     sys.exit(0)
-    return config_data
 
 
 class RichLiveCallbackHandler(BaseCallbackHandler):

aicodebot/helpers.py

@@ -124,14 +124,22 @@ def exec_and_get_output(command):
     return result.stdout
 
 
+def get_config_file():
+    if "AICODEBOT_CONFIG_FILE" in os.environ:
+        config_file = Path(os.getenv("AICODEBOT_CONFIG_FILE"))
+    else:
+        config_file = Path(Path.home() / ".aicodebot.yaml")
+    logger.debug(f"Using config file {config_file}")
+    return config_file
+
+
 def read_config():
-    if read_config.CONFIG_FILE.exists():
-        logger.debug(f"Config file {read_config.CONFIG_FILE} exists")
-        with Path(read_config.CONFIG_FILE).open("r") as f:
+    """Read the config file and return its contents as a dictionary."""
+    config_file = get_config_file()
+    if config_file.exists():
+        logger.debug(f"Config file {config_file} exists")
+        with Path(config_file).open("r") as f:
             return yaml.safe_load(f)
     else:
-        logger.debug(f"Config file {read_config.CONFIG_FILE} does not exist")
+        logger.debug(f"Config file {config_file} does not exist")
         return None
-
-
-read_config.CONFIG_FILE = Path(Path.home() / ".aicodebot.yaml")

tests/test_cli.py

@@ -1,5 +1,7 @@
 from aicodebot import version as aicodebot_version
 from aicodebot.cli import cli
+from aicodebot.helpers import read_config
+from pathlib import Path
 import os, pytest
 
 
@@ -33,3 +35,63 @@ def test_debug_failure(cli_runner):
     result = cli_runner.invoke(cli, ["debug", "ls", "-9"])
     assert result.exit_code > 0, f"Output: {result.output}"
     assert "ls -9" in result.output
+
+
+def test_setup(cli_runner, tmp_path, monkeypatch):
+    # Unset the environment variable for the OpenAI API key
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    assert not os.getenv("OPENAI_API_KEY")
+
+    temp_config_file = Path(tmp_path / ".aicodebot.test.yaml")
+    # set AICODEBOT_CONFIG_FILE to the temp config file
+    monkeypatch.setenv("AICODEBOT_CONFIG_FILE", str(temp_config_file))
+
+    assert os.getenv("AICODEBOT_CONFIG_FILE") == str(temp_config_file)
+
+    assert read_config() is None
+
+    # Run the setup command
+    result = cli_runner.invoke(cli, ["setup", "--openai-api-key", "fake_api_key", "--gpt-4-supported"])
+
+    # Check if the command was successful
+    assert result.exit_code == 0, f"Output: {result.output}"
+
+    # Check if the config file was created
+    assert Path(temp_config_file).exists()
+
+    # Load the config file
+    config_data = read_config()
+
+    # Check if the config file contains the correct data
+    assert config_data == {"config_version": 1, "OPENAI_API_KEY": "fake_api_key", "gpt_4_supported": True}
+
+    # Run the setup command again, should fail because the config file already exists
+    result = cli_runner.invoke(cli, ["setup", "--openai-api-key", "fake_api_key", "--gpt-4-supported"])
+    assert result.exit_code == 1, f"Output: {result.output}"
+    assert "Setup cancelled" in result.output
+
+
+def test_setup_with_openai_key(cli_runner, tmp_path, monkeypatch):
+    monkeypatch.setenv("OPENAI_API_KEY", "fake_api_key2")
+    assert os.getenv("OPENAI_API_KEY") == "fake_api_key2"
+
+    temp_config_file = Path(tmp_path / ".aicodebot.test.yaml")
+    # set AICODEBOT_CONFIG_FILE to the temp config file
+    monkeypatch.setenv("AICODEBOT_CONFIG_FILE", str(temp_config_file))
+
+    assert read_config() is None
+
+    # Run the setup command
+    result = cli_runner.invoke(cli, ["setup"])
+
+    # Check if the command was successful
+    assert result.exit_code == 0, f"Output: {result.output}"
+
+    # Check if the config file was created
+    assert Path(temp_config_file).exists()
+
+    # Load the config file
+    config_data = read_config()
+
+    # Check if the config file contains the correct data
+    assert config_data == {"config_version": 1, "OPENAI_API_KEY": "fake_api_key2", "gpt_4_supported": False}