app.py

"""
Streamlit + Groq API - 8種 RAG 策略 PDF 問答系統
安裝依賴: pip install streamlit groq pypdf sentence-transformers numpy faiss-cpu scikit-learn
執行方式: streamlit run rag_streamlit.py
"""

import streamlit as st
from groq import Groq
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
from pypdf import PdfReader
import re
from sklearn.feature_extraction.text import TfidfVectorizer
import io

# ─────────────────────────────────────────────
# RAG 核心類別
# ─────────────────────────────────────────────

class MultiStrategyRAG:
    def __init__(self, api_key: str):
        self.client = Groq(api_key=api_key)
        self.embedding_model = SentenceTransformer(
            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
        )
        self.chunks: list[str] = []
        self.embeddings = None
        self.index = None
        self.tfidf_vectorizer = None
        self.tfidf_matrix = None

    # ── 文件載入 ──────────────────────────────

    def load_pdf(self, pdf_bytes: bytes) -> str:
        try:
            reader = PdfReader(io.BytesIO(pdf_bytes))
            full_text = "".join(
                (page.extract_text() or "") + "\n" for page in reader.pages
            )

            self.chunks = self._split_text(full_text, chunk_size=800, overlap=150)

            self.embeddings = self.embedding_model.encode(
                self.chunks, convert_to_numpy=True
            )

            dim = self.embeddings.shape[1]
            self.index = faiss.IndexFlatL2(dim)
            self.index.add(self.embeddings.astype("float32"))

            self.tfidf_vectorizer = TfidfVectorizer(max_features=1000)
            self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(self.chunks)

            return f"✅ 成功載入 PDF！共 {len(reader.pages)} 頁，分割為 {len(self.chunks)} 個片段"
        except Exception as e:
            return f"❌ 載入失敗：{e}"

    def _split_text(self, text: str, chunk_size: int, overlap: int) -> list[str]:
        chunks, start = [], 0
        while start < len(text):
            chunk = re.sub(r"\s+", " ", text[start : start + chunk_size]).strip()
            if chunk:
                chunks.append(chunk)
            start += chunk_size - overlap
        return chunks

    # ── 8 種 RAG 策略 ─────────────────────────

    def strategy_1_basic_similarity(self, query: str, top_k: int = 3) -> list[str]:
        """策略1：基礎語意相似度搜尋"""
        vec = self.embedding_model.encode([query]).astype("float32")
        _, indices = self.index.search(vec, top_k)
        return [self.chunks[i] for i in indices[0]]

    def strategy_2_tfidf(self, query: str, top_k: int = 3) -> list[str]:
        """策略2：TF-IDF 關鍵詞搜尋"""
        qvec = self.tfidf_vectorizer.transform([query])
        scores = (self.tfidf_matrix * qvec.T).toarray().flatten()
        top_idx = scores.argsort()[-top_k:][::-1]
        return [self.chunks[i] for i in top_idx]

    def strategy_3_hybrid(self, query: str, top_k: int = 3) -> list[str]:
        """策略3：混合搜尋（語意 + TF-IDF）"""
        vec = self.embedding_model.encode([query]).astype("float32")
        _, sem_idx = self.index.search(vec, top_k * 2)

        qvec = self.tfidf_vectorizer.transform([query])
        tfidf_scores = (self.tfidf_matrix * qvec.T).toarray().flatten()
        tfidf_idx = tfidf_scores.argsort()[-top_k * 2 :][::-1]

        combined = list(dict.fromkeys(sem_idx[0].tolist() + tfidf_idx.tolist()))
        return [self.chunks[i] for i in combined[:top_k]]

    def strategy_4_reranking(self, query: str, top_k: int = 3) -> list[str]:
        """策略4：重新排序（LLM 評分重排）"""
        candidates = self.strategy_1_basic_similarity(query, top_k=top_k * 2)
        reranked = []
        for chunk in candidates:
            prompt = (
                f"問題：{query}\n\n文本：{chunk[:200]}...\n\n"
                "這段文本與問題的相關度(0-10)："
            )
            try:
                resp = self.client.chat.completions.create(
                    model="llama-3.1-8b-instant",
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=10,
                    temperature=0,
                )
                raw = resp.choices[0].message.content.strip()
                nums = re.findall(r"\d+", raw)
                score = float(nums[0]) if nums else 0
            except Exception:
                score = 0
            reranked.append((chunk, score))
        reranked.sort(key=lambda x: x[1], reverse=True)
        return [c for c, _ in reranked[:top_k]]

    def strategy_5_multi_query(self, query: str, top_k: int = 3) -> list[str]:
        """策略5：多查詢擴展"""
        prompt = f"將以下問題改寫成3個相關但不同角度的問題，用換行分隔：\n{query}"
        try:
            resp = self.client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=200,
                temperature=0.7,
            )
            queries = [query] + resp.choices[0].message.content.strip().split("\n")[:3]
        except Exception:
            queries = [query]

        all_chunks: list[str] = []
        for q in queries:
            all_chunks.extend(self.strategy_1_basic_similarity(q, top_k=2))
        return list(dict.fromkeys(all_chunks))[:top_k]

    def strategy_6_contextual_compression(self, query: str, top_k: int = 3) -> list[str]:
        """策略6：上下文壓縮"""
        chunks = self.strategy_1_basic_similarity(query, top_k=top_k)
        compressed = []
        for chunk in chunks:
            prompt = (
                f"從以下文本中提取與問題「{query}」最相關的1-2句話：\n\n{chunk}"
            )
            try:
                resp = self.client.chat.completions.create(
                    model="llama-3.1-8b-instant",
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=150,
                    temperature=0,
                )
                compressed.append(resp.choices[0].message.content.strip())
            except Exception:
                compressed.append(chunk[:300])
        return compressed

    def strategy_7_parent_child(self, query: str, top_k: int = 3) -> list[str]:
        """策略7：父子文檔（小片段對應大上下文）"""
        small_chunks = self._split_text(" ".join(self.chunks), chunk_size=300, overlap=50)
        small_emb = self.embedding_model.encode(small_chunks, convert_to_numpy=True)
        small_index = faiss.IndexFlatL2(small_emb.shape[1])
        small_index.add(small_emb.astype("float32"))

        vec = self.embedding_model.encode([query]).astype("float32")
        _, indices = small_index.search(vec, top_k)

        results = []
        for idx in indices[0]:
            for big in self.chunks:
                if small_chunks[idx] in big:
                    results.append(big)
                    break
        return list(dict.fromkeys(results))[:top_k]

    def strategy_8_hypothetical_answer(self, query: str, top_k: int = 3) -> list[str]:
        """策略8：假設性答案（HyDE）"""
        prompt = f"請對以下問題給出一個假設性的答案（即使不確定）：\n{query}"
        try:
            resp = self.client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=200,
                temperature=0.7,
            )
            hypo = resp.choices[0].message.content
        except Exception:
            hypo = query

        vec = self.embedding_model.encode([hypo]).astype("float32")
        _, indices = self.index.search(vec, top_k)
        return [self.chunks[i] for i in indices[0]]

    # ── 答案生成 ──────────────────────────────

    def generate_answer(
        self, query: str, strategy: str, top_k: int = 3
    ) -> tuple[str, str]:
        if not self.chunks:
            return "❌ 請先上傳 PDF 檔案！", ""

        strategy_map = {
            "1. 基礎語意搜尋":       self.strategy_1_basic_similarity,
            "2. TF-IDF 關鍵詞":      self.strategy_2_tfidf,
            "3. 混合搜尋":           self.strategy_3_hybrid,
            "4. 重新排序":           self.strategy_4_reranking,
            "5. 多查詢擴展":         self.strategy_5_multi_query,
            "6. 上下文壓縮":         self.strategy_6_contextual_compression,
            "7. 父子文檔":           self.strategy_7_parent_child,
            "8. 假設性答案 (HyDE)":  self.strategy_8_hypothetical_answer,
        }
        retrieval_fn = strategy_map.get(strategy, self.strategy_1_basic_similarity)
        relevant_chunks = retrieval_fn(query, top_k)
        context = "\n\n---\n\n".join(relevant_chunks)

        prompt = (
            "請根據以下上下文回答問題。如果上下文中沒有相關資訊，請說明無法回答。\n\n"
            f"上下文：\n{context}\n\n問題：{query}\n\n請用繁體中文詳細回答："
        )

        try:
            resp = self.client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {"role": "system", "content": "你是專業的文件分析助手。"},
                    {"role": "user",   "content": prompt},
                ],
                max_tokens=1024,
                temperature=0.3,
            )
            answer = resp.choices[0].message.content
            source_info = (
                f"📚 使用策略：{strategy}\n"
                f"📄 檢索片段數：{len(relevant_chunks)}\n\n"
                + "=" * 50 + "\n相關文本片段：\n" + "=" * 50
                + f"\n\n{context}"
            )
            return answer, source_info
        except Exception as e:
            return f"❌ 生成答案失敗：{e}", ""


