diff --git a/codeaide/logic/chat_handler.py b/codeaide/logic/chat_handler.py index 6fa588c..05beafd 100644 --- a/codeaide/logic/chat_handler.py +++ b/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): diff --git a/codeaide/utils/api_utils.py b/codeaide/utils/api_utils.py index 2a01aae..0a53e99 100644 --- a/codeaide/utils/api_utils.py +++ b/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