rag2line / main.py
hazelhh's picture
Update main.py
c4283a0 verified
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, Request, Header, BackgroundTasks, HTTPException, status, staticfiles
from gradio_client import Client
import json
import os
from linebot import (
LineBotApi, WebhookHandler
)
from linebot.exceptions import (
InvalidSignatureError
)
from linebot.models import (
MessageEvent, TextMessage, TextSendMessage, ImageSendMessage, AudioMessage
)
# 設定 GeminiRAG 的HF網址
client = Client(os.environ["GeminiRAGapi"], hf_token=os.environ['HF_TOKEN'])
# 設定 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"
# ==========================
# LangChain 代理人設定
# ==========================
# 結合所有工具
tools = [generate_and_upload_image, analyze_image_with_text]
# 建立 LLM 模型實例
llm = ChatGoogleGenerativeAI(google_api_key=google_api, model="gemini-2.5-flash-lite", temperature=0.2)
# 建立提示模板
prompt_template = ChatPromptTemplate([
("system", "你是一個強大的虛擬穿搭助理,可以根據用戶的請求使用提供的工具。當你執行 generate_and_upload_image 工具\
成功後,你將會獲得一個圖片的包含https的完整網址,你的任務是將這個包含https的完整網址回傳。如果工具有產生錯誤訊息請解讀回應。"),
("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 = FastAPI()
# app.mount("/static", staticfiles.StaticFiles(directory="static"), name="static")
# 設定 CORS,允許跨域請求
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 處理根路徑請求
@app.get("/")
def root():
return {"title": "RAG 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"
# 處理文字訊息事件
@line_handler.add(MessageEvent, message=TextMessage)
def handle_message(event):
global working_status
# 檢查事件類型和訊息類型
if event.type != "message" or event.message.type != "text":
# 回覆錯誤訊息
line_bot_api.reply_message(
event.reply_token,
TextSendMessage(text="Event type error:[No message or the message does not contain text]")
)
# 檢查使用者是否輸入 "再見"
elif event.message.text == "再見,掰掰,結束,":
# 回覆 "ByeBye!"
line_bot_api.reply_message(
event.reply_token,
TextSendMessage(text="Bye!")
)
return
# 檢查是否正在與使用者交談
elif working_status:
try:
# 取得使用者輸入的文字
prompt = event.message.text
# 使用 GeminiRAGapi
completion = client.predict(question=prompt, api_name="/predict")
# 檢查生成結果是否為空
if (completion != None):
# 取得生成結果
out = completion
# 判斷如果是文字
elif type=='text':
msg = json_data['events'][0]['message']['text'] # 取得 LINE 收到的文字訊息
reply = msg
# 判斷如果是圖片
elif type == 'image':
msgID = json_data['events'][0]['message']['id'] # 取得訊息 id
message_content = line_bot_api.get_message_content(msgID) # 根據訊息 ID 取得訊息內容
# 在同樣的資料夾中建立以訊息 ID 為檔名的 .jpg 檔案
with open(f'{msgID}.jpg', 'wb') as fd:
fd.write(message_content.content) # 以二進位的方式寫入檔案
reply = '圖片儲存完成!' # 設定要回傳的訊息
else:
reply = '你傳的不是文字或圖片呦~'
print(reply)
line_bot_api.reply_message(tk,TextSendMessage(reply)) # 回傳訊息
except:
# 處理錯誤
out = "Gemini執行出錯!請換個說法!"
# 回覆生成結果
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)