# ─────────────────────────────────────────────
# Streamlit 頁面
# ─────────────────────────────────────────────

STRATEGY_DESCRIPTIONS = {
    "1. 基礎語意搜尋":      "使用向量餘弦相似度，將問題與文件片段對比，找出語意最接近的段落。",
    "2. TF-IDF 關鍵詞":     "基於詞頻-逆文檔頻率（TF-IDF）統計，適合精確關鍵字匹配場景。",
    "3. 混合搜尋":          "同時執行語意搜尋與 TF-IDF，融合兩者結果，兼顧語意與關鍵字。",
    "4. 重新排序":          "先用語意搜尋召回候選片段，再讓 LLM 為每段打分重新排序。",
    "5. 多查詢擴展":        "讓 LLM 將問題改寫為多個角度的問題，再分別搜尋合併結果。",
    "6. 上下文壓縮":        "先語意搜尋，再請 LLM 從每段中萃取與問題最相關的 1-2 句。",
    "7. 父子文檔":          "以更小的子片段搜尋，但回傳包含該子片段的原始大段文本。",
    "8. 假設性答案 (HyDE)": "先讓 LLM 生成一個假設答案，用此假設答案的向量搜尋文件。",
}

EXAMPLE_QUESTIONS = [
    "這份文件的主要內容是什麼？",
    "文件中提到哪些重要概念？",
    "有哪些關鍵數據或統計資料？",
    "文件的結論是什麼？",
]


