Update app.py

app.py

@@ -16,8 +16,56 @@ from prompts import QP_MS_TRANSCRIPTION_PROMPT, get_grading_prompt
 from supabase import create_client, Client 
 
 # ---------------- CONFIG ---------------- 
-# Create client with new SDK
-client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+# Multi-API Key Configuration for handling RESOURCE_EXHAUSTED errors
+class GeminiClientManager:
+    """Manages multiple Gemini API keys with automatic rotation on quota exhaustion."""
+    
+    def __init__(self):
+        # Load all three API keys from environment
+        self.api_keys = [
+            os.getenv("GEMINI_API_KEY_1"),
+            os.getenv("GEMINI_API_KEY_2"),
+            os.getenv("GEMINI_API_KEY_3")
+        ]
+        
+        # Filter out None values
+        self.api_keys = [key for key in self.api_keys if key]
+        
+        if not self.api_keys:
+            raise ValueError("❌ No API keys found! Please set at least GEMINI_API_KEY_1")
+        
+        print(f"✅ Loaded {len(self.api_keys)} Gemini API key(s)")
+        
+        # Current key index (0 = primary)
+        self.current_key_index = 0
+        
+        # Create clients for all keys
+        self.clients = [genai.Client(api_key=key) for key in self.api_keys]
+        
+    def get_current_client(self):
+        """Get the currently active client."""
+        return self.clients[self.current_key_index]
+    
+    def rotate_to_next_key(self):
+        """Rotate to the next available API key."""
+        if len(self.api_keys) == 1:
+            print("⚠️ Only one API key available, cannot rotate")
+            return False
+        
+        old_index = self.current_key_index
+        self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
+        print(f"🔄 Rotating from API key #{old_index + 1} to API key #{self.current_key_index + 1}")
+        return True
+    
+    def reset_to_primary(self):
+        """Reset to primary (first) API key."""
+        if self.current_key_index != 0:
+            print(f"🔙 Resetting to primary API key #1")
+            self.current_key_index = 0
+
+# Initialize the client manager
+client_manager = GeminiClientManager()
+client = client_manager.get_current_client()  # For backward compatibility
 GRID_ROWS, GRID_COLS = 20, 14
 
 # Supabase configuration
@@ -501,16 +549,18 @@ def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
 def upload_to_gemini(path, display_name=None):
     """
     Upload a file to Gemini using the NEW google-genai SDK.
+    Uses the current active API key from client_manager.
     """
     print(f"📤 Uploading {path} to Gemini...")
     try:
-        uploaded_file = client.files.upload(file=path)
+        current_client = client_manager.get_current_client()
+        uploaded_file = current_client.files.upload(file=path)
         
         # Wait for processing to complete
         print(f"⏳ Waiting for file processing: {uploaded_file.name}")
         while uploaded_file.state.name == "PROCESSING":
             time.sleep(2)
-            uploaded_file = client.files.get(name=uploaded_file.name)
+            uploaded_file = current_client.files.get(name=uploaded_file.name)
         
         if uploaded_file.state.name == "FAILED":
             raise Exception(f"File processing failed: {uploaded_file.name}")
@@ -534,6 +584,7 @@ def merge_pdfs(paths, output_path):
 def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, model_name="gemini-2.5-pro", fallback_model="gemini-2.5-flash"): 
     """ 
     Send prompt_text and optionally an uploaded file (or an image object/list) to the model using NEW SDK.
+    Automatically rotates through available API keys on RESOURCE_EXHAUSTED errors.
     Returns textual response and prints progress. 
     """ 
     contents = [prompt_text] 
@@ -558,29 +609,81 @@ def gemini_generate_content(prompt_text, file_upload_obj=None, image_obj=None, m
                 
     print("📡 Sending request to Gemini (prompt length:", len(prompt_text), "chars )") 
     
-    try:
-        response = client.models.generate_content(
-            model=model_name,
-            contents=contents
-        )
-        raw_text = response.text
-        print("📥 Received response (chars):", len(raw_text)) 
-        return raw_text
-    except Exception as e:
-        print(f"❌ Generation failed: {e}")
-        # Try fallback model
-        print(f"⚡ Trying fallback model: {fallback_model}")
+    # Try with all available API keys
+    max_attempts = len(client_manager.api_keys)
+    attempt = 0
+    
+    while attempt < max_attempts:
+        current_client = client_manager.get_current_client()
+        current_key_num = client_manager.current_key_index + 1
+        
         try:
-            response = client.models.generate_content(
-                model=fallback_model,
+            print(f"🔑 Using API key #{current_key_num} with model {model_name}")
+            response = current_client.models.generate_content(
+                model=model_name,
                 contents=contents
             )
             raw_text = response.text
-            print("📥 Received response (chars):", len(raw_text)) 
+            print(f"📥 Received response (chars): {len(raw_text)}") 
+            
+            # Success! Reset to primary key for next request
+            client_manager.reset_to_primary()
             return raw_text
-        except Exception as e2:
-            print(f"❌ Fallback also failed: {e2}")
-            raise
+            
+        except Exception as e:
+            error_str = str(e)
+            print(f"❌ Generation failed with API key #{current_key_num}: {e}")
+            
+            # Check if it's a RESOURCE_EXHAUSTED error
+            if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str:
+                print(f"⚠️ Quota exhausted for API key #{current_key_num}")
+                
+                # Try to rotate to next key
+                if client_manager.rotate_to_next_key():
+                    attempt += 1
+                    print(f"🔄 Retrying with next API key (attempt {attempt + 1}/{max_attempts})...")
+                    continue
+                else:
+                    # Only one key available, try fallback model
+                    print(f"⚡ Trying fallback model: {fallback_model}")
+                    try:
+                        response = current_client.models.generate_content(
+                            model=fallback_model,
+                            contents=contents
+                        )
+                        raw_text = response.text
+                        print(f"📥 Received response (chars): {len(raw_text)}") 
+                        client_manager.reset_to_primary()
+                        return raw_text
+                    except Exception as e2:
+                        print(f"❌ Fallback also failed: {e2}")
+                        raise Exception(f"All API keys exhausted. Error: {e2}")
+            else:
+                # Not a quota error, try fallback model with same key
+                print(f"⚡ Trying fallback model: {fallback_model}")
+                try:
+                    response = current_client.models.generate_content(
+                        model=fallback_model,
+                        contents=contents
+                    )
+                    raw_text = response.text
+                    print(f"📥 Received response (chars): {len(raw_text)}") 
+                    client_manager.reset_to_primary()
+                    return raw_text
+                except Exception as e2:
+                    print(f"❌ Fallback also failed: {e2}")
+                    # If we have more keys, try them
+                    if attempt < max_attempts - 1:
+                        client_manager.rotate_to_next_key()
+                        attempt += 1
+                        print(f"🔄 Trying next API key (attempt {attempt + 1}/{max_attempts})...")
+                        continue
+                    else:
+                        raise Exception(f"All attempts failed. Last error: {e2}")
+    
+    # If we exhausted all attempts
+    raise Exception(f"❌ All {max_attempts} API key(s) exhausted. Please check your quota or try again later.")
+
 
 # ---------------- PARSERS ---------------- 
 def extract_question_ids_from_qpms(text: str):