Update model selection logic to support GPT-4 🚀

This commit introduces changes to the model selection logic in the AI code bot. Previously, the bot was hardcoded to use the GPT-3.5-turbo model. Now, the bot checks if GPT-4 is supported by the provided OpenAI API key and uses it if available. If GPT-4 is not supported, the bot falls back to using GPT-3.5-turbo.

The model selection is also now dynamic based on the token size of the request. The bot selects the largest model that supports the required token size. This ensures that the bot can handle larger contexts while still using the most powerful model available.

Additionally, a new environment variable `GPT_4_SUPPORTED` has been added to the `.aicodebot.template` file to store the GPT-4 support status of the API key. This value is set during the setup process when the API key is validated.

This update enhances the bot's capabilities and prepares it for future improvements in the OpenAI models.

Author: TechNickAI

aicodebot/.aicodebot.template

@@ -1,3 +1,4 @@
 # Obtain your API key from https://platform.openai.com/account/api-keys
 # sk-...
 OPENAI_API_KEY=
+GPT_4_SUPPORTED=false

aicodebot/cli.py

@@ -4,6 +4,7 @@
 from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
 from langchain.prompts import load_prompt
+from openai.api_resources import engine
 from pathlib import Path
 from rich.console import Console
 from rich.style import Style
@@ -13,7 +14,6 @@
 
 DEFAULT_MAX_TOKENS = 1024
 DEFAULT_TEMPERATURE = 0.1
-DEFAULT_MODEL = "gpt-3.5-turbo"  # Can't wait to use GPT-4, as the results are much better. On the waitlist.
 DEFAULT_SPINNER = "point"
 
 # ----------------------- Setup for rich console output ---------------------- #
@@ -50,9 +50,8 @@ def alignment(verbose):
     prompt = load_prompt(Path(__file__).parent / "prompts" / "alignment.yaml")
 
     # Set up the language model
-    llm = ChatOpenAI(
-        model=DEFAULT_MODEL, temperature=DEFAULT_TEMPERATURE, max_tokens=DEFAULT_MAX_TOKENS, verbose=verbose
-    )
+    model = get_llm_model(get_token_length(prompt.template))
+    llm = ChatOpenAI(model=model, temperature=DEFAULT_TEMPERATURE, max_tokens=DEFAULT_MAX_TOKENS, verbose=verbose)
 
     # Set up the chain
     chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose)
@@ -100,18 +99,10 @@ def commit(verbose, response_token_size, yes, skip_pre_commit):
         sys.exit(0)
 
     # Check the size of the diff context and adjust accordingly
-    prompt_token_size = get_token_length(diff_context) + get_token_length(prompt.template)
+    request_token_size = get_token_length(diff_context) + get_token_length(prompt.template)
+    model = get_llm_model(request_token_size)
     if verbose:
-        console.print(f"Diff context token size: {prompt_token_size}")
-
-    if prompt_token_size + response_token_size > 16_000:
-        # Bigger models coming soon
-        console.print("The diff context is too large to review. 😞")
-        sys.exit(1)
-    elif prompt_token_size + response_token_size > 3_500:  # It's actually 4k, but we want a buffer
-        model = "gpt-3.5-turbo-16k"  # supports 16k tokens but is a bit slower and more expensive
-    else:
-        model = DEFAULT_MODEL  # gpt-3.5-turbo supports 4k tokens
+        console.print(f"Diff context token size: {request_token_size}, using model: {model}")
 
     # Set up the language model
     llm = ChatOpenAI(model=model, temperature=DEFAULT_TEMPERATURE, max_tokens=DEFAULT_MAX_TOKENS, verbose=verbose)
@@ -149,39 +140,40 @@ def debug(command, verbose):
     """Run a command and debug the output."""
     setup_environment()
 
-    # Load the prompt
-    prompt = load_prompt(Path(__file__).parent / "prompts" / "debug.yaml")
-
-    # Set up the language model
-    llm = ChatOpenAI(
-        model=DEFAULT_MODEL, temperature=DEFAULT_TEMPERATURE, max_tokens=DEFAULT_MAX_TOKENS, verbose=verbose
-    )
-
-    # Set up the chain
-    chat_chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose)
-    command_str = " ".join(command)
-
     # Run the command and capture its output
+    command_str = " ".join(command)
     console.print(f"Executing the command:\n{command_str}")
     process = subprocess.run(command_str, shell=True, capture_output=True, text=True)  # noqa: S602
 
     # Print the output of the command
     output = f"Standard Output:\n{process.stdout}\nStandard Error:\n{process.stderr}"
     console.print(output)
 
-    # Print a message about the exit status
+    # If it succeeded, exit
     if process.returncode == 0:
         console.print("✅ The command completed successfully.")
-    else:
-        console.print(f"The command exited with status {process.returncode}.")
+        sys.exit(0)
 
     # If the command failed, send its output to ChatGPT for analysis
-    if process.returncode != 0:
-        error_output = process.stderr
-        with console.status("Debugging", spinner=DEFAULT_SPINNER):
-            response = chat_chain.run(error_output)
-            console.print(response, style=bot_style)
-        sys.exit(process.returncode)
+    error_output = process.stderr
+
+    console.print(f"The command exited with status {process.returncode}.")
+
+    # Load the prompt
+    prompt = load_prompt(Path(__file__).parent / "prompts" / "debug.yaml")
+
+    # Set up the language model
+    model = get_llm_model(get_token_length(error_output) + get_token_length(prompt.template))
+    llm = ChatOpenAI(model=model, temperature=DEFAULT_TEMPERATURE, max_tokens=DEFAULT_MAX_TOKENS, verbose=verbose)
+
+    # Set up the chain
+    chat_chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose)
+
+    with console.status("Debugging", spinner=DEFAULT_SPINNER):
+        response = chat_chain.run(error_output)
+        console.print(response, style=bot_style)
+
+    sys.exit(process.returncode)
 
 
 @cli.command()
