Fixed api key prompting logic

codeaide/logic/chat_handler.py

@@ -48,25 +48,21 @@ def __init__(self):
             self.current_model
         ]["max_tokens"]
 
+        self.api_key_valid, self.api_key_message = self.check_api_key()
+
     def check_api_key(self):
         """
         Check if the API key is set and valid.
 
-        Args:
-            None
-
         Returns:
             tuple: A tuple containing a boolean indicating if the API key is valid and a message.
         """
-        if self.api_client is None:
-            self.api_client = get_api_client(self.current_provider, self.current_model)
+        self.api_client = get_api_client(self.current_provider, self.current_model)
+        self.api_key_set = self.api_client is not None
 
-        if self.api_client:
-            self.api_key_set = True
-            return True, None
-        else:
-            self.api_key_set = False
+        if not self.api_key_set:
             return False, self.get_api_key_instructions(self.current_provider)
+        return True, None
 
     def get_api_key_instructions(self, provider):
         """
@@ -91,6 +87,18 @@ def get_api_key_instructions(self, provider):
                 "This file is already in .gitignore, so it won't be shared if you push your code to a repository.\n\n"
                 "Please paste your Anthropic API key now:"
             )
+        elif provider == "openai":
+            return (
+                "It looks like you haven't set up your OpenAI API key yet. "
+                "Here's how to get started:\n\n"
+                "1. Go to https://platform.openai.com/api-keys and sign in to your OpenAI account or create an account if you don't have one.\n"
+                "2. Generate a new API key.\n"
+                "3. Add some funds to your account to cover the cost of using the API (start with as little as $1).\n"
+                "4. Copy the API key and paste it here.\n\n"
+                "Once you've pasted your API key, I'll save it securely in a .env file in the root of your project. "
+                "This file is already in .gitignore, so it won't be shared if you push your code to a repository.\n\n"
+                "Please paste your OpenAI API key now:"
+            )
         else:
             return f"Please enter your API key for {provider.capitalize()}:"
 
@@ -120,37 +128,31 @@ def validate_api_key(self, api_key):
 
     def handle_api_key_input(self, api_key):
         """
-        Handle the input of the API key.
+        Handle the API key input from the user.
 
         Args:
             api_key (str): The API key entered by the user.
 
         Returns:
-            tuple: A tuple containing a boolean indicating if the API key was saved successfully and a message.
-        """
-        cleaned_key = api_key.strip().strip("'\"")  # Remove quotes and whitespace
-        is_valid, error_message = self.validate_api_key(cleaned_key)
-        if is_valid:
-            if save_api_key(self.current_provider, cleaned_key):
-                # Try to get a new API client with the new key
-                self.api_client = get_api_client(
-                    self.current_provider, self.current_model
+            tuple: A tuple containing a boolean indicating success, a message, and a boolean indicating if waiting for API key.
+        """
+        if save_api_key(self.current_provider, api_key):
+            self.api_client = get_api_client(self.current_provider, self.current_model)
+            self.api_key_set = self.api_client is not None
+            if self.api_key_set:
+                return (
+                    True,
+                    "Great! Your API key has been saved. What would you like to work on?",
+                    False,
                 )
-                if self.api_client:
-                    self.api_key_set = True
-                    return True, "API key saved and verified successfully."
-                else:
-                    return (
-                        False,
-                        "API key saved, but verification failed. Please check your key and try again.",
-                    )
             else:
                 return (
                     False,
-                    "Failed to save the API key. Please check your permissions and try again.",
+                    "Failed to initialize API client with the provided key.",
+                    True,
                 )
         else:
-            return False, f"Invalid API key format: {error_message}. Please try again."
+            return False, "Failed to save the API key.", True
 
     def process_input(self, user_input):
         """
@@ -162,8 +164,9 @@ def process_input(self, user_input):
         Returns:
             dict: A response dictionary containing the type and content of the response.
         """
+        print(f"Processing input: {user_input}")
         try:
-            if not self.check_and_set_api_key():
+            if not self.api_key_set:
                 return {
                     "type": "api_key_required",
                     "message": self.get_api_key_instructions(self.current_provider),
@@ -200,19 +203,6 @@ def process_input(self, user_input):
         except Exception as e:
             return self.handle_unexpected_error(e)
 
-    def check_and_set_api_key(self):
-        """
-        Check if the API key is set and valid.
-
-        Args:
-            None
-
-        Returns:
-            bool: True if the API key is valid, False otherwise.
-        """
-        api_key_valid, _ = self.check_api_key()
-        return api_key_valid
-
     def add_user_input_to_history(self, user_input):
         """
         Add user input to the conversation history with version information.
@@ -399,7 +389,7 @@ def add_error_prompt_to_history(self, error_message):
         Returns:
             None
         """
-        error_prompt = f"\n\nThere was an error in your last response: {error_message}. Please ensure you're using proper JSON formatting to avoid this error and others like it."
+        error_prompt = f"\n\nThere was an error in your last response: {error_message}. Please ensure you're using proper JSON formatting to avoid this error and others like it. Please don't apologize for the error because it will be hidden from the end user."
         self.conversation_history[-1]["content"] += error_prompt
 
     def handle_unexpected_error(self, e):
@@ -489,19 +479,24 @@ def is_task_in_progress(self):
     def set_model(self, provider, model):
         if provider not in AI_PROVIDERS:
             print(f"Invalid provider: {provider}")
-            return False
+            return False, None
         if model not in AI_PROVIDERS[provider]["models"]:
             print(f"Invalid model {model} for provider {provider}")
-            return False
+            return False, None
 
         self.current_provider = provider
         self.current_model = model
         self.max_tokens = AI_PROVIDERS[self.current_provider]["models"][
             self.current_model
         ]["max_tokens"]
-        self.api_client = get_api_client(self.current_provider, self.current_model)
-        self.api_key_set = self.api_client is not None
-        return self.api_key_set
+
+        # Check API key when setting a new model
+        api_key_valid, message = self.check_api_key()
+        if not api_key_valid:
+            return False, message
+
+        print(f"Model {model} for provider {provider} set successfully.")
+        return True, None
 
     def clear_conversation_history(self):
         self.conversation_history = []

codeaide/ui/chat_window.py

@@ -1,7 +1,7 @@
 import signal
 import sys
 import traceback
-
+import time
 from PyQt5.QtCore import Qt, QTimer
 from PyQt5.QtGui import QColor
 from PyQt5.QtWidgets import (
@@ -48,8 +48,13 @@ def __init__(self, chat_handler):
         self.code_popup = None
         self.waiting_for_api_key = False
         self.setup_ui()
-        self.add_to_chat("AI", INITIAL_MESSAGE)
-        self.check_api_key()
+
+        # Check API key status
+        if not self.chat_handler.api_key_valid:
+            self.waiting_for_api_key = True
+            self.add_to_chat("AI", self.chat_handler.api_key_message)
+        else:
+            self.add_to_chat("AI", INITIAL_MESSAGE)
 
         self.input_text.setTextColor(QColor(CHAT_WINDOW_FG))
 
@@ -143,14 +148,35 @@ def eventFilter(self, obj, event):
 
     def on_submit(self):
         user_input = self.input_text.toPlainText().strip()
-        if user_input:
+        if not user_input:
+            return
+
+        # Clear the input field immediately
+        self.input_text.clear()
+
+        if self.waiting_for_api_key:
+            (
+                success,
+                message,
+                self.waiting_for_api_key,
+            ) = self.chat_handler.handle_api_key_input(user_input)
+            self.remove_thinking_messages()
+            self.add_to_chat("AI", message)
+            if success:
+                self.enable_ui_elements()
+        else:
+            # Immediately display user input and "Thinking..." message
             self.add_to_chat("User", user_input)
-            self.input_text.clear()
-            self.display_thinking()
             self.disable_ui_elements()
-            QTimer.singleShot(100, lambda: self.process_input(user_input))
-        else:
-            print("ChatWindow: Empty input, not submitting")
+            self.add_to_chat("AI", "Thinking... 🤔")
+
+            # Use QTimer to process the input after the UI has updated
+            QTimer.singleShot(100, lambda: self.call_process_input_async(user_input))
+
+    def call_process_input_async(self, user_input):
+        # Process the input
+        response = self.chat_handler.process_input(user_input)
+        self.handle_response(response)
 
     def on_modify(self):
         self.input_text.ensureCursorVisible()
@@ -166,42 +192,6 @@ def add_to_chat(self, sender, message):
     def display_thinking(self):
         self.add_to_chat("AI", "Thinking... 🤔")
 
-    def process_input(self, user_input):
-        try:
-            if self.waiting_for_api_key:
-                self.handle_api_key_input(user_input)
-            else:
-                response = self.chat_handler.process_input(user_input)
-                self.handle_response(response)
-        except Exception as e:
-            error_message = f"An unexpected error occurred: {str(e)}. Please check the console window for the full traceback."
-            self.add_to_chat("AI", error_message)
-            print("Unexpected error in ChatWindow process_input:", file=sys.stderr)
-            traceback.print_exc()
-        finally:
-            self.enable_ui_elements()
-
-    def check_api_key(self):
-        api_key_valid, message = self.chat_handler.check_api_key()
-        if not api_key_valid:
-            self.add_to_chat("AI", message)
-            self.waiting_for_api_key = True
-        else:
-            self.waiting_for_api_key = False
-
-    def handle_api_key_input(self, api_key):
-        success, message = self.chat_handler.handle_api_key_input(api_key)
-        self.remove_thinking_messages()
-        if success:
-            self.waiting_for_api_key = False
-            self.add_to_chat(
-                "AI",
-                "Great! Your API key has been saved. What would you like to work on?",
-            )
-        else:
-            self.add_to_chat("AI", message)
-        self.enable_ui_elements()
-
     def handle_response(self, response):
         self.enable_ui_elements()
         self.remove_thinking_messages()
@@ -313,12 +303,17 @@ def update_chat_handler(self):
             return
 
         current_version = self.chat_handler.get_latest_version()
-        success = self.chat_handler.set_model(provider, model)
+        success, message = self.chat_handler.set_model(provider, model)
+
         if not success:
-            self.add_to_chat(
-                "System",
-                f"Failed to set model {model} for provider {provider}. Please check your API key.",
-            )
+            if message:  # This indicates that an API key is required
+                self.waiting_for_api_key = True
+                self.add_to_chat("AI", message)
+            else:
+                self.add_to_chat(
+                    "System",
+                    f"Failed to set model {model} for provider {provider}. Please check your API key.",
+                )
             return
 
         self.chat_handler.clear_conversation_history()
@@ -338,9 +333,3 @@ def update_chat_handler(self):
 {'='*50}
 """,
         )
-
-        self.check_api_key()
-
-    def increment_version(self, version):
-        major, minor = map(int, version.split("."))
-        return f"{major}.{minor + 1}"

codeaide/utils/api_utils.py

@@ -25,6 +25,13 @@ def __init__(self, service):
 
 def get_api_client(provider=DEFAULT_PROVIDER, model=DEFAULT_MODEL):
     try:
+        root_dir = os.path.dirname(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        )
+
+        # Use AutoConfig to automatically find and load the .env file in the project root
+        config = AutoConfig(search_path=root_dir)
+
         api_key_name = AI_PROVIDERS[provider]["api_key_name"]
         api_key = config(api_key_name, default=None)
         print(f"Attempting to get API key for {provider} with key name: {api_key_name}")

codeaide/utils/constants.py

@@ -3,9 +3,9 @@
     "anthropic": {
         "api_key_name": "ANTHROPIC_API_KEY",
         "models": {
-            "claude-3-opus-20240229": {"max_tokens": 4096},
             "claude-3-5-sonnet-20240620": {"max_tokens": 8192},
             "claude-3-haiku-20240307": {"max_tokens": 4096},
+            "claude-3-opus-20240229": {"max_tokens": 4096},
         },
     },
     "openai": {