app.py

import os
import re
import gradio as gr
import aisuite as ai
import zipfile
import time # 為了讓檔案解壓縮有時間

# RAG 相關
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

from huggingface_hub import login
# 移除 Colab 專用的 userdata 匯入

# =========================================================
# 1. 檔案解壓縮與環境變數載入 (Hugging Face Secrets 會自動注入 OS 環境變數)
# =========================================================

# 1.0 透過 Token 登入 Hugging Face (【新增的關鍵程式碼】)
# =========================================================
# 1. 檔案解壓縮與環境變數載入 (Hugging Face Secrets 會自動注入 OS 環境變數)
# =========================================================

# 1.0 透過 Token 登入 Hugging Face (【關鍵修正區塊】)
groq_api_key = os.environ.get('GROQ')
hf_token = os.environ.get('HuggingFace')

if hf_token:
    try:
        login(token=hf_token)
        print("✅ Hugging Face 登入成功，已驗證模型存取權。")
    except Exception as e:
        print(f"❌ 警告：Hugging Face Token 登入失敗。請檢查 Secrets 裡面的 'HuggingFace' Token 是否有讀取權限或是否過期。錯誤: {e}")

if not groq_api_key:
    print("❌ 警告：GROQ API Key 未設定在 Hugging Face Secrets 中！")
    
# 由於 aisuite 依賴 os.environ，我們在這裡確保 Groq Key 再次被設定
os.environ['GROQ_API_KEY'] = groq_api_key if groq_api_key else ""

# 1.1 啟動時解壓縮 FAISS 資料庫

# 1.1 啟動時解壓縮 FAISS 資料庫
if not os.path.exists("faiss_dbV2"):
    print("正在解壓縮 faiss_dbV2.zip...")
    try:
        with zipfile.ZipFile("faiss_dbV2.zip", 'r') as zip_ref:
            zip_ref.extractall(".")
        print("✅ faiss_db 解壓縮完成。")
    except FileNotFoundError:
        print("❌ 錯誤：未找到 faiss_dbV2.zip，請檢查是否已上傳。")
        exit()

# 1.2 檢查 API Key
#groq_api_key = os.environ.get('GROQ') # 注意：從 Secrets 讀取時 Key 要大寫，如 GROQ
#hf_token = os.environ.get('HuggingFace')

if not groq_api_key:
    print("❌ 警告：GROQ API Key 未設定在 Hugging Face Secrets 中！")
if not hf_token:
    print("❌ 警告：HuggingFace Token 未設定在 Secrets 中！")

# 由於 aisuite 依賴 os.environ，我們在這裡確保 Groq Key 再次被設定
os.environ['GROQ_API_KEY'] = groq_api_key if groq_api_key else ""


# =========================================================
# 2. 模型與設定 (與 Colab 相同)
# =========================================================

# Embedding 模型配置
class EmbeddingGemmaEmbeddings(HuggingFaceEmbeddings):
    def __init__(self, **kwargs):
        super().__init__(
            model_name="google/embeddinggemma-300m",
            encode_kwargs={"normalize_embeddings": True},
            **kwargs
        )

    def embed_documents(self, texts):
        texts = [f"title: none | text: {t}" for t in texts]
        return super().embed_documents(texts)

    def embed_query(self, text):
        return super().embed_query(f"task: search result | query: {text}")

embedding_model = EmbeddingGemmaEmbeddings()

# FAISS 資料庫載入
print("正在載入 FAISS 資料庫...")
vectorstore = FAISS.load_local(
    "faiss_db",
    embeddings=embedding_model,
    allow_dangerous_deserialization=True
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 8})
print("✅ RAG 系統初始化完成。")


# LLM Agent 配置 (使用您最後成功的配置)
provider_system_planner = "groq"
model_system_planner = "llama-3.3-70b-versatile"
provider_system_writer = "groq"
model_system_writer = "llama-3.3-70b-versatile"
provider_system_reviewer = "groq"
model_system_reviewer = "openai/gpt-oss-120b"
provider_system_refiner = "groq"
model_system_refiner = "llama-3.3-70b-versatile"

# 初始化 AI Client
client = ai.Client()


# =========================================================
# 3. 核心函式定義 (與 Colab 相同)
# =========================================================

