Update app.py

app.py

@@ -4,40 +4,29 @@ import gradio as gr
 from langchain.chains import ConversationalRetrievalChain
 from langchain.text_splitter import CharacterTextSplitter
 from langchain_community.document_loaders import PyMuPDFLoader, PyPDFLoader
-from langchain_community.vectorstores import Chroma
+from langchain.vectorstores import Chroma
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.chat_models import ChatOpenAI
-from dotenv import load_dotenv
-
-# Load environment variables
-load_dotenv()
-os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
-
-# Validate OpenAI API Key
-api_key = os.getenv('OPENAI_API_KEY')
-if not api_key:
-    raise ValueError("Please set the 'OPENAI_API_KEY' environment variable")
-
-# OpenAI API key
-openai_api_key = api_key
-
-# Transform chat history for LangChain format
-def transform_history_for_langchain(history):
-    return [(chat[0], chat[1]) for chat in history if chat[0]]
-
-# Transform chat history for OpenAI format
-def transform_history_for_openai(history):
-    new_history = []
-    for chat in history:
-        if chat[0]:
-            new_history.append({"role": "user", "content": chat[0]})
-        if chat[1]:
-            new_history.append({"role": "assistant", "content": chat[1]})
-    return new_history
-
-# Load and process documents function
-def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader'):
+import shutil  # 用於文件複製
+
+# 獲取 OpenAI API 密鑰（初始不使用固定密鑰）
+api_key_env = os.getenv("OPENAI_API_KEY")
+if api_key_env:
+    openai.api_key = api_key_env
+else:
+    print("未設置固定的 OpenAI API 密鑰。將使用使用者提供的密鑰。")
+
+# 確保向量資料庫目錄存在且有寫入權限
+VECTORDB_DIR = os.path.abspath("./data")
+os.makedirs(VECTORDB_DIR, exist_ok=True)
+os.chmod(VECTORDB_DIR, 0o755)
+
+# 定義載入和處理 PDF 文件的函數
+def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader', api_key=None):
+    if not api_key:
+        raise ValueError("未提供 OpenAI API 密鑰。")
     documents = []
+
     for file_path in file_paths:
         if not os.path.exists(file_path):
             continue
@@ -55,99 +44,165 @@ def load_and_process_documents(file_paths, loader_type='PyMuPDFLoader'):
             continue
 
     if not documents:
-        raise ValueError("No documents found or could not load any documents.")
+        raise ValueError("沒有找到任何 PDF 文件或 PDF 文件無法載入。")
 
-    # Split long texts
+    # 分割長文本
     text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
     documents = text_splitter.split_documents(documents)
 
     if not documents:
-        raise ValueError("Document list is empty after splitting. Please check the content of the files.")
+        raise ValueError("分割後的文檔列表為空。請檢查 PDF 文件內容。")
 
-    # Initialize vector database
-    vectordb = Chroma.from_documents(
-        documents,
-        embedding=OpenAIEmbeddings(),
-        persist_directory="./tmp"
-    )
-    return vectordb
+    # 初始化向量資料庫
+    try:
+        embeddings = OpenAIEmbeddings(openai_api_key=api_key)  # 使用使用者的 API 密鑰
+    except Exception as e:
+        raise ValueError(f"初始化 OpenAIEmbeddings 時出現錯誤: {e}")
 
-# Initialize vector database as a global variable
-if 'vectordb' not in globals():
-    vectordb = load_and_process_documents(["./sample_docs/sample.pdf"])
+    try:
+        vectordb = Chroma.from_documents(
+            documents,
+            embedding=embeddings,
+            persist_directory=VECTORDB_DIR
+        )
+    except Exception as e:
+        raise ValueError(f"初始化 Chroma 向量資料庫時出現錯誤: {e}")
+
+    return vectordb
 
-# Define query handling function for RAG
-def handle_query(user_message, temperature, chat_history):
+# 定義聊天處理函數
+def handle_query(user_message, chat_history, vectordb, api_key):
     try:
         if not user_message:
-            return chat_history  # Return unchanged chat history
+            return chat_history
 
-        # Use LangChain's ConversationalRetrievalChain to handle the query
+        # 添加角色指令前綴
         preface = """
 Instruction: Answer in Traditional Chinese, within 200 characters.
 If the question is unrelated to the documents, respond with: 此事無可奉告,話說這件事須請教海虔王...
 """
-        query = f"{preface} Query content: {user_message}"
-
-        # Extract previous answers as context, converting them to LangChain format
-        previous_answers = transform_history_for_langchain(chat_history)
+        query = f"{preface} 查詢內容：{user_message}"
 
+        # 初始化 ConversationalRetrievalChain，並傳遞 openai_api_key
         pdf_qa = ConversationalRetrievalChain.from_llm(
-            ChatOpenAI(temperature=temperature, model_name='gpt-4'),
+            ChatOpenAI(temperature=0.7, model="gpt-4", openai_api_key=api_key),
             retriever=vectordb.as_retriever(search_kwargs={'k': 6}),
-            return_source_documents=True,
-            verbose=False
+            return_source_documents=True
         )
 
