Update main.py

main.py

@@ -2,7 +2,7 @@ import json, os
 import gradio as gr
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import FastAPI, Request,  Header, BackgroundTasks, HTTPException, status
-import google.generativeai as genai
+#import google.generativeai as genai
 import base64
 from collections import defaultdict
 from linebot import LineBotApi, WebhookHandler
@@ -10,15 +10,22 @@ from linebot.exceptions import InvalidSignatureError
 from linebot.models import MessageEvent, TextMessage, TextSendMessage, ImageSendMessage, AudioMessage, ImageMessage
 import PIL.Image
 
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationChain
+
 # 設定 Google AI API 金鑰
-genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+ChatGoogleGenerativeAI.configure(api_key=os.environ["GOOGLE_API_KEY"])
 
 # 設定生成文字的參數
-generation_config = genai.types.GenerationConfig(max_output_tokens=2048, temperature=0.2, top_p=0.5, top_k=16)
+generation_config = ChatGoogleGenerativeAI.types.GenerationConfig(max_output_tokens=2048, temperature=0.2, top_p=0.5, top_k=16)
 
 # 使用 Gemini-1.5-flash 模型
 model = genai.GenerativeModel('gemini-2.0-flash-exp', system_instruction="請用繁體中文回答。你現在是個專業助理，職稱為OPEN小助理，個性活潑、樂觀，願意回答所有問題", generation_config=generation_config)
 chat = model.start_chat(history=[])
+# 2. 初始化 ConversationChain 和 Memory
+memory = ConversationBufferMemory()
+conversation = ConversationChain(llm=model, memory=memory)
 
 # 設定 Line Bot 的 API 金鑰和秘密金鑰
 line_bot_api = LineBotApi(os.environ["CHANNEL_ACCESS_TOKEN"])
@@ -97,6 +104,8 @@ def analyze_with_gemini(image_path, user_text):
         
         organ = PIL.Image.open(image_path)
         response = chat.send_message([user_text, organ])
+        #chain = get_or_create_chain(user_id)
+        #completion = chain.run(prompt)
         
         # 提取回應內容
         return response.text
@@ -114,6 +123,14 @@ def get_previous_message(user_id):
         return user_message_history[user_id][-1]
     return None
 
+# 分開不同ID的歷史
+chat_memories = {}
+def get_or_create_chain(user_id):
+    if user_id not in chat_memories:
+        memory = ConversationBufferMemory()
+        chat_memories[user_id] = ConversationChain(llm=model, memory=memory)
+    return chat_memories[user_id]
+
 @line_handler.add(MessageEvent, message=(ImageMessage,TextMessage))
 def handle_image_message(event):
     user_id = event.source.user_id
@@ -167,7 +184,9 @@ def handle_image_message(event):
                 store_user_message(user_id, "text", prompt) 
    
                 # 使用 Gemini 模型生成文字
-                completion = chat.send_message(prompt)
+                #completion = chat.send_message(prompt)
+                chain = get_or_create_chain(user_id)
+                completion = chain.run(prompt) # 使用帶對話歷史的方法
                 # 檢查生成結果是否為空
                 if (completion.text != None):
                     # 取得生成結果