Update app.py

app.py

@@ -4,15 +4,19 @@ import time
 import json
 import markdown
 from flask import Flask, request, jsonify, render_template_string
-import google.generativeai as genai # 👈 CLASSIC LIBRARY (STABLE)
+import google.generativeai as genai
 
 # ==========================================
 # 👇 API KEYS SETUP 👇
 # ==========================================
 keys_string = os.environ.get("API_KEYS", "")
+# Clean up keys (remove spaces/newlines)
 API_KEYS = [k.strip() for k in keys_string.replace(',', ' ').replace('\n', ' ').split() if k.strip()]
 
-print(f"✅ Loaded {len(API_KEYS)} API Keys.") # Debug Log
+# 👇 MODELS TO TRY (Priority Order) 👇
+MODELS = ["gemini-1.5-flash", "gemini-1.5-pro", "gemini-pro"]
+
+print(f"✅ Loaded {len(API_KEYS)} API Keys.")
 
 current_key_index = 0
 app = Flask(__name__)
@@ -45,67 +49,77 @@ PERMANENT RULES:
 4. **BRIEF:** Keep explanations short and crisp.
 """
 
+def get_next_key():
+    """Rotates to the next available API Key"""
+    global current_key_index
+    current_key_index = (current_key_index + 1) % len(API_KEYS)
+    return API_KEYS[current_key_index]
+
 def generate_with_retry(prompt, history_messages=[]):
     """
-    Uses Standard Library (google-generativeai)
-    Model: gemini-1.5-flash (Proven & Stable)
+    SMART LOGIC:
+    - If 429 (Quota): Switch KEY immediately.
+    - If 404 (Not Found): Switch MODEL immediately.
     """
     global current_key_index
     
-    if not API_KEYS:
-        return "❌ Error: API Keys Missing in Secrets."
+    if not API_KEYS: return "❌ Error: API Keys Missing in Secrets."
 
-    # CONVERT HISTORY TO OLD SDK FORMAT
-    # Old SDK expects: [{'role': 'user', 'parts': ['text']}, {'role': 'model', 'parts': ['text']}]
+    # History Formatting
     recent_history = history_messages[-10:] 
     formatted_history = []
-    
     for msg in recent_history:
         role = "user" if msg["role"] == "user" else "model"
-        formatted_history.append({
-            "role": role,
-            "parts": [msg["content"]]
-        })
+        formatted_history.append({"role": role, "parts": [msg["content"]]})
 
     last_error = ""
+    
+    # We allow (Number of Keys * Number of Models) attempts total
+    max_attempts = len(API_KEYS) * len(MODELS)
 
-    # RETRY LOOP
-    for attempt in range(len(API_KEYS) * 2):
+    for attempt in range(max_attempts):
         key = API_KEYS[current_key_index]
         
-        try:
-            # 1. Configure Key
-            genai.configure(api_key=key)
-            
-            # 2. Initialize Model (Classic Way)
-            model = genai.GenerativeModel(
-                model_name="gemini-1.5-flash",
-                system_instruction=SYSTEM_INSTRUCTION
-            )
-            
-            # 3. Start Chat & Send Message
-            chat = model.start_chat(history=formatted_history)
-            response = chat.send_message(prompt)
-            
-            return response.text
-            
-        except Exception as e:
-            error_msg = str(e)
-            print(f"⚠️ Key #{current_key_index+1} Failed: {error_msg}")
-            
-            if "404" in error_msg:
-                # If 404 happens here, it's purely a Model Name issue, but gemini-1.5-flash IS correct.
-                # So it usually means key access issue.
-                pass 
+        # Try current key with all models priority
+        for model_name in MODELS:
+            try:
+                genai.configure(api_key=key)
+                model = genai.GenerativeModel(
+                    model_name=model_name,
+                    system_instruction=SYSTEM_INSTRUCTION
+                )
+                
+                chat = model.start_chat(history=formatted_history)
+                response = chat.send_message(prompt)
+                
+                # If success, return text
+                return response.text
 
-            last_error = error_msg
-            current_key_index = (current_key_index + 1) % len(API_KEYS)
-            time.sleep(1)
-            continue
+            except Exception as e:
+                error_str = str(e).lower()
+                
+                # CASE 1: QUOTA EXCEEDED (429) -> SWITCH KEY
+                if "429" in error_str or "exhausted" in error_str:
+                    print(f"⚠️ Key #{current_key_index+1} Quota Full ({model_name}). Switching Key...")
+                    get_next_key() # Rotate Key
+                    time.sleep(1)  # Cool down
+                    break # Break inner model loop, try new key
+                
+                # CASE 2: MODEL NOT FOUND (404) -> TRY NEXT MODEL (SAME KEY)
+                elif "404" in error_str or "not found" in error_str:
+                    print(f"⚠️ {model_name} not found. Trying next model...")
+                    continue # Try next model in the list
+                
+                # CASE 3: OTHER ERRORS -> SWITCH KEY
+                else:
+                    print(f"⚠️ Error on {model_name}: {error_str}")
+                    last_error = error_str
+                    get_next_key()
+                    break 
 
-    return f"⚠️ **Connection Error:** {last_error}<br>Try restarting the Space."
+    return f"⚠️ **System Busy:** All keys/models failed. (Last Error: {last_error})"
 
-# --- UI TEMPLATE (KEEPING THE WORKING UI) ---
+# --- UI TEMPLATE (STABLE - NO CHANGES) ---
 HTML_TEMPLATE = """
 <!DOCTYPE html>
 <html lang="en">
@@ -314,7 +328,7 @@ def chat():
     try:
         previous_messages = user_db[user][cid]['messages'][:-1]
         
-        # Call the classic generation function
+        # Call the SMART generation function
         raw_response = generate_with_retry(msg, previous_messages)
         
         formatted_html = markdown.markdown(raw_response, extensions=['fenced_code', 'codehilite'])