def get_rag(api_key: str) -> MultiStrategyRAG:
    """在 session_state 中快取 RAG 實例（避免重複載入模型）"""
    if "rag" not in st.session_state:
        st.session_state.rag = MultiStrategyRAG(api_key=api_key)
    return st.session_state.rag


def main():
    st.set_page_config(
        page_title="多策略 RAG PDF 問答系統",
        page_icon="🤖",
        layout="wide",
    )

    # ── 標題 ─────────────────────────────────
    st.title("🤖 多策略 RAG PDF 問答系統")
    st.markdown(
        "採用 **8 種不同的 RAG 策略**，為您的 PDF 文件提供智能問答服務！"
    )
    st.divider()

    # ── API 金鑰（側邊欄） ────────────────────
    with st.sidebar:
        st.header("⚙️ 設定")
        api_key = st.text_input(
            "Groq API Key",
            value="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",
            type="password",
            help="請輸入您的 Groq API 金鑰",
        )

        st.divider()
        st.subheader("📖 策略說明")
        for name, desc in STRATEGY_DESCRIPTIONS.items():
            with st.expander(name):
                st.write(desc)

    # ── 初始化 RAG ────────────────────────────
    if not api_key:
        st.warning("⚠️ 請在側邊欄輸入 Groq API Key 後繼續。")
        st.stop()

    rag = get_rag(api_key)

    # ── 步驟 1：上傳 PDF ──────────────────────
    st.subheader("📤 步驟 1：上傳 PDF")
    uploaded_file = st.file_uploader("選擇 PDF 檔案", type=["pdf"])

    if uploaded_file is not None:
        # 只在檔名改變時重新載入
        if st.session_state.get("loaded_filename") != uploaded_file.name:
            with st.spinner("🔄 正在載入並建立索引，請稍候…"):
                status = rag.load_pdf(uploaded_file.read())
            st.session_state["loaded_filename"] = uploaded_file.name
            st.session_state["load_status"] = status

        status_msg = st.session_state.get("load_status", "")
        if "✅" in status_msg:
            st.success(status_msg)
        else:
            st.error(status_msg)

    st.divider()

    # ── 步驟 2 & 3：策略選擇 + 提問 ──────────
    col_left, col_right = st.columns([1, 2])

    with col_left:
        st.subheader("⚙️ 步驟 2：RAG 策略")
        strategy = st.selectbox(
            "選擇策略",
            list(STRATEGY_DESCRIPTIONS.keys()),
            index=0,
            label_visibility="collapsed",
        )
        st.caption(STRATEGY_DESCRIPTIONS[strategy])

        top_k = st.slider(
            "檢索片段數量（Top-K）",
            min_value=1,
            max_value=10,
            value=3,
            step=1,
        )

    with col_right:
        st.subheader("💬 步驟 3：提問")

        # 範例問題快速填入
        st.caption("💡 快速填入範例問題：")
        example_cols = st.columns(len(EXAMPLE_QUESTIONS))
        for col, q in zip(example_cols, EXAMPLE_QUESTIONS):
            if col.button(q[:10] + "…", key=f"ex_{q}", use_container_width=True, help=q):
                st.session_state["question_input"] = q

        question = st.text_area(
            "輸入您的問題",
            value=st.session_state.get("question_input", ""),
            placeholder="例如：這份文件的主要內容是什麼？",
            height=100,
            key="question_input",
        )

        ask_clicked = st.button(
            "🔍 提問",
            type="primary",
            use_container_width=True,
            disabled=(not rag.chunks),
        )

        if not rag.chunks:
            st.info("ℹ️ 請先上傳 PDF 文件，才能開始提問。")

    # ── 答案輸出 ──────────────────────────────
    if ask_clicked:
        if not question.strip():
            st.warning("⚠️ 請輸入問題後再送出。")
        else:
            with st.spinner("🧠 正在思考中，請稍候…"):
                answer, source_info = rag.generate_answer(question, strategy, top_k)

            st.divider()
            st.subheader("💡 AI 回答")

            if answer.startswith("❌"):
                st.error(answer)
            else:
                st.markdown(answer)

            if source_info:
                with st.expander("📚 查看檢索到的文本片段"):
                    st.text(source_info)


if __name__ == "__main__":
    main()