feat/add g-drive coontection

app.py

@@ -30,7 +30,11 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         index_file = gr.File(
-            file_count="multiple", file_types=["pdf"], label="Upload PDF file"
+            file_count="multiple", file_types=["pdf"], label="Upload PDF file", scale=3
+        )
+        upload_to_db = gr.CheckboxGroup(
+            ["Upload to Database"],
+            label="是否上傳至資料庫", info="將資料上傳至資料庫時，資料庫會自動建立索引，下次使用時可以直接檢索，預設為僅作這次使用", scale=1
         )
 
     with gr.Row():
@@ -42,7 +46,8 @@ with gr.Blocks() as demo:
             3. 可以根據下方的摘要內容來提問
             4. 每次對話會根據第一個問題的內容來檢索所有文件，並挑選最能回答問題的文件來回覆
             5. 要切換檢索的文件，請點選「清除」按鈕後再重新提問
-        """
+
+        """,
         )
 
     with gr.Row():
@@ -80,6 +85,7 @@ with gr.Blocks() as demo:
         **bot_args
     ).then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
 
+
     # defining workflow of clear state
     clear_state_args = dict(
         fn=clear_state,
@@ -98,7 +104,7 @@ with gr.Blocks() as demo:
 
     bulid_knowledge_base_args = dict(
         fn=build_knowledge_base,
-        inputs=[user_chatbot, index_file],
+        inputs=[user_chatbot, index_file, upload_to_db],
         outputs=None,
     )
 
@@ -118,6 +124,5 @@ with gr.Blocks() as demo:
 
     video_text_input.submit(video_bot, [test_video_chabot, video_text_input], video_text_output, api_name="video_bot")
 
-
 if __name__ == "__main__":
     demo.launch()

utils/chatbot.py

@@ -1,31 +1,62 @@
-import json
+import io
 import os
+import json
+import logging
+import secrets
 
+import gradio as gr
+import numpy as np
 import openai
 import pandas as pd
-import numpy as np
-import gradio as gr
+from google.oauth2.service_account import Credentials
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
 from openai.embeddings_utils import distances_from_embeddings
 
-from .work_flow_controller import WorkFlowController
 from .gpt_processor import QuestionAnswerer
+from .work_flow_controller import WorkFlowController
 
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+openai.api_key = OPENAI_API_KEY
 
 class Chatbot:
     def __init__(self) -> None:
         self.history = []
         self.upload_state = "waiting"
+        self.uid = self.__generate_uid()
 
+        self.g_drive_service = self.__init_drive_service()
         self.knowledge_base = None
         self.context = None
         self.context_page_num = None
         self.context_file_name = None
 
-    def build_knowledge_base(self, files):
-        work_flow_controller = WorkFlowController(files)
+    def build_knowledge_base(self, files, upload_mode="僅作這次使用"):
+        work_flow_controller = WorkFlowController(files, self.uid)
         self.csv_result_path = work_flow_controller.csv_result_path
         self.json_result_path = work_flow_controller.json_result_path
 
+        if upload_mode == "上傳至資料庫":
+            self.knowledge_base = self.__get_db_knowledge_base()
+        else:
+            self.knowledge_base = self.__get_local_knowledge_base()
+
+    def __get_db_knowledge_base(self):
+        filename = "knowledge_base.csv"
+        db = self.__read_db(self.g_drive_service)
+        cur_content = pd.read_csv(self.csv_result_path)
+        for _ in range(10):
+            try:
+                self.__write_into_db(self.g_drive_service, db, cur_content)
+                break
+            except Exception as e:
+                logging.error(e)
+                logging.error("Failed to upload to database, retrying...")
+                continue
+        self.knowledge_base = db
+        self.upload_state = "done"
+    
+    def __get_local_knowledge_base(self):
         with open(self.csv_result_path, "r", encoding="UTF-8") as fp:
             knowledge_base = pd.read_csv(fp)
         knowledge_base["page_embedding"] = (
@@ -35,10 +66,81 @@ class Chatbot:
         self.knowledge_base = knowledge_base
         self.upload_state = "done"
 
+    def __write_into_db(self, service, db: pd.DataFrame, cur_content: pd.DataFrame):
+        # db = pd.concat([db, cur_content], ignore_index=True)
+        # db.to_csv(f"{self.uid}_knowledge_base.csv", index=False)
+        cur_content.to_csv(f"{self.uid}_knowledge_base.csv", index=False)
+        media = MediaFileUpload(f"{self.uid}_knowledge_base.csv", resumable=True)
+        request = service.files().update(fileId="1m3ozrphHP221hhdCFMFX9-10nzSDfNyW", media_body=media).execute()
+    
+    def __init_drive_service(self):
+        SCOPES = ['https://www.googleapis.com/auth/drive']
+        SERVICE_ACCOUNT_FILE = os.getenv("CREDENTIALS")
+
+        creds = Credentials.from_service_account_file(SERVICE_ACCOUNT_FILE, scopes=SCOPES)
+        
+        return build('drive', 'v3', credentials=creds)
+    
+    def __read_db(self, service):
+        request = service.files().get_media(fileId="1m3ozrphHP221hhdCFMFX9-10nzSDfNyW")
+        fh = io.BytesIO()
+        downloader = MediaIoBaseDownload(fh, request)
+
+        done = False
+        while done is False:
+            status, done = downloader.next_chunk()
+            print(f"Download {int(status.progress() * 100)}%.")
+
+        # file_content = fh.getvalue().decode('utf-8')
+        fh.seek(0)
+
+        return pd.read_csv(fh)
+    
+    def __read_file(self, service, filename) -> pd.DataFrame:
+        query = f"name='{filename}'"
+        results = service.files().list(q=query).execute()
+        files = results.get('files', [])
+
+        file_id = files[0]['id']
+
+        request = service.files().get_media(fileId=file_id)
+        fh = io.BytesIO()
+        downloader = MediaIoBaseDownload(fh, request)
+
+        done = False
+        while done is False:
+            status, done = downloader.next_chunk()
+            print(f"Download {int(status.progress() * 100)}%.")
+
+        # file_content = fh.getvalue().decode('utf-8')
+        fh.seek(0)
+
+        return pd.read_csv(fh)
+    
+    def __upload_file(self, service):
+        results = service.files().list(pageSize=10).execute()
+        items = results.get('files', [])
+        if not items:
+            print('No files found.')
+        else:
+            print('Files:')
+            for item in items:
+                print(f"{item['name']} ({item['id']})")
+
+        media = MediaFileUpload(self.csv_result_path, resumable=True)
+        filename_prefix = 'ex_bot_database_'
+        filename = filename_prefix + self.uid + '.csv'
+        request = service.files().create(media_body=media, body={
+            'name': filename,
+            'parents': ["1Lp21EZlVlqL-c27VQBC6wTbUC1YpKMsG"]  # Optional, to place the file in a specific folder
+        }).execute()
+        
+
     def clear_state(self):
         self.context = None
         self.context_page_num = None
         self.context_file_name = None
+        self.knowledge_base = None
         self.upload_state = "waiting"
         self.history = []
 
@@ -130,9 +232,12 @@ class Chatbot:
             self.context_page_num = self.knowledge_base["page_num"].values[0]
             self.context_file_name = self.knowledge_base["file_name"].values[0]
 
+    def __generate_uid(self):
+        return secrets.token_hex(8)
+
+
 class VideoChatbot:
     def __init__(self) -> None:
-        openai.api_key = os.getenv("OPENAI_API_KEY")
         self.metadata_keys = ["標題", "逐字稿", "摘要", "關鍵字"]
         self.metadata = {
             "c2fK-hxnPSY":{      

utils/work_flow_controller.py

@@ -20,10 +20,11 @@ processors = {
 
 
 class WorkFlowController:
-    def __init__(self, file_src) -> None:
+    def __init__(self, file_src, uid) -> None:
         # check if the file_path is list
         # self.file_paths = self.__get_file_name(file_src)
         self.file_paths = [x.name for x in file_src]
+        self.uid = uid
 
         print(self.file_paths)
 
@@ -83,6 +84,7 @@ class WorkFlowController:
 
         for i, _ in enumerate(file["file_content"]):
             # use i+1 to meet the index of file_content
+            print("Translating page: " + str(i + 1))
             file["file_content"][i + 1][
                 "page_content"
             ] = translator.translate_to_chinese(
@@ -97,33 +99,34 @@ class WorkFlowController:
         # process file content
         # return processed data
         if not file["is_chinese"]:
+            print("Translating to chinese...")
             file = self.__translate_to_chinese(file)
+        print("Getting embedding...")
         file = self.__get_embedding(file)
+        print("Getting summary...")
         file = self.__get_summary(file)
         return file
 
     def __dump_to_json(self):
         with open(
-            os.path.join(os.getcwd(), "knowledge_base.json"), "w", encoding="utf-8"
+            os.path.join(os.getcwd(), f"{self.uid}_knowledge_base.json"), "w", encoding="utf-8"
         ) as f:
             print(
                 "Dumping to json, the path is: "
-                + os.path.join(os.getcwd(), "knowledge_base.json")
+                + os.path.join(os.getcwd(), f"{self.uid}_knowledge_base.json")
             )
-            self.json_result_path = os.path.join(os.getcwd(), "knowledge_base.json")
+            self.json_result_path = os.path.join(os.getcwd(), f"{self.uid}_knowledge_base.json")
             json.dump(self.files_info, f, indent=4, ensure_ascii=False)
 
     def __construct_knowledge_base_dataframe(self):
         rows = []
         for file_path, content in self.files_info.items():
-            file_full_content = content["file_full_content"]
             for page_num, page_details in content["file_content"].items():
                 row = {
                     "file_name": content["file_name"],
                     "page_num": page_details["page_num"],
                     "page_content": page_details["page_content"],
                     "page_embedding": page_details["page_embedding"],
-                    "file_full_content": file_full_content,
                 }
                 rows.append(row)
 
@@ -132,19 +135,18 @@ class WorkFlowController:
             "page_num",
             "page_content",
             "page_embedding",
-            "file_full_content",
         ]
         df = pd.DataFrame(rows, columns=columns)
         return df
 
     def __dump_to_csv(self):
         df = self.__construct_knowledge_base_dataframe()
-        df.to_csv(os.path.join(os.getcwd(), "knowledge_base.csv"), index=False)
+        df.to_csv(os.path.join(os.getcwd(), f"{self.uid}_knowledge_base.csv"), index=False)
         print(
             "Dumping to csv, the path is: "
-            + os.path.join(os.getcwd(), "knowledge_base.csv")
+            + os.path.join(os.getcwd(), f"{self.uid}_knowledge_base.csv")
         )
-        self.csv_result_path = os.path.join(os.getcwd(), "knowledge_base.csv")
+        self.csv_result_path = os.path.join(os.getcwd(), f"{self.uid}_knowledge_base.csv")
 
     def __get_file_name(self, file_src):
         file_paths = [x.name for x in file_src]