-        # Invoke the model to handle the query
-        result = pdf_qa.invoke({"question": query, "chat_history": previous_answers})
-
-        # Ensure 'answer' is present in the result
-        if "answer" not in result:
-            return chat_history + [("System", "Sorry, an error occurred.")]
-
-        # Update the AI response in chat history
-        chat_history[-1] = (user_message, result["answer"])  # Update the last record, pairing user input with AI response
+        # 呼叫模型並處理查詢
+        result = pdf_qa.invoke({"question": query, "chat_history": chat_history})
 
+        if "answer" in result:
+            chat_history = chat_history + [(user_message, result["answer"])]
+        else:
+            chat_history = chat_history + [(user_message, "抱歉，未能獲得有效回應。")] 
         return chat_history
 
     except Exception as e:
-        return chat_history + [("System", f"An error occurred: {str(e)}")]
-
-# Create a custom chat interface using Gradio Blocks API
+        return chat_history + [("系統", f"出現錯誤: {str(e)}")]
+
+# 定義保存 API 密鑰的函數
+def save_api_key(api_key, state):
+    if not api_key.startswith("sk-"):
+        return "請輸入有效的 OpenAI API 密鑰。", state
+    state['api_key'] = api_key
+    return "API 密鑰已成功保存。您現在可以上傳 PDF 文件並開始提問。", state
+
+# 定義 Gradio 的處理函數
+def process_files(files, state):
+    if files:
+        try:
+            api_key = state.get('api_key', None)
+            if not api_key:
+                return "請先輸入並保存您的 OpenAI API 密鑰。", state
+
+            saved_file_paths = []
+            for idx, file_data in enumerate(files):
+                filename = f"uploaded_{idx}.pdf"
+                save_path = os.path.join(VECTORDB_DIR, filename)
+                with open(save_path, "wb") as f:
+                    f.write(file_data)
+                saved_file_paths.append(save_path)
+            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:
+            return f"處理文件時出現錯誤: {e}", 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, "請先輸入並保存您的 OpenAI API 密鑰。"
+
+    updated_history = handle_query(user_message, chat_history, vectordb, api_key)
+    return updated_history, state, ""
+
+# 設計 Gradio 介面
 with gr.Blocks(css="body { background-color: #EBD6D6; }") as demo:
-    gr.Markdown("<h1 style='text-align: center;'>AI Assistant for AI Forum</h1>")
+    gr.Markdown("<h1 style='text-align: center;'>AI論壇助理</h1>")
+
+    state = gr.State({"vectordb": None, "api_key": None})
+
+    # API 密鑰輸入框
+    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)
+
+    # 上傳 PDF 文件
+    gr.Markdown("<span style='font-size: 1.5em; font-weight: bold;'>請上傳AI論壇相關文檔，提供AI相關問題解答</span>")
+    upload = gr.File(
+        file_count="multiple",
+        file_types=[".pdf"],
+        label="上傳AI論壇 PDF 文件",
+        interactive=True,
+        type="binary"
+    )
+    upload_btn = gr.Button("上傳並處理")
+    upload_status = gr.Textbox(label="上傳狀態", interactive=False)
 
+    # 智能諮詢
+    gr.Markdown("### AI論壇助理")
     chatbot = gr.Chatbot()
-    state = gr.State([])
-
-    with gr.Row():
-        with gr.Column(scale=0.85):
-            txt = gr.Textbox(show_label=False, placeholder="Please enter your question...")
-        with gr.Column(scale=0.15, min_width=0):
-            submit_btn = gr.Button("Ask")
-
-    # Immediately show user input without response part, and clear input box
-    def user_input(user_message, history):
-        history.append((user_message, ""))  # Show user message, response part as empty string
-        return history, "", history  # Return cleared input box and updated chat history
-
-    # Handle AI response, update response part
-    def bot_response(history):
-        user_message = history[-1][0]  # Get the latest user input
-        history = handle_query(user_message, 0.7, history)  # Call the query handler
-        return history, history  # Return updated chat history
-
-    # First show user message, then handle AI response, clear input box
-    submit_btn.click(user_input, [txt, state], [chatbot, txt, state], queue=False).then(
-        bot_response, state, [chatbot, state]
+
+    txt = gr.Textbox(show_label=False, placeholder="請輸入您的AI問題...")
+    submit_btn = gr.Button("提問")
+
+    # 綁定事件
+    save_api_key_btn.click(
+        save_api_key,
+        inputs=[api_key_input, state],
+        outputs=[api_key_status, state]
+    )
+
+    upload_btn.click(
+        process_files,
+        inputs=[upload, state],
+        outputs=[upload_status, state]
+    )
+
+    submit_btn.click(
+        chat_interface,
+        inputs=[txt, chatbot, state],
+        outputs=[chatbot, state, txt]
     )
 
-    # Support pressing "Enter" to submit question, immediately show user input, clear input box
-    txt.submit(user_input, [txt, state], [chatbot, txt, state], queue=False).then(
-        bot_response, state, [chatbot, state]
+    txt.submit(
+        chat_interface,
+        inputs=[txt, chatbot, state],
+        outputs=[chatbot, state, txt]
     )
 
-# Launch Gradio app
+# 啟動 Gradio 應用
 demo.launch()