Update ai_service.py

ai_service.py

@@ -1,15 +1,36 @@
-# ai_service.py (Updated PWS tool to use the new MCP-pws-running space)
+# ai_service.py (Implements API Key rotation)
 import json
+import os
+import random
 from datetime import datetime
 import google.generativeai as genai
 from gradio_client import Client
 
 # 從設定檔匯入金鑰和 URL
-from config import GEMINI_API_KEY, MCP_SERVER_URL
+from config import MCP_SERVER_URL #不再直接從config匯入GEMINI_API_KEY
 
-# --- 1. 設定 Gemini API 金鑰 (一次性設定) ---
-if GEMINI_API_KEY and "YOUR_GEMINI_API_KEY" not in GEMINI_API_KEY:
-    genai.configure(api_key=GEMINI_API_KEY)
+# --- 1. 設定 Gemini API 金鑰 (金鑰輪替邏輯) ---
+
+# [核心修改] 讀取並輪流使用兩組 API 金鑰
+# 從環境變數讀取您設定的兩組金鑰
+key1 = os.getenv("GEMINI_API_KEY")
+key2 = os.getenv("GEMINI_API_KEY2")
+
+# 建立一個包含所有有效金鑰的列表
+valid_keys = []
+if key1 and "YOUR_GEMINI_API_KEY" not in key1:
+    valid_keys.append(key1)
+if key2 and "YOUR_GEMINI_API_KEY" not in key2:
+    valid_keys.append(key2)
+
+# 如果列表中有可用的金鑰，就隨機選取一個來設定
+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. ---")
 
 # --- 2. 工具函式 (Tool Functions) ---
 
@@ -44,13 +65,8 @@ def call_mcp_pws_search() -> str:
     """從遠端伺服器查詢最新的 PWS (Public Weather Service) 發布情形。"""
     try:
         print(f"--- 正在呼叫遠端 PWS MCP 伺服器 (由 Gemini 觸發) ---")
-        
-        # [*** 核心修正 ***]
-        # 將 PWS 工具的 Hugging Face Space 位址更新為新的 running 版本
         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}"
@@ -58,25 +74,17 @@ def call_mcp_pws_search() -> str:
 # --- 3. 向 Gemini 定義工具 (Tool Declarations) ---
 
 earthquake_search_tool_declaration = {
-    "name": "call_earthquake_search_tool",
-    "description": "從台灣中央氣象署的資料庫中搜尋地震事件。",
-    "parameters": {
-        "type": "OBJECT",
-        "properties": {
+    "name": "call_earthquake_search_tool", "description": "從台灣中央氣象署的資料庫中搜尋地震事件。",
+    "parameters": { "type": "OBJECT", "properties": {
             "start_date": {"type": "STRING", "description": "搜尋的開始日期 (格式 'YYYY-MM-DD')。"},
             "end_date": {"type": "STRING", "description": "搜尋的結束日期 (格式 'YYYY-MM-DD')。"},
-        },
-        "required": ["start_date", "end_date"]
+        }, "required": ["start_date", "end_date"]
     }
 }
 
 pws_search_tool_declaration = {
-    "name": "call_mcp_pws_search",
-    "description": "查詢最新的 PWS (Public Weather Service) 公共天氣服務發布情形。",
-    "parameters": {
-        "type": "OBJECT",
-        "properties": {}
-    }
+    "name": "call_mcp_pws_search", "description": "查詢最新的 PWS (Public Weather Service) 公共天氣服務發布情形。",
+    "parameters": {"type": "OBJECT", "properties": {}}
 }
 
 available_tools = {
@@ -86,7 +94,8 @@ available_tools = {
 
 # --- 4. 建立 Gemini 模型 ---
 model = None
-if GEMINI_API_KEY and "YOUR_GEMINI_API_KEY" not in GEMINI_API_KEY:
+# [核心修改] 只有在成功選取到key之後才建立模型
+if selected_key:
     try:
         system_instruction = (
             "You are a helpful AI assistant. You must answer in Traditional Chinese."
@@ -110,7 +119,6 @@ def generate_ai_text(user_prompt: str) -> str:
     try:
         chat = model.start_chat()
         response = chat.send_message(user_prompt)
-        
         for part in response.parts:
             if part.function_call:
                 function_call = part.function_call
@@ -126,7 +134,6 @@ 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}"