Add personality selection to AICodeBot setup

This commit introduces a new feature to the AICodeBot setup process. Now, users can select a personality for their AI assistant from a predefined list. The chosen personality is stored in the configuration file and used to generate the AI's responses. This change also includes a refactoring of the personality prompts, moving them into a dictionary for easier access and management. Ain't no half-steppin' here, we're making this bot a real cool cat.

Author: TechNickAI

aicodebot/cli.py

@@ -8,7 +8,7 @@
     logger,
     read_config,
 )
-from aicodebot.prompts import generate_files_context, get_prompt
+from aicodebot.prompts import PERSONALITIES, generate_files_context, get_prompt
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
@@ -404,8 +404,6 @@ def setup_config(openai_api_key=None, gpt_4_supported=None):
         logger.info("Using OPENAI_API_KEY environment variable")
         openai.api_key = os.getenv("OPENAI_API_KEY")
 
-    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")
 
@@ -421,7 +419,6 @@ def setup_config(openai_api_key=None, gpt_4_supported=None):
 
         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
     if gpt_4_supported is None:
         try:
@@ -436,9 +433,22 @@ def setup_config(openai_api_key=None, gpt_4_supported=None):
         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}
+    # Pull the choices from the name from each of the PERSONALITIES
+    personality_choices = "\nHow would you like your AI to act? You can choose from the following personalities:\n"
+    for key, personality in PERSONALITIES.items():
+        personality_choices += f"\t{key} - {personality.description}\n"
+    console.print(personality_choices)
+
+    personality = click.prompt(
+        "Please choose a personality", type=click.Choice(PERSONALITIES.keys(), case_sensitive=False)
+    )
 
-    logger.info(f"openai.api_key 3 {openai.api_key}")
+    config_data = {
+        "config_version": 1,
+        "OPENAI_API_KEY": openai.api_key,
+        "gpt_4_supported": gpt_4_supported,
+        "personality": personality,
+    }
 
     with Path.open(config_file, "w") as f:
         yaml.dump(config_data, f)

aicodebot/prompts.py

@@ -1,6 +1,7 @@
-from aicodebot.helpers import get_token_length, logger
+from aicodebot.helpers import get_token_length, logger, read_config
 from langchain import PromptTemplate
 from pathlib import Path
+from types import SimpleNamespace
 import os
 
 # ---------------------------------------------------------------------------- #
@@ -42,23 +43,28 @@
 Speak like Morpheus.
 """
 
+PERSONALITIES = {
+    "HER": SimpleNamespace(name="Her", prompt=HER, description="The AI character from the movie Her"),
+    "JULES": SimpleNamespace(name="Jules", prompt=JULES, description="Samuel L. Jackson's character from Pulp Fiction"),
+    "SHERLOCK": SimpleNamespace(name="Sherlock", prompt=SHERLOCK, description="Sherlock Holmes"),
+    "THE_DUDE": SimpleNamespace(name="The Dude", prompt=THE_DUDE, description="The Dude from The Big Lebowski"),
+    "MORPHEUS": SimpleNamespace(name="Morpheus", prompt=MORPHEUS, description="Morpheus from The Matrix"),
+}
+
 
 def get_personality_prompt():
     """Generates a prompt for the sidekick personality."""
-    personality_map = {
-        "HER": HER,
-        "JULES": JULES,
-        "SHERLOCK": SHERLOCK,
-        "THE_DUDE": THE_DUDE,
-        "MORPHEUS": MORPHEUS,
-    }
+    if os.getenv("AICODEBOT_PERSONALITY"):
+        personality = os.getenv("AICODEBOT_PERSONALITY")
+    else:
+        config = read_config()
+        personality = config["personality"]
 
-    personality = os.getenv("AICODEBOT_PERSONALITY", "HER")
-    try:
-        logger.debug(f"Using personality {personality}")
-        return personality_map[personality]
-    except KeyError as e:
-        raise ValueError(f"Personality {personality} not found") from e
+    if personality not in PERSONALITIES:
+        raise ValueError(f"Personality {personality} not found")
+
+    logger.debug(f"Using personality {personality}")
+    return PERSONALITIES[personality].prompt
 
 
 # ---------------------------------------------------------------------------- #