Update main.py

main.py

@@ -1,38 +1,50 @@
+import os
+import io
+import tempfile
+from collections import defaultdict
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi import FastAPI, Request,  Header, BackgroundTasks, HTTPException, status
+from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException
 from fastapi.staticfiles import StaticFiles
 from google import genai
+from google.genai import types
 from linebot import LineBotApi, WebhookHandler
 from linebot.exceptions import InvalidSignatureError
-from linebot.models import MessageEvent, TextMessage, TextSendMessage, ImageSendMessage, AudioMessage, ImageMessage
-import json, os
-import io
+from linebot.models import (
+    MessageEvent,
+    TextMessage,
+    TextSendMessage,
+    ImageSendMessage,
+    ImageMessage,
+)
 import PIL.Image
-from Image_text_generation import Image_text_Generator
-from Uploading_images_file import get_image_url, store_user_message, analyze_with_gemini, get_previous_message
+import uvicorn
+
+# LangChain 相關匯入
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import tool
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.agents import create_tool_calling_agent, AgentExecutor
 
-#==========================
-#  API 金鑰
-#==========================
 
-# 設定 Google AI API 金鑰
-client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
+# ==========================
+#  環境設定與工具函式
+# ==========================
 
-# 設定生成文字的參數
-generation_config = genai.types.GenerateContentConfig(max_output_tokens=256, temperature=0.5, top_p=0.5, top_k=16)
+# 設置 Google AI API 金鑰
+os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
 
-# 設定 Line Bot 的 API 金鑰和秘密金鑰
+# 設置 Line Bot 的 API 金鑰和秘密金鑰
 line_bot_api = LineBotApi(os.environ["CHANNEL_ACCESS_TOKEN"])
 line_handler = WebhookHandler(os.environ["CHANNEL_SECRET"])
 
-# 設定是否正在與使用者交談
-working_status = os.getenv("DEFALUT_TALKING", default = "true").lower() == "true"
+# 使用字典模擬用戶訊息歷史存儲
+user_message_history = defaultdict(list)
 
 # 建立 FastAPI 應用程式
 app = FastAPI()
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# 設定 CORS，允許跨域請求
+# 設定 CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -41,141 +53,208 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 處理根路徑請求
+def get_image_url_from_line(message_id):
+    """
+    從 Line 訊息 ID 獲取圖片內容並儲存到暫存檔案。
+    """
+    try:
+        message_content = line_bot_api.get_message_content(message_id)
+        file_path = f"/tmp/{message_id}.png"
+        with open(file_path, "wb") as f:
+            for chunk in message_content.iter_content():
+                f.write(chunk)
+        print(f"✅ 圖片成功儲存到：{file_path}")
+        return file_path
+    except Exception as e:
+        print(f"❌ 圖片取得失敗：{e}")
+        return None
+
+def store_user_message(user_id, message_type, message_content):
+    """
+    儲存用戶的訊息。
+    """
+    user_message_history[user_id].append(
+        {"type": message_type, "content": message_content}
+    )
+
+def get_previous_message(user_id):
+    """
+    獲取用戶的上一則訊息。
+    """
+    if user_id in user_message_history and len(user_message_history[user_id]) > 0:
+        return user_message_history[user_id][-1]
+    return {"type": "text", "content": "No message!"}
+
+
+# ==========================
+#  LangChain 工具定義
+# ==========================
+
+@tool
+def generate_and_upload_image(prompt: str) -> str:
+    """
+    這個工具可以根據文字提示生成圖片，並將其上傳到伺服器。
+    
+    Args:
+        prompt: 用於生成圖片的文字提示。
+        
+    Returns:
+        生成圖片的 URL。
+    """
+    try:
+        genai_client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
+        response = genai_client.models.generate_content(
+            model="gemini-2.0-flash-exp",
+            contents=prompt,
+            config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
+        )
+        
+        image_binary = None
+        for part in response.candidates[0].content.parts:
+            if part.inline_data is not None:
+                image_binary = part.inline_data.data
+                break
+        
+        if image_binary:
+            image = PIL.Image.open(io.BytesIO(image_binary))
+            # 隨機生成一個檔案名以避免衝突
+            file_name = f"static/{os.urandom(16).hex()}.png"
+            image.save(file_name, format="PNG")
+            
+            image_url = os.path.join(os.getenv("HF_SPACE"), file_name)
+            return image_url
+            
+        return "圖片生成失敗。"
+    except Exception as e:
+        return f"圖片生成與上傳失敗: {e}"
+
+@tool
+def analyze_image_with_text(image_path: str, user_text: str) -> str:
+    """
+    這個工具可以根據一張圖片和一個文字提示來回答問題。
+    
+    Args:
+        image_path: 圖片在本地端儲存的路徑。
+        user_text: 針對圖片提出的文字問題。
+        
+    Returns:
+        模型針對圖片和文字提示給出的回應。
+    """
+    try:
+        if not os.path.exists(image_path):
+            return "圖片路徑無效，無法進行分析。"
+            
+        img_user = PIL.Image.open(image_path)
+        
+        # 建立一個新的聊天模型實例，以確保狀態獨立
+        chat_model = ChatGoogleGenerativeAI(
+            model="gemini-1.5-pro",
+            temperature=0.5,
+            top_p=0.5,
+            top_k=16
+        )
+        
+        response = chat_model.invoke([
+            {"role": "user", "content": [
+                {"type": "image_url", "image_url": {"url": f"file://{image_path}"}},
+                {"type": "text", "text": user_text}
+            ]}
+        ])
+        
+        return response.content
+    except Exception as e:
+        return f"圖片分析失敗: {e}"
+
+
+# ==========================
+#  LangChain 代理人設定
+# ==========================
+
+# 結合所有工具
+tools = [generate_and_upload_image, analyze_image_with_text]
+
+# 建立 LLM 模型實例
+llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro", temperature=0.5)
+
+# 建立提示模板
+prompt_template = ChatPromptTemplate.from_messages([
+    ("system", "你是一個強大的助理，可以根據用戶的請求使用提供的工具。"),
+    ("user", "{input}"),
+    ("placeholder", "{agent_scratchpad}"),
+])
+
+# 建立代理人
+agent = create_tool_calling_agent(llm, tools, prompt_template)
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+
+
+# ==========================
+#  FastAPI 路由
+# ==========================
+
 @app.get("/")
 def root():
     return {"title": "Line Bot"}
 
