Update ai_service.py

ai_service.py

@@ -1,57 +1,114 @@
-# ai_service.py
-import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-from config import LLM_MODEL, LLM_MAX_NEW_TOKENS, LLM_TEMPERATURE, LLM_TOP_K
+# ai_service.py (Gemini 最終版)
+import json
+from datetime import datetime
+import google.generativeai as genai
+from gradio_client import Client
 
-_LLM = {"loaded": False, "ok": False, "err": None, "model": None, "tokenizer": None, "device": "cpu"}
+# 從設定檔匯入金鑰和 URL
+from config import GEMINI_API_KEY, MCP_SERVER_URL
 
-def _ensure_llm():
-    """在首次使用時載入 Flan-T5 模型與分詞器。"""
-    if _LLM["loaded"]:
-        return _LLM["ok"], _LLM["err"]
-    _LLM["loaded"] = True
+# --- 1. 設定 Gemini API 金鑰 (一次性設定) ---
+if GEMINI_API_KEY and "YOUR_GEMINI_API_KEY" not in GEMINI_API_KEY:
+    genai.configure(api_key=GEMINI_API_KEY)
 
+# --- 2. 工具函式 (用於地震查詢) ---
+def call_mcp_earthquake_search(
+    start_date: str,
+    end_date: str,
+    min_magnitude: float = 4.5,
+    max_magnitude: float = 8.0
+) -> str:
+    """根據指定的條件（時間、規模）從遠端伺服器搜尋地震事件。"""
     try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        # 載入 T5 專用的分詞器和模型
-        tokenizer = T5Tokenizer.from_pretrained(LLM_MODEL)
-        model = T5ForConditionalGeneration.from_pretrained(LLM_MODEL).to(device)
-        
-        _LLM.update({"ok": True, "model": model, "tokenizer": tokenizer, "device": device})
-        print(f"Flan-T5 model '{LLM_MODEL}' loaded successfully on {device}.")
-        return True, None
+        print(f"--- 正在呼叫遠端地震 MCP 伺服器 (由 Gemini 觸發) ---")
+        print(f"    查詢條件: {start_date} 到 {end_date}, 規模 {min_magnitude} 以上")
+
+        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"
+        )
+        dataframe_dict = result[0]
+        data = dataframe_dict.get('data', [])
+
+        if not data:
+            print("--- MCP 伺服器回傳：未找到符合條件的地震 ---")
+            return "查詢完成，但未找到任何符合條件的地震資料。"
+
+        headers = dataframe_dict.get('headers', [])
+        formatted_results = [dict(zip(headers, row)) for row in data]
+        print(f"--- MCP 伺服器成功回傳 {len(data)} 筆資料 ---")
+        return json.dumps(formatted_results, indent=2, ensure_ascii=False)
     except Exception as e:
-        _LLM["err"] = f"{e}"
-        _LLM["ok"] = False
-        return False, _LLM["err"]
+        print(f"呼叫 MCP 伺服器失敗: {e}")
+        return f"工具執行失敗，錯誤訊息: {e}"
 
-def generate_ai_text(user_prompt: str) -> str:
-    """使用已載入的 Flan-T5 模型生成文字回應。"""
-    ok, err = _ensure_llm()
-    if not ok:
-        return f"🤖 AI 模型無法使用。\n詳細錯誤：{err}"
+# --- 3. 向 Gemini 定義工具 ---
+earthquake_search_tool_declaration = {
+    "name": "call_earthquake_search_tool",
+    "description": "根據指定的條件（時間、地點、規模等）從台灣中央氣象署的資料庫中搜尋地震事件。預設搜尋台灣周邊地區。",
+    "parameters": {
+        "type": "OBJECT", "properties": {
+            "start_date": {"type": "STRING", "description": "搜尋的開始日期，格式為 'YYYY-MM-DD'。"},
+            "end_date": {"type": "STRING", "description": f"搜尋的結束日期，格式為 'YYYY-MM-DD'。預設為今天: {datetime.now().strftime('%Y-%m-%d')}。"},
+            "min_magnitude": {"type": "NUMBER", "description": "要搜尋的最小地震規模。預設為 4.5。"},
+            "max_magnitude": {"type": "NUMBER", "description": "要搜尋的最大地震規模。預設為 8.0。"},
+        }, "required": ["start_date", "end_date", "min_magnitude"]
+    }
+}
 
-    tokenizer = _LLM["tokenizer"]
-    model = _LLM["model"]
-    device = _LLM["device"]
+available_tools = {"call_earthquake_search_tool": call_mcp_earthquake_search}
 
-    # 為 Flan-T5 建立一個通用的問答指令
-    input_text = f"請用繁體中文回答以下問題: {user_prompt}"
+# --- 4. 建立 Gemini 模型 (Singleton 模式) ---
+model = None
+if GEMINI_API_KEY and "YOUR_GEMINI_API_KEY" not in GEMINI_API_KEY:
+    try:
+        model = genai.GenerativeModel(
+            model_name="gemini-1.5-flash",
+            tools=[earthquake_search_tool_declaration]
+        )
+    except Exception as e:
+        print(f"建立 Gemini 模型失敗: {e}")
 
+# --- 5. 主要的 AI 文字生成函式 ---
+def generate_ai_text(user_prompt: str) -> str:
+    """使用 Gemini 模型生成回應，並在需���時觸發工具呼叫。"""
+    if not model:
+        return "🤖 AI (Gemini) 服務尚未設定 API 金鑰，或金鑰無效。"
     try:
-        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
-        
-        with torch.no_grad():
-            outputs = model.generate(
-                input_ids,
-                max_new_tokens=LLM_MAX_NEW_TOKENS,
-                do_sample=True,
-                temperature=LLM_TEMPERATURE,
-                top_k=LLM_TOP_K
-            )
-        
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response.strip() or "（AI 沒有產生任何內容）"
+        print(f"--- 開始 Gemini 對話，使用者輸入: '{user_prompt}' ---")
+        chat = model.start_chat()
+        response = chat.send_message(user_prompt)
+
+        # 檢查模型是否要求呼叫工具
+        try:
+            function_call = response.candidates[0].content.parts[0].function_call
+        except (IndexError, AttributeError):
+            function_call = None
+
+        if not function_call:
+            print("--- Gemini 直接回覆文字 ---")
+            return response.text
+
+        # 處理工具呼叫
+        print(f"--- Gemini 要求呼叫工具: {function_call.name} ---")
+        tool_function = available_tools.get(function_call.name)
+        if not tool_function:
+            return f"錯誤：模型嘗試呼叫一個不存在的工具 '{function_call.name}'。"
+
+        tool_result = tool_function(**dict(function_call.args))
+        print("--- 將工具結果回傳給 Gemini ---")
+        response = chat.send_message(
+            genai.Part(function_response={"name": function_call.name, "response": {"result": tool_result}}),
+        )
+        print("--- Gemini 根據工具結果生成最終回覆 ---")
+        return response.text
+
     except Exception as e:
-        return f"AI 產生內容時發生錯誤：{e}"
+        print(f"與 Gemini AI 互動時發生錯誤: {e}")
+        return f"🤖 AI 服務發生錯誤: {e}"