Update main.py

main.py

@@ -1,6 +1,11 @@
-from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException, status
+import os
+import io
+from collections import defaultdict
 from fastapi.middleware.cors import CORSMiddleware
+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 (
@@ -10,62 +15,36 @@ from linebot.models import (
     ImageSendMessage,
     ImageMessage,
 )
-from google import genai
-from google.genai import types
-from PIL import Image
-from collections import defaultdict
-import os
-import io
-import requests
+import PIL.Image
 import uvicorn
-import logging
-import base64
+
+# LangChain 相關匯入
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.tools import tool
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain.agents import AgentExecutor, create_tool_calling_agent
-from pydantic import BaseModel, Field
-from typing import List
 
-# ==========================#  環境設定與工具函式# ==========================#
-# 設置日誌記錄，級別為 INFO，格式包含時間、級別和訊息
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# 檢查必要的環境變數是否都已設定
-if not all(k in os.environ for k in ["CHANNEL_ACCESS_TOKEN", "CHANNEL_SECRET", "GOOGLE_API_KEY", "HF_SPACE"]):
-    logging.error("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and HF_SPACE.")
-    raise ValueError("Missing environment variables. Please set CHANNEL_ACCESS_TOKEN, CHANNEL_SECRET, GOOGLE_API_KEY, and HF_SPACE.")
+# ==========================
+#  環境設定與工具函式
+# ==========================
 
-# 獲取環境變數中的金鑰和 URL
+# 設置 Google AI API 金鑰
 google_api = os.environ["GOOGLE_API_KEY"]
-line_channel_access_token = os.environ["CHANNEL_ACCESS_TOKEN"]
-line_channel_secret = os.environ["CHANNEL_SECRET"]
-HF_SPACE_URL = os.environ["HF_SPACE"]
-
-# Line Bot API 設定
-line_bot_api = LineBotApi(line_channel_access_token)
-line_handler = WebhookHandler(line_channel_secret)
-
-# Google AI API 設定
 genai_client = genai.Client(api_key=google_api)
 
-# 使用者狀態追蹤，用來儲存已上傳的衣物圖片 URL
-user_states = defaultdict(lambda: {
-    "upper_body_images": [],
-    "lower_body_images": [],
-    "current_mode": None,
-})
+# 設置 Line Bot 的 API 金鑰和秘密金鑰
+line_bot_api = LineBotApi(os.environ["CHANNEL_ACCESS_TOKEN"])
+line_handler = WebhookHandler(os.environ["CHANNEL_SECRET"])
 
-MAX_IMAGES_PER_TYPE = 3
-IMAGIN_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent"
+# 使用字典模擬用戶訊息歷史存儲
+user_message_history = defaultdict(list)
 
 # 建立 FastAPI 應用程式
 app = FastAPI()
-
-# 設定靜態文件服務，用於託管圖片
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# 設定 CORS 跨域請求
+# 設定 CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -74,134 +53,138 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-def get_base64_from_url(image_url: str):
-    """從 URL 下載圖片並轉換為 Base64 編碼。"""
-    try:
-        response = requests.get(image_url)
-        response.raise_for_status()
-        return base64.b64encode(response.content).decode('utf-8')
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Error fetching image from URL: {image_url}, error: {e}")
-        return None
-        
-def save_image_locally(image_binary: bytes):
-    """
-    將二進位圖片資料儲存到本地，並返回一個可供外部存取的 URL。
-    """
-    try:
-        # 確保 'static' 資料夾存在
-        if not os.path.exists("static"):
-            os.makedirs("static")
-        
-        image = 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(HF_SPACE_URL, file_name)
-        logging.info(f"Image successfully saved locally: {image_url}")
-        return image_url
-    except Exception as e:
-        logging.error(f"Error saving image locally: {e}")
-        return None
-
 def get_image_url_from_line(message_id):
     """
     從 Line 訊息 ID 獲取圖片內容並儲存到暫存檔案。
     """
     try:
         message_content = line_bot_api.get_message_content(message_id)
-        # 取得二進位圖片資料
-        image_binary = message_content.content
-        # 使用 save_image_locally 函式儲存圖片並取得 URL
-        return save_image_locally(image_binary)
+        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
 
-# ==========================#  LangChain 工具定義# ==========================#
-# 定義工具的輸入模型，明確指定參數的型別和描述
-class OutfitInput(BaseModel):
-    """用於生成穿搭圖片的輸入參數。"""
-    upper_body_urls: List[str] = Field(..., description="上衣圖片的 URL 列表。")
-    lower_body_urls: List[str] = Field(..., description="褲子/裙子圖片的 URL 列表。")
+def store_user_message(user_id, message_type, message_content):
+    """
+    儲存用戶的訊息。
+    """
+    user_message_history[user_id].append(
+        {"type": message_type, "content": message_content}
+    )
 
-@tool(args_schema=OutfitInput)
-def generate_outfit_from_clothes(upper_body_urls: list, lower_body_urls: list) -> str:
+def get_previous_message(user_id):
+    """
+    獲取用戶的上一則訊息。
     """
-    這個工具可以根據提供的上衣和褲子/裙子圖片 URLs，生成一套全新的穿搭圖片。
+    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:
-        upper_body_urls: 一組上衣圖片的 URL 列表。
-        lower_body_urls: 一組褲子/裙子圖片的 URL 列表。
+        prompt: 用於生成圖片的文字提示。
         
     Returns:
         回傳生成圖片的 URL。
     """
-    logging.info("Attempting to generate a new outfit image.")
-    prompt = "使用提供的上衣和褲子/裙子圖片，生成一套完整且時尚的穿搭圖片。請將衣服呈現在一個有模特兒穿著或是在平面上呈現的完整畫面中。風格應與提供的衣物相符。"
-    
-    all_images_base64 = (
-        [get_base64_from_url(url) for url in upper_body_urls] +
-        [get_base64_from_url(url) for url in lower_body_urls]
-    )
-
-    payload = {
-        "contents": [
-            {
-                "parts": [
-                    {"text": prompt}
-                ] + [
-                    {"inlineData": {"mimeType": "image/jpeg", "data": img_base64}} for img_base64 in all_images_base64
-                ]
-            }
-        ],
-        "generationConfig": {
-            "responseModalities": ['IMAGE']
-        }
-    }
+    try:
+        response = genai_client.models.generate_content(
+            model="gemini-2.0-flash-preview-image-generation",
+            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:
-        response = requests.post(f"{IMAGIN_API_URL}?key={os.environ['GOOGLE_API_KEY']}", json=payload, timeout=120)
-        response.raise_for_status()
-        response_data = response.json()
-        if 'candidates' in response_data and response_data['candidates']:
-            image_part = response_data['candidates'][0]['content']['parts'][0]
-            if image_part and 'inlineData' in image_part:
-                generated_image_base64 = image_part['inlineData']['data']
-                generated_image_url = save_image_locally(base64.b64decode(generated_image_base64))
-                if generated_image_url:
-                    logging.info(f"Successfully generated and saved new outfit image: {generated_image_url}")
-                    return generated_image_url
-                else:
-                    logging.error("Failed to save generated outfit image locally.")
-                    return None
-            else:
-                logging.error("Gemini image response format is invalid for outfit generation.")
-                return None
+        if not os.path.exists(image_path):
+            return "圖片路徑無效，無法進行分析。"
+            
+        img_user = PIL.Image.open(image_path)
+        response = genai_client.models.generate_content(
+                    model="gemini-2.5-flash",
+                    contents=[img_user, user_text]
+        )
+        if (response.text != None):
+            out = response.text
         else:
-            logging.error(f"Gemini outfit generation response has no candidates: {response.text}")
-            return None
-    except requests.exceptions.RequestException as e:
-        logging.error(f"Gemini outfit generation API request failed: {e}")
-        return None
+            out = "Gemini沒答案!請換個說法！"
+    except Exception as e:
+        # 處理錯誤
+        out = f"Gemini執行出錯: {e}"
+        
+    return out
+
+
+# ==========================
+#  LangChain 代理人設定
+# ==========================
 
-# ==========================#  LangChain 代理人設定# ==========================#
 # 結合所有工具
-tools = [generate_outfit_from_clothes]
+tools = [generate_and_upload_image, analyze_image_with_text]
+
 # 建立 LLM 模型實例
 llm = ChatGoogleGenerativeAI(google_api_key=google_api, model="gemini-2.5-flash", temperature=0.2)
+
 # 建立提示模板
 prompt_template = ChatPromptTemplate([
-    ("system", "你是一個強大的圖像生成與問答助理，可以根據用戶的請求使用提供的工具。當你執行 generate_outfit_from_clothes 工具成功後會獲得一個 URL，然後你回答的 output 要包含有這個 URL 的完整資訊。如果工具有產生錯誤訊息請解讀並回應。"),
+    ("system", "你是一個強大的圖像生成與問答助理，可以根據用戶的請求使用提供的工具。當你執行 generate_and_upload_image 工具\
+    成功後會獲得一個 URL，然後你回答的 output 要包含有這個 URL 的完整資訊。如果工具有產生錯誤訊息請解讀並回應。"),
     ("user", "{input}"),
     ("placeholder", "{agent_scratchpad}"),
 ])
+
 # 建立代理人
 agent = create_tool_calling_agent(llm, tools, prompt_template)
 agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
-# ==========================#  FastAPI 路由# ==========================#
+# ==========================
+#  FastAPI 路由
+# ==========================
+
 @app.get("/")
 def root():
     return {"title": "Line Bot"}
@@ -227,118 +210,53 @@ def handle_message(event):
     
     # 處理圖片上傳
     if event.message.type == "image":
-        image_url = get_image_url_from_line(event.message.id)
-        if image_url:
-            mode = user_states[user_id]["current_mode"]
-            if not mode:
-                line_bot_api.reply_message(
-                    event.reply_token, TextSendMessage(text="請先輸入「上衣」或「褲子」來選擇要上傳的衣服類型。")
-                )
-                return
-
-            if mode == "upper":
-                if len(user_states[user_id]["upper_body_images"]) < MAX_IMAGES_PER_TYPE:
-                    user_states[user_id]["upper_body_images"].append(image_url)
-                    reply_text = f"已接收上衣。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。"
-                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text=reply_text))
-                else:
-                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text="上衣數量已滿。"))
-            else: # mode == "lower"
-                if len(user_states[user_id]["lower_body_images"]) < MAX_IMAGES_PER_TYPE:
-                    user_states[user_id]["lower_body_images"].append(image_url)
-                    reply_text = f"已接收褲子/裙子。上衣: {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE}，褲子/裙子: {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE}。"
-                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text=reply_text))
-                else:
-                    line_bot_api.reply_message(event.reply_token, TextSendMessage(text="褲子/裙子數量已滿。"))
-
-            # 檢查是否所有圖片都已上傳完畢
-            if (len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and
-                len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE):
-                line_bot_api.push_message(
-                    user_id,
-                    TextSendMessage(text="所有衣物圖片已收集完畢！\n\n現在您可以輸入「**生成穿搭**」或「**圖片推薦**」來獲得一套新的圖片穿搭。")
-                )
+        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.reply_message(
                 event.reply_token, TextSendMessage(text="沒有接收到圖片~")
             )
-            
+    
     # 處理文字訊息
     elif event.message.type == "text":
-        user_text = event.message.text.lower()
-        reply_token = event.reply_token
-        
-        # 處理重置功能
-        if user_text in ["重置", "重來", "重新開始", "再一次"]:
-            user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None})[user_id]
-            line_bot_api.reply_message(
-                reply_token, TextSendMessage(text="狀態已重置。請先輸入「上衣」或「褲子」來開始。")
-            )
-            return
-            
-        # 處理衣物上傳模式切換
-        if user_text == "上衣":
-            user_states[user_id]["current_mode"] = "upper"
-            line_bot_api.reply_message(
-                reply_token, TextSendMessage(text=f"請上傳三件上衣圖片，您已上傳 {len(user_states[user_id]['upper_body_images'])}/{MAX_IMAGES_PER_TYPE} 張。")
-            )
-            return
-        elif user_text in ["褲子", "裙子"]:
-            user_states[user_id]["current_mode"] = "lower"
-            line_bot_api.reply_message(
-                reply_token, TextSendMessage(text=f"請上傳三件褲子/裙子圖片，您已上傳 {len(user_states[user_id]['lower_body_images'])}/{MAX_IMAGES_PER_TYPE} 張。")
-            )
-            return
-
-        # 處理圖片生成穿搭
-        if user_text in ["生成穿搭", "圖片推薦"]:
-            if (len(user_states[user_id]["upper_body_images"]) == MAX_IMAGES_PER_TYPE and 
-                len(user_states[user_id]["lower_body_images"]) == MAX_IMAGES_PER_TYPE):
-                try:
-                    line_bot_api.reply_message(reply_token, TextSendMessage(text="好的，正在為您生成一套新的穿搭圖片，請稍候..."))
-                    
-                    # 直接呼叫生成工具，並傳入已收集的圖片 URLs
-                    generated_image_url = generate_outfit_from_clothes(
-                        user_states[user_id]["upper_body_images"],
-                        user_states[user_id]["lower_body_images"]
-                    )
-                    
-                    if generated_image_url:
-                        line_bot_api.push_message(
-                            user_id,
-                            ImageSendMessage(original_content_url=generated_image_url, preview_image_url=generated_image_url)
-                        )
-                        line_bot_api.push_message(
-                            user_id,
-                            TextSendMessage(text="這是根據您的衣物生成的圖片推薦。如果想再次使用，請輸入「重置」。")
-                        )
-                    else:
-                         line_bot_api.push_message(user_id, TextSendMessage(text="圖片生成失敗，請稍後再試。"))
-                    
-                except Exception as e:
-                    logging.error(f"Error generating outfit image: {e}")
-                    line_bot_api.push_message(user_id, TextSendMessage(text=f"圖片生成失敗，請稍後再試。錯誤：{e}"))
-
-                # 生成後重置狀態以便下一次使用
-                user_states[user_id] = defaultdict(lambda: {"upper_body_images": [], "lower_body_images": [], "current_mode": None})[user_id]
-                return
-            else:
-                line_bot_api.reply_message(
-                    reply_token, TextSendMessage(text="請先上傳三件上衣和三件褲子/裙子圖片，再輸入「生成穿搭」來獲得圖片推薦。")
-                )
-                return
+        user_text = event.message.text
+        previous_message = get_previous_message(user_id)
+        print(previous_message)
 
-        # 如果都不是特定指令，則交給代理人處理
-        agent_input = {"input": user_text}
+        # 根據上一則訊息類型，動態傳遞給代理人
+        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:
+            agent_input = {"input": user_text}
         try:
             # 運行代理人
             response = agent_executor.invoke(agent_input)
             out = response["output"]
-            line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))
+            if 'https' in out:
+                img_tmp = 'https'+out.split('https')[1]
+                image_url = img_tmp.split('png')[0]+'png'
+                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.reply_message(event.reply_token, TextSendMessage(text=out))
         except Exception as e:
             print(f"代理人執行出錯: {e}")
             out = f"代理人執行出錯!錯誤訊息：{e}"
             line_bot_api.reply_message(event.reply_token, TextSendMessage(text=out))
 
 if __name__ == "__main__":
-    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)