Update backend/agents/puzzle_agent.py

backend/agents/puzzle_agent.py

@@ -44,20 +44,41 @@ MAZE_TOOLBOX_XML = (
 )
 
 
-def _get_best_flash_model() -> str:
+def _get_best_flash_model(api_key: Optional[str] = None) -> str:
+    """
+    Auto-detects the best available Gemini model.
+    Dynamically parses the list of supported models and selects the highest-version Flash model.
+    Falls back to a default if query fails or no flash models are found.
+    """
+    key = api_key or os.getenv("GEMINI_API_KEY")
+    if not key:
+        return "gemini-2.0-flash"
     try:
-        available_models = []
-        for m in genai.list_models():
-            if 'generateContent' in m.supported_generation_methods and 'flash' in m.name.lower():
-                available_models.append(m.name)
-        if available_models:
-            available_models.sort(reverse=True)
-            return available_models[0].replace('models/', '')
-        for m in genai.list_models():
-            if 'generateContent' in m.supported_generation_methods and 'gemini' in m.name.lower():
-                return m.name.replace('models/', '')
+        import google.generativeai as genai
+        import re
+        genai.configure(api_key=key)
+        
+        # Get all available models for this key
+        models = [m.name for m in genai.list_models()]
+        
+        # Filter for flash models
+        flash_candidates = [m for m in models if "flash" in m.lower()]
+        
+        if flash_candidates:
+            # Sort by version number extracted from the name
+            def _extract_version(name: str) -> float:
+                match = re.search(r'(\d+\.\d+)', name)
+                return float(match.group(1)) if match else 0.0
+                
+            flash_candidates.sort(key=_extract_version, reverse=True)
+            return flash_candidates[0].replace("models/", "")
+            
+        # If no flash models found, fall back to standard candidates list
+        for candidate in ["models/gemini-2.0-flash", "models/gemini-1.5-flash", "models/gemini-1.0-pro"]:
+            if candidate in models:
+                return candidate.replace("models/", "")
     except Exception as e:
-        print(f"Warning: Could not list models ({e}). Using hardcoded default.")
+        print(f"[PuzzleAgent] Error listing models: {e}. Defaulting to gemini-2.0-flash.")
     return "gemini-2.0-flash"
 
 
@@ -108,7 +129,7 @@ class PuzzleAgent:
             self.model = None
         else:
             genai.configure(api_key=self.api_key)
-            best_model = _get_best_flash_model()
+            best_model = _get_best_flash_model(self.api_key)
             print(f"[PuzzleAgent] Selected Model: {best_model}")
             self.model = genai.GenerativeModel(best_model)