Update ai_service.py

ai_service.py

@@ -1,128 +1,104 @@
-# ai_service.py (Instructs AI to assume 'largest' means by magnitude)
+# ai_service.py (Implements robust API key validation and rotation)
 import json
 import os
 import random
 from datetime import datetime
 import google.generativeai as genai
+from google.api_core import exceptions as google_exceptions
 from gradio_client import Client
 
 # 從設定檔匯入金鑰和 URL
 from config import MCP_SERVER_URL
 
-# --- 1. 設定 Gemini API 金鑰 (金鑰輪替邏輯) ---
+# --- 1. [核心修改] API 金鑰健康檢查與輪替機制 ---
+
+# 從環境變數讀取您設定的兩組金鑰
 key1 = os.getenv("GEMINI_API_KEY")
 key2 = os.getenv("GEMINI_API_KEY2")
-valid_keys = []
+
+all_keys = []
 if key1 and "YOUR_GEMINI_API_KEY" not in key1:
-    valid_keys.append(key1)
+    all_keys.append(key1)
 if key2 and "YOUR_GEMINI_API_KEY" not in key2:
-    valid_keys.append(key2)
+    all_keys.append(key2)
+
+# 建立一個存放 "健康" (有效) 金鑰的列表
+healthy_keys = []
+if all_keys:
+    print(f"--- Found {len(all_keys)} Gemini API Key(s). Starting validation... ---")
+    for key in all_keys:
+        try:
+            # 嘗試用每個 key 來設定並初始化一個模型，以驗證其有效性
+            genai.configure(api_key=key)
+            # 嘗試發送一個非常簡單的測試請求
+            test_model = genai.GenerativeModel('gemini-1.5-flash')
+            test_model.generate_content("test", request_options={'timeout': 10})
+            
+            # 如果成功，將此 key 加入健康列表
+            healthy_keys.append(key)
+            print(f"--- Key ending in ...{key[-4:]} is VALID and added to the pool. ---")
+        except google_exceptions.PermissionDenied as e:
+            print(f"--- Key ending in ...{key[-4:]} is INVALID (Permission Denied). Skipping. ---")
+            print(f"    Error details: {e}")
+        except Exception as e:
+            # 捕捉其他可能的錯誤，例如超時、額度用盡等
+            print(f"--- Key ending in ...{key[-4:]} failed validation. Skipping. ---")
+            print(f"    Error details: {e}")
 
-selected_key = None
-if valid_keys:
-    selected_key = random.choice(valid_keys)
-    print(f"--- Configuring Gemini. Selected key ending in: ...{selected_key[-4:]} ---")
-    genai.configure(api_key=selected_key)
-else:
-    print("--- No valid Gemini API Key found. AI Service will be disabled. ---")
+if not healthy_keys:
+    print("--- CRITICAL: No valid Gemini API Keys found. AI Service will be disabled. ---")
 
-# --- 2. 工具函式 (Tool Functions) ---
+# --- 2. 工具函式 (Tool Functions) --- (此區塊維持不變)
 
-def call_mcp_earthquake_search(
-    start_date: str,
-    end_date: str,
-    min_magnitude: float = 4.0, #降低最小規模門檻以獲取更多資料
-    max_magnitude: float = 9.0
-) -> str:
-    """根據指定的條件（時間、規模）從遠端伺服器搜尋地震事件。"""
+def call_mcp_earthquake_search(start_date: str, end_date: str, min_magnitude: float = 4.0, max_magnitude: float = 9.0) -> str:
     try:
-        print(f"--- 正在呼叫遠端地震 MCP 伺服器 (由 Gemini 觸發) ---")
         client = Client(src=MCP_SERVER_URL)
-        result = client.predict(
-            param_0=start_date, param_1="00:00:00",
-            param_2=end_date, param_3="23:59:59",
-            param_4=21.0, param_5=26.0,
-            param_6=119.0, param_7=123.0,
-            param_8=0.0, param_9=100.0,
-            param_10=min_magnitude, param_11=max_magnitude,
-            api_name="/gradio_fetch_and_plot_data"
-        )
+        result = client.predict(param_0=start_date, param_1="00:00:00", param_2=end_date, param_3="23:59:59", param_4=21.0, param_5=26.0, param_6=119.0, param_7=123.0, param_8=0.0, param_9=100.0, param_10=min_magnitude, param_11=max_magnitude, api_name="/gradio_fetch_and_plot_data")
         data = result[0].get('data', [])
-        if not data:
-            return "查詢完成，但未找到任何符合條件的地震資料。"
+        if not data: return "查詢完成，但未找到任何符合條件的地震資料。"
         headers = result[0].get('headers', [])
         return json.dumps([dict(zip(headers, row)) for row in data], indent=2, ensure_ascii=False)
-    except Exception as e:
-        return f"工具執行失敗，錯誤訊息: {e}"
+    except Exception as e: return f"工具執行失敗，錯誤訊息: {e}"
 
 def call_mcp_pws_search() -> str:
-    """從遠端伺服器查詢最新的 PWS (Public Weather Service) 發布情形。"""
     try:
-        print(f"--- 正在呼叫遠端 PWS MCP 伺服器 (由 Gemini 觸發) ---")
         client = Client("cwadayi/MCP-pws-running")
         result = client.predict(None, api_name="/get_disaster_warnings")
         return result[0] if isinstance(result, tuple) and len(result) > 0 else str(result)
