Deal with quota exceeded error

codeaide/logic/chat_handler.py

@@ -7,6 +7,7 @@
     send_api_request,
     get_api_client,
     save_api_key,
+    QuotaExceededException,
 )
 from codeaide.utils.constants import (
     MAX_RETRIES,
@@ -224,18 +225,20 @@ def process_input(self, user_input):
             self.add_user_input_to_history(user_input)
 
             for attempt in range(MAX_RETRIES):
-                response = self.get_ai_response()
-                if response is None:
-                    if self.is_last_attempt(attempt):
-                        return self.create_error_response(
-                            "Failed to get a response from the AI. Please try again."
-                        )
-                    continue
+                try:
+                    response = self.get_ai_response()
+                    if response is None:
+                        if self.is_last_attempt(attempt):
+                            return self.create_error_response(
+                                "Failed to get a response from the AI. Please try again."
+                            )
+                        continue
 
-                self.cost_tracker.log_request(response)
+                    self.cost_tracker.log_request(response)
 
-                try:
                     return self.process_ai_response(response)
+                except QuotaExceededException as e:
+                    return self.create_error_response(str(e))
                 except ValueError as e:
                     self.logger.error(f"ValueError: {str(e)}\n")
                     if not self.is_last_attempt(attempt):

codeaide/utils/api_utils.py

@@ -6,6 +6,7 @@
 import hjson
 import re
 from google.generativeai.types import GenerationConfig
+from google.api_core import exceptions as google_exceptions
 
 from codeaide.utils.constants import (
     AI_PROVIDERS,
@@ -119,24 +120,29 @@ def send_api_request(api_client, conversation_history, max_tokens, model, provid
             if not response.choices:
                 return None
         elif provider.lower() == "google":
-            # Convert conversation history to the format expected by Google Gemini
-            prompt = ""
-            for message in conversation_history:
-                role = message["role"]
-                content = message["content"]
-                prompt += f"{role.capitalize()}: {content}\n\n"
-
-            # Create a GenerationConfig object
-            generation_config = GenerationConfig(
-                max_output_tokens=max_tokens,
-                temperature=0.7,  # You can adjust this as needed
-                top_p=0.95,  # You can adjust this as needed
-                top_k=40,  # You can adjust this as needed
-            )
-
-            response = api_client.generate_content(
-                contents=prompt, generation_config=generation_config
-            )
+            try:
+                prompt = ""
+                for message in conversation_history:
+                    role = message["role"]
+                    content = message["content"]
+                    prompt += f"{role.capitalize()}: {content}\n\n"
+
+                # Create a GenerationConfig object
+                generation_config = GenerationConfig(
+                    max_output_tokens=max_tokens,
+                    temperature=0.7,  # You can adjust this as needed
+                    top_p=0.95,  # You can adjust this as needed
+                    top_k=40,  # You can adjust this as needed
+                )
+
+                response = api_client.generate_content(
+                    contents=prompt, generation_config=generation_config
+                )
+            except google_exceptions.ResourceExhausted:
+                logger.error("Google API quota exceeded")
+                raise QuotaExceededException(
+                    "Your quota has been exceeded. You might need to wait briefly before trying again or try using a different model."
+                )
         else:
             raise NotImplementedError(f"API request for {provider} not implemented")
 
@@ -145,6 +151,8 @@ def send_api_request(api_client, conversation_history, max_tokens, model, provid
         return response
     except Exception as e:
         logger.error(f"Error in API request to {provider}: {str(e)}")
+        if isinstance(e, QuotaExceededException):
+            raise
         return None
 
 
@@ -240,3 +248,8 @@ def check_api_connection():
 if __name__ == "__main__":
     success, message = check_api_connection()
     logger.info(f"Connection {'successful' if success else 'failed'}: {message}")
+
+
+# Add this new exception class at the end of the file
+class QuotaExceededException(Exception):
+    pass