Final changes and update

modules/ai.py

@@ -1,5 +1,4 @@
 import os
-import sys
 import json
 import time
 import warnings
@@ -77,7 +76,7 @@ def select_best_models(force_refresh=False):
 
     try:
         utils.log("SYSTEM", "Refreshing AI model list from API...")
-        models = [m.name for m in genai.list_models() if 'generateContent' in m.supported_generation_methods]
+        models = [m.name for m in genai.list_models() if 'generateContent' in m.supported_generation_methods and 'gemini' in m.name.lower()]
         
         bootstrapper = "models/gemini-1.5-flash"
         if bootstrapper not in models:
@@ -88,8 +87,12 @@ def select_best_models(force_refresh=False):
         model = genai.GenerativeModel(bootstrapper)
         prompt = f"Analyze this list of available Google Gemini models:\n{json.dumps(models)}\n\nSelect the best model for each of these three roles based on these criteria:\n- Most recent version with best features and ability.\n- Beta versions are okay, but avoid 'experimental' if a stable beta/prod version exists.\n- Consider quota efficiency (Flash is cheaper/faster, Pro is smarter).\n\nROLES:\n1. LOGIC: For complex reasoning, JSON structuring, and plot planning.\n2. WRITER: For creative fiction writing, prose generation, and speed.\n3. ARTIST: For generating visual art prompts and design instructions.\n\nAlso provide a 'ranking' list of ALL models analyzed, ordered from best/most useful to worst/least useful, with a short reason.\n\nReturn JSON: {{ 'logic': {{ 'model': 'model_name', 'reason': 'reasoning' }}, 'writer': {{ 'model': 'model_name', 'reason': 'reasoning' }}, 'artist': {{ 'model': 'model_name', 'reason': 'reasoning' }}, 'ranking': [ {{ 'model': 'model_name', 'reason': 'reasoning' }} ] }}"
         
-        response = model.generate_content(prompt)
-        selection = json.loads(utils.clean_json(response.text))
+        try:
+            response = model.generate_content(prompt)
+            selection = json.loads(utils.clean_json(response.text))
+        except Exception as e:
+            utils.log("SYSTEM", f"Model selection generation failed (Safety/Format): {e}")
+            raise e
         
         if not os.path.exists(config.DATA_DIR): os.makedirs(config.DATA_DIR)
         with open(cache_path, 'w') as f: