Update main.py

main.py

@@ -1,34 +1,38 @@
+from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi.security import APIKeyHeader
+from pydantic import BaseModel
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
-import PyPDF2
-import os
-import gradio as gr
-import google.generativeai as genai
 from langchain.chains import ConversationalRetrievalChain
-from langchain_huggingface import HuggingFaceEmbeddings
+import pdfplumber
+import os
+
+# 初始化 FastAPI 應用
+app = FastAPI()
 
 print('程式初始化')
 
 # 設定 Google API 金鑰
-genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
+api_key = os.getenv("GOOGLE_API_KEY")
+if not api_key:
+    raise ValueError("GOOGLE_API_KEY is not set")
+genai.configure(api_key=api_key)
 
 # 選擇模型
 llm_model = 'gemini-1.5-flash'
 embeddings_model = "models/embedding-001"
 pdf_dir = 'data'
 
-# 讀取 PDF 檔案
+# 讀取 PDF 檔案並初始化資料
 print('-' * 21, '讀取資料', '-' * 21)
 docs = ""
 for filename in os.listdir(pdf_dir):
     if filename.endswith('.pdf'):
         print(filename)
-        with open(os.path.join(pdf_dir, filename), 'rb') as pdf_file:
-            pdf_reader = PyPDF2.PdfReader(pdf_file)
-            for i in range(len(pdf_reader.pages)):  
-                page = pdf_reader.pages[i]  
-                docs += page.extract_text() 
+        with pdfplumber.open(os.path.join(pdf_dir, filename)) as pdf:
+            for page in pdf.pages:
+                docs += page.extract_text()
 
 print('-' * 21, '讀取完成', '-' * 21)
 
@@ -39,40 +43,72 @@ if docs:
 
     # 建立嵌入模型和檢索器
     embeddings = GoogleGenerativeAIEmbeddings(
-        model=embeddings_model, google_api_key=os.getenv("GOOGLE_API_KEY")
+        model=embeddings_model, google_api_key=api_key
     )
     retriever = Chroma.from_texts(texts, embeddings).as_retriever(search_kwargs={"k": 1})
     print('分割文本完成')
 
     # 初始化 Gemini 模型
     llm = ChatGoogleGenerativeAI(
-        model=llm_model, temperature=0.1, google_api_key=os.getenv("GOOGLE_API_KEY")
+        model=llm_model, temperature=0.1, google_api_key=api_key
     )
     print('模型載入完成')
 
+else:
+    raise ValueError("No documents found in the 'data' directory.")
 
-# 定義 invoke 函數
+# 初始化 chat_history 作為全局變數
 chat_history = []
-def invoke(question):
+
+# 定義 invoke 函數
+def invoke(question: str):
     print('invoke 函數觸發')
-    if docs:
-        system_prompt = (
-            "You are an assistant for question-answering tasks. "
-            "Use the following pieces of retrieved context to answer the question. "
-            
-        )
-        #"If you don't know the answer, say that you don't know."
+    try:
         # 初始化 ConversationalRetrievalChain
         qa_chain = ConversationalRetrievalChain.from_llm(
             llm=llm, retriever=retriever
         )
 
-        # 调用链并传递 chat_history
+        # 呼叫 QA chain 並處理回應
         response = qa_chain.invoke({"question": question, "chat_history": chat_history})
+        return response['answer']
+
+    except Exception as e:
+        print(f"Error during invoke: {e}")
+        return "An error occurred while processing the request."
+
+
+# 定義請求資料模型
+class InvokeRequest(BaseModel):
+    question: str
+
+# API 金鑰設定
+API_KEY_NAME = "X-API-Key"
+API_KEY = os.getenv("API_KEY", "your_secret_api_key")
+api_key_header = APIKeyHeader(name=API_KEY_NAME)
+
+# 定義 API 金鑰驗證的依賴函數
+def get_api_key(api_key: str = Depends(api_key_header)):
+    if api_key != API_KEY:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Invalid API Key",
+        )
+    return api_key
+
+# 根端點，用於檢查服務是否運行正常
+@app.get("/")
+async def read_root():
+    return {"message": "Hello, World!"}
+
+# 定義 POST 端點，讓使用者發送問題並獲得模型回應
+@app.post("/invoke")
+async def rag_endpoint(request: InvokeRequest, api_key: str = Depends(get_api_key)):
+    result = invoke(request.question)
+    return {"result": result}
+
+# 啟動應用程式：執行命令 `uvicorn main:app --reload`
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
 
-        # 更新 chat_history，保留上下文
-        # chat_history.append((question, response['answer']))
-    else:
-        response = 'No context!'
-    
-    return response['answer']