-# 處理 Line Webhook 請求
 @app.post("/webhook")
 async def webhook(
     request: Request,
     background_tasks: BackgroundTasks,
     x_line_signature=Header(None),
 ):
-    # 取得請求內容
     body = await request.body()
     try:
-        # 將處理 Line 事件的任務加入背景工作
         background_tasks.add_task(
             line_handler.handle, body.decode("utf-8"), x_line_signature
         )
     except InvalidSignatureError:
-        # 處理無效的簽章錯誤
         raise HTTPException(status_code=400, detail="Invalid signature")
     return "ok"
-    
 
-#==========================
-# 主程式(圖片與文字)
-#==========================
-# 建立 chat_sessions 字典
-chat_sessions = {}
 @line_handler.add(MessageEvent, message=(ImageMessage, TextMessage))
-def handle_image_message(event):
+def handle_message(event):
     user_id = event.source.user_id
-    user_text = event.message.text if event.message.type == "text" else None
-    previous_message = get_previous_message(user_id)
-
-    if event.message.type != "text" and event.message.type != "image":
-        line_bot_api.reply_message(event.reply_token, TextSendMessage(text="請輸入文字或圖片~"))
-        return
-    elif user_text == "再見":
-        line_bot_api.reply_message(event.reply_token, TextSendMessage(text="Bye!"))
-        return
-        
-    # ========
-    # 生成圖片
-    # ========
-    elif user_text and user_text.startswith("生成圖片"):
-        prompt = user_text.replace("生成圖片", "").strip()
-
-         # 先立即回覆避免token過期
-        line_bot_api.reply_message(event.reply_token, TextSendMessage(text="圖片生成中～ 請稍候.....✨"))
-        image_generator = Image_text_Generator(user_id)
-        # 生成圖片
-        image_binary = image_generator.generate_image_with_gemini(prompt)
-
-        if image_binary:
-            image_url = os.path.join(os.getenv("HF_SPACE"),image_generator.upload_image_to_tmp(image_binary)) ### os.getenv("HF_SPACE") => https://xxxxx.hf.space
-
-            if image_url:
-                # 使用 push message 發送圖片，避免 reply token 超時
-                line_bot_api.push_message(
-                    event.source.user_id,
-                    [
-                        TextSendMessage(text="✨ 這是我為你生成的圖片喔～"),
-                        ImageSendMessage(original_content_url=image_url, preview_image_url=image_url)
-                    ]
-                )
-            else:
-                line_bot_api.push_message(
-                    event.source.user_id,
-                    TextSendMessage(text="⚠️ 圖片上傳失敗，請稍後再試~")
-                )
+    
+    # 處理圖片上傳
+    if event.message.type == "image":
+        image_path = get_image_url_from_line(event.message.id)
+        if image_path:
+            store_user_message(user_id, "image", image_path)
+            line_bot_api.reply_message(
+                event.reply_token, TextSendMessage(text="圖片已接收成功囉，幫我輸入你想詢問的問題喔~")
+            )
         else:
