Update ai_service.py

ai_service.py

@@ -1,12 +1,9 @@
-# ai_service.py (Definitive fix for the ImportError)
+# ai_service.py (已整合 PWS 查詢工具)
 import json
 from datetime import datetime
 import google.generativeai as genai
 from gradio_client import Client
 
-# [修正] 移除 'Part' 的 import，因為它導致了錯誤
-# from google.generativeai.types import Part
-
 # 從設定檔匯入金鑰和 URL
 from config import GEMINI_API_KEY, MCP_SERVER_URL
 
@@ -14,7 +11,9 @@ from config import GEMINI_API_KEY, MCP_SERVER_URL
 if GEMINI_API_KEY and "YOUR_GEMINI_API_KEY" not in GEMINI_API_KEY:
     genai.configure(api_key=GEMINI_API_KEY)
 
-# --- 2. 工具函式 (用於地震查詢) ---
+# --- 2. 工具函式 (Tool Functions) ---
+
+# [既有] 地震查詢工具函式
 def call_mcp_earthquake_search(
     start_date: str,
     end_date: str,
@@ -51,21 +50,59 @@ def call_mcp_earthquake_search(
         print(f"呼叫 MCP 伺服器失敗: {e}")
         return f"工具執行失敗，錯誤訊息: {e}"
 
-# --- 3. 向 Gemini 定義工具 ---
+# [新增] PWS 查詢工具函式
+def call_mcp_pws_search() -> str:
+    """從遠端伺服器查詢最新的 PWS (Public Weather Service) 發布情形。"""
+    try:
+        print(f"--- 正在呼叫遠端 PWS MCP 伺服器 (由 Gemini 觸發) ---")
+        # 連線到使用者指定的 Gradio 應用程式
+        client = Client("cwadayi/MCP-pws")
+        # 假設 Gradio API 不需要參數，直接呼叫 predict
+        result = client.predict(api_name="/predict")
+        
+        # Gradio client 回傳的結果通常是元組 (tuple)，取出第一個元素
+        if isinstance(result, tuple) and len(result) > 0:
+            report = result[0]
+        else:
+            report = str(result)
+
+        print(f"--- PWS MCP 伺服器成功回傳 ---")
+        return report
+    except Exception as e:
+        print(f"呼叫 PWS MCP 伺服器失敗: {e}")
+        return f"工具執行失敗，錯誤訊息: {e}"
+
+# --- 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年』推斷出 '2024-01-01'。"},
-            "end_date": {"type": "STRING", "description": "搜尋的結束日期 (格式 'YYYY-MM-DD')。模型應根據使用者問題推斷此日期，例如從『昨天』或『2024年』推斷出 '2024-12-31'。"},
-            "min_magnitude": {"type": "NUMBER", "description": "要搜尋的最小地震規模。如果使用者未指定，請使用預設值 4.5。"},
-            "max_magnitude": {"type": "NUMBER", "description": "要搜尋的最大地震規模。預設為 8.0。"},
+            "start_date": {"type": "STRING", "description": "搜尋的開始日期 (格式 'YYYY-MM-DD')。"},
+            "end_date": {"type": "STRING", "description": "搜尋的結束日期 (格式 'YYYY-MM-DD')。"},
+            "min_magnitude": {"type": "NUMBER", "description": "要搜尋的最小地震規模。"},
+            "max_magnitude": {"type": "NUMBER", "description": "要搜尋的最大地震規模。"},
         }, "required": ["start_date", "end_date"]
     }
 }
 
-available_tools = {"call_earthquake_search_tool": call_mcp_earthquake_search}
+# [新增] PWS 查詢工具定義
+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
+}
 
 # --- 4. 建立 Gemini 模型 ---
 model = None
@@ -73,20 +110,19 @@ if GEMINI_API_KEY and "YOUR_GEMINI_API_KEY" not in GEMINI_API_KEY:
     try:
         system_instruction = (
             "You are a helpful AI assistant. You must answer in Traditional Chinese."
-            "You have access to tools. When a tool returns data in JSON format, "
-            "you must analyze the JSON data to fully answer the user's question. "
-            "For example, if the user asks for the largest earthquake, use the search tool for the relevant date range "
-            "and then find the entry with the highest magnitude from the JSON results before answering."
+            "You have access to tools. When a tool returns data, "
+            "you must analyze the data to fully answer the user's question."
         )
+        # [修改] 初始化模型時，提供所有工具的定義
         model = genai.GenerativeModel(
             model_name="gemini-1.5-flash",
-            tools=[earthquake_search_tool_declaration],
+            tools=[earthquake_search_tool_declaration, pws_search_tool_declaration],
             system_instruction=system_instruction
         )
     except Exception as e:
         print(f"建立 Gemini 模型失敗: {e}")
 
-# --- 5. 主要的 AI 文字生成函式 ---
+# --- 5. 主要的 AI 文字生成函式 (此部分維持不變) ---
 def generate_ai_text(user_prompt: str) -> str:
     if not model:
         return "🤖 AI (Gemini) 服務尚未設定 API 金鑰，或金鑰無效。"
@@ -110,7 +146,6 @@ def generate_ai_text(user_prompt: str) -> str:
         tool_result = tool_function(**dict(function_call.args))
         print("--- 將工具結果回傳給 Gemini ---")
         
-        # [修正] 直接傳送包含 function_response 的字典，不再使用 Part 類別
         response = chat.send_message(
             {"function_response": {"name": function_call.name, "response": {"result": tool_result}}}
         )
@@ -119,5 +154,4 @@ def generate_ai_text(user_prompt: str) -> str:
         return response.text
     except Exception as e:
         print(f"與 Gemini AI 互動時發生錯誤: {e}")
-        return f"🤖 AI 服務發生錯誤: {e}"
-
+        return f"🤖 AI 服務發生錯誤: {e}"