def call_ai(system_prompt, user_prompt, provider, model):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    try:
        # **注意**：我們在這裡使用了 os.environ 裡面的 Key，
        # 如果 Key 失敗，這裡就會報錯。
        response = client.chat.completions.create(
            model=f"{provider}:{model}",
            messages=messages
        )
        return response.choices[0].message.content
    except Exception as e:
        # 如果模型呼叫失敗，拋出詳細錯誤
        raise Exception(f"【AI Agent 呼叫失敗: {provider}:{model}】 錯誤原因: {str(e)}")


# A. CoT Planner: 發想 5 種思路
system_planner = """
你是一位精通實變函數論 (Real Analysis) 的數學家。
你的任務不是直接給答案，而是進行「廣度優先思考」。
針對使用者的問題與提供的課本內容，請提出 5 種不同的解釋路徑或證明策略，盡可能詳細證明。
"""
# B. CoT Selector & Writer: 選擇最佳並撰寫初稿
system_writer = """
你是一位嚴謹的數學教授。
請參考 Planner 提供的 5 種思路，選出「最符合實變函數論嚴謹定義」的一種。
請依此思路撰寫完整的解答：
1. 必須使用 LaTeX 格式書寫所有數學符號。
2. 【關鍵規則】：所有行內公式 (Inline Math) 的 $ 符號前後必須留一個空格。
3. 【關鍵規則】：複雜或長度超過 10 個字元的公式，請務必使用雙錢字號 $$ ... $$ 獨立一行顯示。
4. 【新增要求】：在 $$...$$ 區塊公式的**上下**必須各留**一個空行**，以確保 Markdown 渲染器能正確識別。
5. 引用相關的定義與定理 (Definition/Theorem)。
6. 使用繁體中文 (台灣習慣) 進行解說。
"""
# C. Reflexion Reviewer: 2diff 審查
system_reviewer = """
你是一位負責嚴格數學審稿人。
請檢查這份數學解答。
檢查重點：
1. 定義引用是否正確？
2. 證明邏輯是否有漏洞 (Logical Gaps)？
3. 符號使用是否規範？
4. 是否有濫用直觀而犧牲嚴謹性的情況？
請給出具體的修改建議。如果解答正確，請回答「PASS」。
"""
# D. Reflexion Refiner: 根據建議修正
system_refiner = """
你負責根據審稿人的建議，修正並潤飾最終的數學解答。
請保留原本正確的部分，針對被指出的錯誤進行修正。
1. 必須使用 LaTeX 格式書寫所有數學符號。
2. 【關鍵規則】：所有行內公式 (Inline Math) 的 $ 符號前後必須留一個空格。
3. 【關鍵規則】：複雜或長度超過 10 個字元的公式，請務必使用雙錢字號 $$ ... $$ 獨立一行顯示。
4. 【新增要求】：在 $$...$$ 區塊公式的**上下**必須各留**一個空行**，以確保 Markdown 渲染器能正確識別。
"""

def fix_latex(text):
    if not text: return text
    # 1. 取代區塊公式 \[ ... \] 為 $$ ... $$
    text = re.sub(r'\\\[(.*?)\\\]', r'\n$$\1$$\n', text, flags=re.DOTALL)
    # 2. 取代行內公式 \( ... \) 為 $ ... $
    text = re.sub(r'\\\((.*?)\\\)', r'$ \1 $', text, flags=re.DOTALL)
    return text


