Update app.py

app.py

@@ -9,27 +9,20 @@ from langchain_community.document_loaders import PyMuPDFLoader, PyPDFLoader
 from langchain.vectorstores import Chroma
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.chat_models import ChatOpenAI
-import shutil  # 用於文件複製
+import shutil
 import logging
 
-# 設置日誌配置
+# 設置日誌
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 獲取 OpenAI API 密鑰
-api_key = os.getenv("OPENAI_API_KEY")
-if not api_key:
-    raise ValueError("未能獲取 OPENAI_API_KEY。請在 Hugging Face Spaces 的 Secrets 中設置它。")
-openai.api_key = api_key
-logger.info("OpenAI API 密鑰已設置。")
-
-# 確保向量資料庫目錄存在且有寫入權限
+# 向量資料庫目錄
 VECTORDB_DIR = os.path.abspath("./data")
 os.makedirs(VECTORDB_DIR, exist_ok=True)
-os.chmod(VECTORDB_DIR, 0o755)  # 設置適當的權限
+os.chmod(VECTORDB_DIR, 0o755)
 logger.info(f"VECTORDB_DIR set to: {VECTORDB_DIR}")
 
-# 定義測試 PDF 加載器的函數
+# 測試 PDF 加載器
 def test_pdf_loader(file_path, loader_type='PyMuPDFLoader'):
     logger.info(f"Testing PDF loader ({loader_type}) with file: {file_path}")
     try:
@@ -49,8 +42,10 @@ def test_pdf_loader(file_path, loader_type='PyMuPDFLoader'):
     except Exception as e:
         logger.error(f"Error loading {file_path} with {loader_type}: {e}")
 
-# 定義載入和處理 PDF 文件的函數
-def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader'):
+# 載入和處理 PDF 文件
+def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader', api_key=None):
+    if not api_key:
+        raise ValueError("未提供 OpenAI API 密鑰。")
     documents = []
     logger.info("開始載入上傳的 PDF 文件。")
 
@@ -70,7 +65,6 @@ def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader'):
             loaded_docs = loader.load()
             if loaded_docs:
                 logger.info(f"載入 {file_path} 成功，包含 {len(loaded_docs)} 個文檔。")
-                # 打印第一個文檔的部分內容以確認
                 logger.info(f"第一個文檔內容: {loaded_docs[0].page_content[:500]}")
                 documents.extend(loaded_docs)
             else:
@@ -93,7 +87,7 @@ def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader'):
 
     # 初始化向量資料庫
     try:
-        embeddings = OpenAIEmbeddings(openai_api_key=api_key)  # 直接傳遞 API 密鑰
+        embeddings = OpenAIEmbeddings(openai_api_key=api_key)
         logger.info("初始化 OpenAIEmbeddings 成功。")
     except Exception as e:
         raise ValueError(f"初始化 OpenAIEmbeddings 時出現錯誤: {e}")
@@ -110,8 +104,8 @@ def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader'):
 
     return vectordb
 
-# 定義聊天處理函數
-def handle_query(user_message, chat_history, vectordb):
+# 聊天處理函數
+def handle_query(user_message, chat_history, vectordb, api_key):
     try:
         if not user_message:
             return chat_history
@@ -144,11 +138,32 @@ def handle_query(user_message, chat_history, vectordb):
         logger.error(f"Error in handle_query: {e}")
         return chat_history + [("系統", f"出現錯誤: {str(e)}")]
 
-# 定義 Gradio 的處理函數
+# 保存 API 密鑰的函數
+def save_api_key(api_key, state):
+    if not api_key.startswith("sk-"):
+        return "請輸入有效的 OpenAI API 密鑰。", state
+    # 嘗試驗證 API 密鑰
+    try:
+        openai.api_key = api_key
+        openai.Engine.list()  # 簡單的 API 請求來驗證密鑰
+    except Exception as e:
+        logger.error(f"Invalid OpenAI API key: {e}")
+        return "無效的 OpenAI API 密鑰。請重新輸入。", state
+    state['api_key'] = api_key
+    logger.info("使用者已保存自己的 OpenAI API 密鑰。")
+    return "API 密鑰已成功保存。您現在可以上傳 PDF 文件並開始提問。", state
+
+# 處理文件上傳的函數
 def process_files(files, state):
     logger.info("process_files called")
     if files:
         try:
+            # 檢查是否已保存 API 密鑰
+            api_key = state.get('api_key', None)
+            if not api_key:
+                logger.error("使用者未提供 OpenAI API 密鑰。")
+                return "請先在「設定 API 密鑰」標籤中輸入並保存您的 OpenAI API 密鑰。", state
+
             logger.info(f"Received {len(files)} files")
             saved_file_paths = []
             for idx, file_data in enumerate(files):
@@ -182,7 +197,7 @@ def process_files(files, state):
             # 列出文件大小
             file_sizes = {file: os.path.getsize(os.path.join(VECTORDB_DIR, file)) for file in saved_files}
             logger.info(f"File sizes in VECTORDB_DIR: {file_sizes}")
-            vectordb = load_and_process_documents(saved_file_paths, loader_type='PyMuPDFLoader')
+            vectordb = load_and_process_documents(saved_file_paths, loader_type='PyMuPDFLoader', api_key=api_key)
             state['vectordb'] = vectordb
             return "PDF 文件已成功上傳並處理。您現在可以開始提問。", state
         except Exception as e:
@@ -191,13 +206,17 @@ def process_files(files, state):
     else:
         return "請上傳至少一個 PDF 文件。", state
 
+# 聊天介面處理函數
 def chat_interface(user_message, chat_history, state):
     vectordb = state.get('vectordb', None)
+    api_key = state.get('api_key', None)
     if not vectordb:
         return chat_history, state, "請先上傳 PDF 文件以進行處理。"
+    if not api_key:
+        return chat_history, state, "請先在「設定 API 密鑰」標籤中輸入並保存您的 OpenAI API 密鑰。"
 
     # 處理查詢
-    updated_history = handle_query(user_message, chat_history, vectordb)
+    updated_history = handle_query(user_message, chat_history, vectordb, api_key)
     return updated_history, state, ""
 
 # 設計 Gradio 介面
@@ -205,7 +224,19 @@ with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center;'>MBTI 與經典調酒 AI 助理</h1>")
 
     # 定義共享的 state
-    state = gr.State({"vectordb": None})
+    state = gr.State({"vectordb": None, "api_key": None})
+
+    with gr.Tab("設定 API 密鑰"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                api_key_input = gr.Textbox(
+                    label="輸入您的 OpenAI API 密鑰",
+                    placeholder="sk-...",
+                    type="password",
+                    interactive=True
+                )
+                save_api_key_btn = gr.Button("保存 API 密鑰")
+                api_key_status = gr.Textbox(label="狀態", interactive=False)
 
     with gr.Tab("上傳 PDF 文件"):
         with gr.Row():
@@ -243,6 +274,13 @@ with gr.Blocks() as demo:
             outputs=[chatbot, state, txt]
         )
 
+    # 綁定保存 API 密鑰按鈕
+    save_api_key_btn.click(
+        save_api_key,
+        inputs=[api_key_input, state],
+        outputs=[api_key_status, state]
+    )
+
     # 綁定上傳按鈕
     upload_btn.click(
         process_files,