-            line_bot_api.push_message(
-                event.source.user_id,
-                TextSendMessage(text="⚠️ 圖片生成失敗，請稍後再試~")
+            line_bot_api.reply_message(
+                event.reply_token, TextSendMessage(text="沒有接收到圖片~")
             )
         return
-        
-    # ========
-    # 純文字
-    # ========
-    elif event.message.type == "text" and previous_message["type"] != "image":
-        try: 
-            user_id = event.source.user_id
-            chat = chat_sessions.get(user_id) or client.chats.create(model="gemini-2.0-flash", config=generation_config)
-            chat_sessions[user_id] = chat
-            # 取得使用者輸入的文字
-            user_input = event.message.text
-            response = chat.send_message(user_input)
-            if (response.text != None):
-                out = response.text
-            else:
-                out = "Gemini沒答案!請換個說法！"
-        except:
-            # 處理錯誤
-            out = "Gemini執行出錯!請換個說法！" 
-        
-    elif previous_message and previous_message["type"] == "image" and event.message.type == "text":
-        image_path = previous_message["content"]
+    
+    # 處理文字訊息
+    if event.message.type == "text":
         user_text = event.message.text
-        store_user_message(user_id, "text", user_text)
-        try:
-            if not os.path.exists(image_path):
-                raise FileNotFoundError(f"圖片路徑無效：{image_path}")
-            previous_img = PIL.Image.open(image_path)
-            user_id = event.source.user_id
-            chat = chat_sessions.get(user_id) or client.chats.create(model="gemini-2.0-flash", config=generation_config)
-            chat_sessions[user_id] = chat
-            # 取得使用者輸入的文字
-            user_input = event.message.text
-            response = chat.send_message([previous_img, user_input])
-            if (response.text != None):
-                out = response.text
-            else:
-                out = "Gemini沒答案!請換個說法！"
-        except:
-            # 處理錯誤
-            out = "Gemini執行出錯!請換個說法！" 
-
-    # ========
-    # 上傳圖片
-    # ========
-    elif event.message.type == "image":
-        image_path = get_image_url(event.message.id)
-        if image_path:
-            store_user_message(user_id, "image", image_path)
-            line_bot_api.reply_message(event.reply_token, TextSendMessage(text="圖片已接收成功囉，幫我輸入你想詢問的問題喔~"))
+        previous_message = get_previous_message(user_id)
+        
+        # 根據上一則訊息類型，動態傳遞給代理人
+        if previous_message["type"] == "image":
+            image_path = previous_message["content"]
+            agent_input = {
+                "input": f"這是一張圖片在 {image_path}，請根據圖片和我的問題回答：{user_text}"
+            }
+            # 清除上一則圖片訊息，避免重複觸發
+            user_message_history[user_id].pop()
         else:
-            line_bot_api.reply_message(event.reply_token, TextSendMessage(text="沒有接收到圖片~"))
-        return
-
-    line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))
-
+            agent_input = {"input": user_text}
+            
+        try:
+            # 運行代理人
+            response = agent_executor.invoke(agent_input)
+            out = response["output"]
+            
+        except Exception as e:
+            print(f"代理人執行出錯: {e}")
+            out = f"代理人執行出錯!請換個說法或稍後再試！錯誤訊息：{e}"
+            
+        line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))
 
 if __name__ == "__main__":
-    # 啟動 FastAPI 應用程式
-    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)