@@ -194,7 +186,8 @@ def fun_fact(verbose):
     prompt = load_prompt(Path(__file__).parent / "prompts" / "fun_fact.yaml")
 
     # Set up the language model
-    llm = ChatOpenAI(model=DEFAULT_MODEL, temperature=0.9, max_tokens=250, verbose=verbose)
+    model = get_llm_model(get_token_length(prompt.template))
+    llm = ChatOpenAI(model=model, temperature=0.9, max_tokens=250, verbose=verbose)
 
     # Set up the chain
     chat_chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose)
@@ -223,21 +216,13 @@ def review(commit, verbose):
 
     # Check the size of the diff context and adjust accordingly
     response_token_size = DEFAULT_MAX_TOKENS / 2
-    prompt_token_size = get_token_length(diff_context) + get_token_length(prompt.template)
+    request_token_size = get_token_length(diff_context) + get_token_length(prompt.template)
+    model = get_llm_model(request_token_size)
     if verbose:
-        console.print(f"Prompt token size: {prompt_token_size}")
-
-    if prompt_token_size + response_token_size > 16_000:
-        # Bigger models coming soon
-        console.print("The diff context is too large to review. 😞")
-        sys.exit(1)
-    elif prompt_token_size + response_token_size > 3_500:  # It's actually 4k, but we want a buffer
-        model = "gpt-3.5-turbo-16k"  # supports 16k tokens but is a bit slower and more expensive
-    else:
-        model = DEFAULT_MODEL  # gpt-3.5-turbo supports 4k tokens
+        console.print(f"Diff context token size: {request_token_size}, using model: {model}")
 
     # Set up the language model
-    llm = ChatOpenAI(model=model, temperature=DEFAULT_TEMPERATURE, max_tokens=DEFAULT_MAX_TOKENS, verbose=verbose)
+    llm = ChatOpenAI(model=model, temperature=DEFAULT_TEMPERATURE, max_tokens=response_token_size, verbose=verbose)
 
     # Set up the chain
     chain = LLMChain(llm=llm, prompt=prompt, verbose=verbose)
@@ -247,7 +232,9 @@ def review(commit, verbose):
         console.print(response, style=bot_style)
 
 
-# ------------------------------- End Commands ------------------------------- #
+# ---------------------------------------------------------------------------- #
+#                               Helper functions                               #
+# ---------------------------------------------------------------------------- #
 
 
 def setup_environment():
@@ -273,17 +260,32 @@ def setup_environment():
     if click.confirm(f"Create the {config_file} file for you?"):
         api_key = click.prompt("Please enter your OpenAI API key")
 
-        # Copy .env.template to .env and insert the API key
+        # Validate the API key and check if it supports GPT-4
+        openai.api_key = api_key
+        try:
+            click.echo("Validating the API key, and checking if GPT-4 is supported...")
+            engines = engine.Engine.list()
+            gpt_4_supported = "true" if "gpt-4" in [engine.id for engine in engines.data] else "false"
+            if gpt_4_supported == "true":
+                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
+
+        # Copy .env.template to .env and insert the API key and gpt_4_supported
         template_file = Path(__file__).parent / ".aicodebot.template"
         with Path.open(template_file, "r") as template, Path.open(config_file, "w") as env:
             for line in template:
                 if line.startswith("OPENAI_API_KEY="):
                     env.write(f"OPENAI_API_KEY={api_key}\n")
+                elif line.startswith("GPT_4_SUPPORTED="):
+                    env.write(f"GPT_4_SUPPORTED={gpt_4_supported}\n")
                 else:
                     env.write(line)
 
         console.print(
-            f"[bold green]Created {config_file} with your OpenAI API key.[/bold green] "
+            f"[bold green]Created {config_file} with your OpenAI API key and GPT-4 support status.[/bold green] "
             "Now, please re-run aicodebot and let's get started!"
         )
         sys.exit(0)
@@ -293,5 +295,32 @@ def setup_environment():
     )
 
 
+def get_llm_model(token_size):
+    # https://platform.openai.com/docs/models/gpt-3-5
+    # We want to use GPT-4, if it is available for this OPENAI_API_KEY, otherwise GPT-3.5
+    # We also want to use the largest model that supports the token size we need
+    model_options = {
+        "gpt-4": 8192,
+        "gpt-4-32k": 32768,
+        "gpt-3.5-turbo": 4096,
+        "gpt-3.5-turbo-16k": 16384,
+    }
+    gpt_4_supported = os.getenv("GPT_4_SUPPORTED") == "true"
+    if gpt_4_supported:
+        if token_size <= model_options["gpt-4"]:
+            return "gpt-4"
+        elif token_size <= model_options["gpt-4-32k"]:
+            return "gpt-4-32k"
+        else:
+            raise click.ClickException("🛑 The context is too large to for the Model. 😞")
+    else:
+        if token_size <= model_options["gpt-3.5-turbo"]:  # noqa: PLR5501
+            return "gpt-3.5-turbo"
+        elif token_size <= model_options["gpt-3.5-turbo-16k"]:
+            return "gpt-3.5-turbo-16k"
+        else:
+            raise click.ClickException("🛑 The context is too large to for the Model. 😞")
+
+
 if __name__ == "__main__":
     cli()