def math_solver_process(user_question):
    try:
        # --- Step 1: RAG 檢索 ---
        docs = retriever.invoke(user_question)
        context = "\n\n".join([doc.page_content for doc in docs])
        rag_info = fix_latex(f"【檢索到的課本內容】:\n{context}")

        # --- Step 2: CoT (Planner) ---
        plan_prompt = f"使用者問題: {user_question}\n\n參考資料:\n{context}\n\n請列出 5 種解題思路。"
        thoughts_5 = fix_latex(call_ai(system_planner, plan_prompt, provider_system_planner, model_system_planner))

        # --- Step 3: Draft (Writer) ---
        draft_prompt = f"使用者問題: {user_question}\n\nPlanner 的思路:\n{thoughts_5}\n\n參考資料:\n{context}\n\n請選擇最佳思路並撰寫解答初稿。"
        first_draft = fix_latex(call_ai(system_writer, draft_prompt, provider_system_writer, model_system_writer))

        # --- Step 4: Reflexion (Reviewer) ---
        review_prompt = f"問題: {user_question}\n\n初稿解答:\n{first_draft}\n\n請審查並給出建議。"
        review_feedback = fix_latex(call_ai(system_reviewer, review_prompt, provider_system_reviewer, model_system_reviewer))

        # --- Step 5: Final Refinement (Refiner) ---
        if "PASS" in review_feedback:
            raw_answer = first_draft
            status = "✅ 審查通過，初稿即完美。"
        else:
            refine_prompt = f"初稿:\n{first_draft}\n\n審查建議:\n{review_feedback}\n\n請產出最終修正版解答。"
            raw_answer = call_ai(system_refiner, refine_prompt, provider_system_refiner, model_system_refiner)
            status = "🔄 經過審查修正後的最終版。"

        final_answer = fix_latex(raw_answer)
        
        return rag_info, thoughts_5, first_draft, review_feedback, final_answer
    
    except Exception as e:
        error_msg = f"系統執行發生錯誤:\n{str(e)}"
        # 如果發生錯誤，將錯誤訊息回傳給 Gradio 介面
        return error_msg, "Error", "Error", "Error", "Error"


    # =========================================================
# 4. Gradio 介面配置 (作為測試介面，確保能跑)
# =========================================================

# 定義 LaTeX 渲染模式（統一所有 Markdown 元件）
LATEX_DELIMITERS = [
    { "left": "$$", "right": "$$", "display": True },
    { "left": "$", "right": "$", "display": False }
]

with gr.Blocks(title="實變函數論 AI Agent") as demo:
    gr.Markdown("# 📐 實變函數論 AI 助教 (Hugging Face Spaces RAG + CoT)")
    gr.Markdown("輸入你的數學問題，AI 將查詢課本、多角度思考、並經由同儕審查機制產出嚴謹證明。")

    with gr.Row():
        input_box = gr.Textbox(label="請輸入問題", placeholder="例如：請解釋 Lebesgue 積分與 Riemann 積分在極限交換時的差異")
        run_btn = gr.Button("開始解題", variant="primary")

    # 步驟一：檢索資料 (RAG) - 使用 Markdown 渲染 LaTeX
    with gr.Accordion("📚 步驟一：檢索資料 (RAG Context)", open=False):
        out_rag = gr.Markdown(
            label="RAG Context", 
            latex_delimiters=LATEX_DELIMITERS
        )

    # 步驟二：思考與選擇 (CoT) - 使用 Markdown 渲染 LaTeX
    with gr.Accordion("🧠 步驟二：思考與選擇 (CoT Planner)", open=False):
        out_thoughts = gr.Markdown(
            label="5 種解題思路", 
            latex_delimiters=LATEX_DELIMITERS
        )

    # 步驟三、四：初稿與審查 - 包裹在一個新的可收合區塊
    with gr.Accordion("📝 步驟三/四：初稿與審查 (Writer / Reviewer)", open=False):
        with gr.Row():
            # 初稿 (Writer) - 使用 Markdown 渲染 LaTeX
            out_draft = gr.Markdown(
                label="初稿 (Writer)", 
                latex_delimiters=LATEX_DELIMITERS
            )
            # 審查意見 (Reviewer) - 使用 Markdown 渲染 LaTeX
            out_review = gr.Markdown(
                label="審查意見 (Reviewer)", 
                latex_delimiters=LATEX_DELIMITERS
            )

    gr.Markdown("### ✨ 最終解答 (Final Answer)")

    # 最終解答 (Final Answer) - 維持不變，確保在最顯眼的位置
    out_final = gr.Markdown(
        label="最終解答",
        latex_delimiters=LATEX_DELIMITERS
    )

    # 按鈕事件綁定 (注意：math_solver_process 必須回傳 5 個輸出)
    run_btn.click(
        math_solver_process,
        inputs=[input_box],
        outputs=[out_rag, out_thoughts, out_draft, out_review, out_final] # 輸出數量與函數回傳匹配
    )

    # 在 Hugging Face 上執行時，只需呼叫 demo.launch()
    demo.launch()