-    except Exception as e:
-        return f"工具執行失敗，錯誤訊息: {e}"
-
-# --- 3. 向 Gemini 定義工具 (Tool Declarations) ---
+    except Exception as e: return f"工具執行失敗，錯誤訊息: {e}"
 
-earthquake_search_tool_declaration = {
-    "name": "call_earthquake_search_tool",
-    "description": "從台灣中央氣象署的資料庫中搜尋地震事件。",
-    "parameters": {
-        "type": "OBJECT",
-        "properties": {
-            "start_date": {
-                "type": "STRING",
-                "description": "搜尋的開始日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題中的相對時間（例如：昨天、上個月、去年）或絕對時間（例如：2024年）來主動推斷此日期。"
-            },
-            "end_date": {
-                "type": "STRING",
-                "description": "搜尋的結束日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題中的相對時間或絕對時間來主動推斷此日期。"
-            },
-        },
-        "required": ["start_date", "end_date"]
-    }
-}
+# --- 3. 向 Gemini 定義工具 (Tool Declarations) --- (此區塊維持不變)
+# ... (此處省略與上一版相同的工具定義程式碼)
+earthquake_search_tool_declaration = { "name": "call_earthquake_search_tool", "description": "從台灣中央氣象署的資料庫中搜尋地震事件。", "parameters": { "type": "OBJECT", "properties": { "start_date": { "type": "STRING", "description": "搜尋的開始日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題中的相對時間（例如：昨天、上個月、去年）或絕對時間（例如：2024年）來主動推斷此日期。" }, "end_date": { "type": "STRING", "description": "搜尋的結束日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題中的相對時間或絕對時間來主動推斷此日期。" }, }, "required": ["start_date", "end_date"] } }
+pws_search_tool_declaration = { "name": "call_mcp_pws_search", "description": "查詢最新的 PWS (Public Weather Service) 公共天氣服務發布情形。", "parameters": { "type": "OBJECT", "properties": {} } }
+available_tools = { "call_earthquake_search_tool": call_mcp_earthquake_search, "call_mcp_pws_search": call_mcp_pws_search }
 
-pws_search_tool_declaration = {
-    "name": "call_mcp_pws_search",
-    "description": "查詢最新的 PWS (Public Weather Service) 公共天氣服務發布情形。",
-    "parameters": {"type": "OBJECT", "properties": {}}
-}
+# --- 4. 主要的 AI 文字生成函式 ---
 
-available_tools = {
-    "call_earthquake_search_tool": call_mcp_earthquake_search,
-    "call_mcp_pws_search": call_mcp_pws_search
-}
-
-# --- 4. 建立 Gemini 模型 ---
-model = None
-if selected_key:
+def generate_ai_text(user_prompt: str) -> str:
+    # [核心修改] 每次呼叫時，都從健康的金鑰池中隨機選取
+    if not healthy_keys:
+        return "🤖 AI (Gemini) 服務錯誤：沒有任何有效的 API 金鑰可供使用。"
+    
     try:
-        # [*** 核心修正 ***]
-        # 新增指令，要求 AI 在遇到「最大」等模糊詞彙時，主動以「規模」作為判斷標準
-        system_instruction = (
-            "You are a helpful AI assistant. You must answer in Traditional Chinese."
-            "You have access to tools. When the user asks for the 'largest' or 'strongest' earthquake, "
-            "you MUST assume they mean by magnitude. You should then analyze the JSON data returned by the tool, "
-            "find the single entry with the highest 'magnitude' value, and present that specific earthquake's details as the answer."
-            "Do not simply show all the data; find the maximum and answer the question directly."
-        )
-        model = genai.GenerateContent(
+        # 從健康池中隨機選擇一個 key
+        selected_key = random.choice(healthy_keys)
+        print(f"--- Handling request with key ending in: ...{selected_key[-4:]} ---")
+        
+        # 使用選定的 key 建立模型實例
+        genai.configure(api_key=selected_key)
+        model = genai.GenerativeModel(
             model_name="gemini-1.5-flash",
             tools=[earthquake_search_tool_declaration, pws_search_tool_declaration],
-            system_instruction=system_instruction
+            system_instruction=(
+                "You are a helpful AI assistant. You must answer in Traditional Chinese."
+                "When the user asks for the 'largest' or 'strongest' earthquake, "
+                "you MUST assume they mean by magnitude. You should then analyze the JSON data returned by the tool, "
+                "find the single entry with the highest 'magnitude' value, and present that specific earthquake's details as the answer."
+                "Do not simply show all the data; find the maximum and answer the question directly."
+            )
         )
-    except Exception as e:
-        print(f"建立 Gemini 模型失敗: {e}")
-
-# --- 5. 主要的 AI 文字生成函式 ---
-def generate_ai_text(user_prompt: str) -> str:
-    if not model:
-        return "🤖 AI (Gemini) 服務尚未設定 API 金鑰，或金鑰無效。"
-    try:
+        
         chat = model.start_chat()
         response = chat.send_message(user_prompt)
         for part in response.parts:
@@ -140,7 +116,7 @@ def generate_ai_text(user_prompt: str) -> str:
                 return final_response.text
         
         return response.text
+
     except Exception as e:
         print(f"與 Gemini AI 互動時發生錯誤: {e}")
         return f"🤖 AI 服務發生錯誤: {e}"
-