app.py

import os, re, base64
from langchain_core.documents import Document
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
import chromadb
import gradio as gr

try:
    from langchain_memory import ConversationBufferMemory
except ImportError:
    try:
        from langchain.memory import ConversationBufferMemory
    except ImportError:
        from langchain_community.memory import ConversationBufferMemory

# === Embedding ===
embedding = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")

# === 載入 QA ===
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
qa_path = os.path.join(BASE_DIR, "QA_v2.txt")
if not os.path.exists(qa_path):
    raise FileNotFoundError(f"❌ 找不到 QA 檔案：{qa_path}")

with open(qa_path, "r", encoding="utf-8") as f:
    text = f.read()

pattern = r"(Ｑ[:：].*?)(?=Ｑ[:：]|$)"
qas = re.findall(pattern, text, flags=re.S)
qa_docs = {"證券": [], "期貨": [], "複委託": []}
for qa in qas:
    if "證券" in qa:
        qa_docs["證券"].append(Document(page_content=qa.strip()))
    elif "期貨" in qa:
        qa_docs["期貨"].append(Document(page_content=qa.strip()))
    elif "複委託" in qa:
        qa_docs["複委託"].append(Document(page_content=qa.strip()))

print("✅ 已成功讀取 QA 並完成分類：")
for k, v in qa_docs.items():
    print(f"　{k}：{len(v)} 筆")

# === 向量資料庫 ===
client = chromadb.PersistentClient(path="./chroma_db")
collections = {"證券": "stocks", "期貨": "futures", "複委託": "overseas"}
vectordbs = {}
for cat, docs in qa_docs.items():
    vectordbs[cat] = Chroma(
        client=client,
        collection_name=collections[cat],
        embedding_function=embedding
    )
    if len(vectordbs[cat].get()["documents"]) == 0:
        vectordbs[cat].add_documents(docs)
    else:
        print(f"⚙️ 已載入現有向量資料庫：{collections[cat]}")
print("✅ 各類別向量資料庫建立完成")

# === Gemini ===
API_KEY = os.getenv("GOOGLE_API_KEY")
if not API_KEY:
    raise ValueError("⚠️ 未設定 GOOGLE_API_KEY。")

llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=API_KEY)
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# === 對話主邏輯 ===
def auto_detect_category(text):
    if any(k in text for k in ["股票", "證券", "開戶", "下單", "交割"]):
        return "證券"
    elif any(k in text for k in ["期貨", "選擇權", "保證金"]):
        return "期貨"
    elif any(k in text for k in ["複委託", "海外", "美股", "港股"]):
        return "複委託"
    else:
        return "證券"

def chat_fn(message, history):
    category = auto_detect_category(message)
    vectordb = vectordbs.get(category)
    docs = vectordb.similarity_search(message, k=2)
    context = "\n\n".join([d.page_content for d in docs]) if docs else "查無資料"
    prompt = f"""
你是一位金融客服人員，根據以下公司QA回答客戶問題：
---
{context}
---
使用者問題：{message}
    """
    try:
        response = llm.invoke(prompt)
        reply = response.content.strip()
    except Exception as e:
        reply = f"⚠️ 生成錯誤：{e}"
    return reply or "請洽營業員"

# === Gradio 介面 ===
logo_path = os.path.join(BASE_DIR, "mega.png")
logo_base64 = ""
if os.path.exists(logo_path):
    with open(logo_path, "rb") as f:
        logo_base64 = base64.b64encode(f.read()).decode("utf-8")

with gr.Blocks(
    theme="soft",
    css="""
        #logo-top {position:fixed;top:12px;left:18px;z-index:1000;
            background:white;border-radius:10px;padding:6px 8px;
            box-shadow:0 0 8px rgba(0,0,0,0.15);}
        #logo-top img{width:120px;height:auto;display:block;}

        #main-title{text-align:center;font-size:28px;font-weight:bold;margin:0;line-height:1.4;}
        @media (max-width:768px){
            #main-title{white-space:pre-line;font-size:24px;}
        }

        #footer{text-align:center;font-size:13px;color:#999;border-top:1px solid #ddd;
            padding-top:8px;margin-top:30px;background:transparent;}

        #input-row{display:flex;align-items:center;gap:8px;margin-top:10px;}
        #user-input{flex-grow:1;border-radius:20px;border:1px solid #ccc;
            padding:6px 12px;font-size:15px;background-color:#fff;}
        #send-btn{background:#00b800;color:white;border:none;border-radius:14px;
            height:26px;padding:0 10px;font-size:13px;font-weight:600;cursor:pointer;}
        #send-btn:hover{background:#00a000;}
    """
) as demo:
    if logo_base64:
        gr.HTML(f"<div id='logo-top'><img src='data:image/png;base64,{logo_base64}'></div>")

    gr.HTML("""
        <h1 id='main-title'>👨‍💼 我是小智<br>您的金融好幫手 🫰</h1>
        <p style='text-align:center;margin-top:8px;color:#666;'>Powered by Gemini & LangChain</p>
    """)

    with gr.Row():
        with gr.Column(scale=4):
            chatbox = gr.Chatbot(label="💬 對話紀錄", type="messages")
            with gr.Row(elem_id="input-row"):
                user_input = gr.Textbox(elem_id="user-input", show_label=False, placeholder="輸入訊息...", scale=8)
                send_btn = gr.Button("輸入", elem_id="send-btn", scale=1)
            def handle_input(message, history):
                if not message.strip():
                    return history, gr.update(value="")
                reply = chat_fn(message, history)
                history += [{"role":"user","content":message},
                            {"role":"assistant","content":reply}]
                return history, gr.update(value="")
            user_input.submit(handle_input,[user_input,chatbox],[chatbox,user_input])
            send_btn.click(handle_input,[user_input,chatbox],[chatbox,user_input])
        with gr.Column(scale=1):
            gr.Markdown("### 👇 快速提問")
            for label,q in [
                ("未成年可以開戶嗎？","未成年可以開戶嗎?"),
                ("法人開戶要準備什麼？","法人開戶要準備什麼?"),
                ("期貨交易保證金是什麼？","期貨交易保證金是什麼?"),
                ("複委託要如何下單？","複委託要如何下單?"),
                ("美股交易時間？","美股交易時間?"),
                ("美股可以定期定額嗎？","美股可以定期定額嗎?")
            ]:
                gr.Button(label).click(lambda h,q=q: handle_input(q,h),[chatbox],[chatbox,user_input])
            def clear_memory():
                memory.clear()
                return [],gr.update(value="",placeholder="輸入訊息...")
            gr.Button("🧹 整理畫面").click(clear_memory,outputs=[chatbox,user_input])

    gr.HTML("<div id='footer'>© Fintech Assistant — 僅業務使用，非官方授權</div>")

demo.launch()