refactor: updated drive upload to use browser localstorage to store id for uploading chats

.gitignore

@@ -1,9 +1,8 @@
 venv
 .env
-log.csv
-utils/__pycache__/doc_utils.cpython-312.pyc
-utils/__pycache__/utils.cpython-312.pyc
+.gradio
+__pycache__
+utils/__pycache__
 test.py
 data/old_KB.py
 creds.json
-utils/__pycache__/google_drive_utils.cpython-312.pyc

app.py

@@ -1,6 +1,7 @@
 import os
 import json
 import base64
+import uuid
 import logging
 
 from typing import Generator
@@ -16,7 +17,8 @@ from utils.utils import (
     get_messages,
     load_knowledge_base,
 )
-from utils.chatLogger import ChatLogger
+
+from utils.chatLogger import ChatUploader
 
 
 def initialize():
@@ -40,9 +42,6 @@ def initialize():
     logger.info("Loading env variables...")
     if not os.getenv("SPACE_ID"):
         load_dotenv()
-    drive_creds_encoded = os.getenv(
-        "GOOGLE_DRIVE_SERVICE_ACCOUNT_CREDENTIALS_BASE64"
-    ).strip()
 
     logger.info("Initializing OpenAI client...")
     openAI_client = OpenAI(
@@ -51,10 +50,13 @@ def initialize():
     )
 
     logger.info("Loading Drive service account details...")
+    drive_creds_encoded = os.getenv(
+        "GOOGLE_DRIVE_SERVICE_ACCOUNT_CREDENTIALS_BASE64"
+    ).strip()
     service_account_json = json.loads(base64.b64decode(drive_creds_encoded).decode())
 
-    logger.info("Initializing ChatLogger instance...")
-    chat_logger = ChatLogger(service_account_json)
+    logger.info("Initializing ChatUploader instance...")
+    chat_uploader = ChatUploader(service_account_json)
 
     logger.info("Ready for user query...")
     return (
@@ -63,7 +65,7 @@ def initialize():
         knowledge_base,
         openAI_client,
         logger,
-        chat_logger,
+        chat_uploader,
     )
 
 
@@ -73,13 +75,14 @@ def initialize():
     knowledge_base,
     openAI_client,
     logger,
-    chat_logger,
+    chat_uploader,
 ) = initialize()
 
 
 def rag_chatbot(
     user_message: str,
-    chat_history: list = None,
+    chat_history: list,
+    browser_id: str,
 ) -> Generator[list, None, None]:
     """
     Retrieves relevant documents to user query and streams LLM response catching errors along the way.
@@ -163,33 +166,44 @@ def rag_chatbot(
 
     try:
         logger.info("Trying to upload chat history to Drive...")
-        chat_logger.upload_chat_history(chat_history)
+        chat_uploader.upload_chat_history(chat_history, browser_id)
     except Exception as e:
-        logger.warning(f"Error during Google Drive upload: {e}")
+        logger.warning(f"Warning: error during Google Drive upload: {e}")
 
     logger.info("Returning chat history...")
     return chat_history
 
 
 # Gradio app code
-demo = gr.ChatInterface(
-    fn=rag_chatbot,
-    title="Matthew Schulz's RAG Chatbot 💬🤖",
-    type="messages",
-    examples=[
-        ["What is Matthew's educational background?"],
-        ["What machine learning projects has Matthew worked on?"],
-        ["What experience does Matthew have in software engineering?"],
-        ["Why did Matthew choose to pursue a degree in computer science?"],
-        ["Does Matthew have any leadership experience?"],
-        ["Has Matthew completed any Summer internships?"],
-        ["Tell me about some real-world projects Matthew has worked on."],
-        ["What is Matthew's greatest strength and weakness?"],
-    ],
-    # save_history=True,
-    run_examples_on_click=False,
-    cache_examples=False,
-)
+with gr.Blocks() as demo:
+    browser_id_state = gr.BrowserState(default_value=None)
+
+    @demo.load(inputs=browser_id_state, outputs=browser_id_state)
+    def load_browser_id(current_id):
+        if current_id is None or current_id == "":
+            new_id = str(uuid.uuid4())
+            return new_id
+        return current_id
+
+    gr.ChatInterface(
+        fn=rag_chatbot,
+        title="Matthew Schulz's RAG Chatbot 💬🤖",
+        additional_inputs=browser_id_state,
+        type="messages",
+        examples=[
+            ["What is Matthew's educational background?"],
+            ["What machine learning projects has Matthew worked on?"],
+            ["What experience does Matthew have in software engineering?"],
+            ["Why did Matthew choose to pursue a degree in computer science?"],
+            ["Does Matthew have any leadership experience?"],
+            ["Has Matthew completed any Summer internships?"],
+            ["Tell me about some real-world projects Matthew has worked on."],
+            ["What is Matthew's greatest strength and weakness?"],
+        ],
+        save_history=True,
+        run_examples_on_click=False,
+        cache_examples=False,
+    )
 
 
 if __name__ == "__main__":

utils/chatLogger.py

@@ -1,64 +1,94 @@
 import io
 import json
-import uuid
 from googleapiclient.discovery import build
-from googleapiclient.http import MediaIoBaseUpload
+from googleapiclient.http import MediaIoBaseUpload, MediaIoBaseDownload
 from google.oauth2 import service_account
 
 
-class ChatLogger:
+class ChatUploader:
     def __init__(
         self,
         service_account_json: dict,
         root_folder_id: str = "1KtfVgL1Rg1iX-ZMHH4Im__ss-pgbaDM9",
     ):
         """
-        Initializes a new chat logger instance using a service account JSON dict.
+        Initializes a new chat uploader instance using a service account JSON dict.
         """
         credentials = service_account.Credentials.from_service_account_info(
             service_account_json, scopes=["https://www.googleapis.com/auth/drive"]
         )
         self.drive_service = build("drive", "v3", credentials=credentials)
         self.root_folder_id = root_folder_id
-        self.session_id = str(uuid.uuid4())
-        self.folder_id = self._create_session_folder()
 
-    def _create_session_folder(self) -> str:
+    def _get_or_create_browser_folder(self, browser_id: str) -> str:
         """
-        Creates a folder in Drive for this session using the UUID.
+        Searches for an existing folder for the given browser_id. If not found, creates a folder named 'browser_{browser_id}' and returns its ID.
         """
-        metadata = {
-            "name": f"session_{self.session_id}",
-            "mimeType": "application/vnd.google-apps.folder",
-            "parents": [self.root_folder_id],
-        }
-        folder = self.drive_service.files().create(body=metadata, fields="id").execute()
-        return folder["id"]
+        folder_name = f"browser_{browser_id}"
+        query = (
+            f"name = '{folder_name}' and '{self.root_folder_id}' in parents and "
+            "mimeType = 'application/vnd.google-apps.folder' and trashed = false"
+        )
+        results = self.drive_service.files().list(q=query, fields="files(id)").execute()
+        folders = results.get("files", [])
+
+        if folders:
+            return folders[0]["id"]
+        else:
+            metadata = {
+                "name": folder_name,
+                "mimeType": "application/vnd.google-apps.folder",
+                "parents": [self.root_folder_id],
+            }
+            folder = (
+                self.drive_service.files().create(body=metadata, fields="id").execute()
+            )
+            return folder["id"]
 
     def upload_chat_history(
-        self, chat_history: list, filename: str = "chat_log.json"
+        self, chat_history: list, browser_id: str, filename: str = "chat_log.json"
     ) -> None:
         """
-        Uploads or updates a single chat log file inside the session's folder.
+        Uploads the chat log file inside the browser's folder. If the folder and/or file exists, it appends the new chat entries to the current log. Otherwise, it creates them.
         """
+        folder_id = self._get_or_create_browser_folder(browser_id)
+
         query = (
-            f"name = '{filename}' and '{self.folder_id}' in parents and "
-            f"mimeType = 'application/json' and trashed = false"
+            f"name = '{filename}' and '{folder_id}' in parents and "
+            "mimeType = 'application/json' and trashed = false"
         )
         results = self.drive_service.files().list(q=query, fields="files(id)").execute()
         files = results.get("files", [])
 
-        content = json.dumps(chat_history, indent=2)
-
-        media = MediaIoBaseUpload(
-            io.BytesIO(content.encode()), mimetype="application/json"
-        )
-
         if files:
             file_id = files[0]["id"]
+
+            request = self.drive_service.files().get_media(fileId=file_id)
+            existing_stream = io.BytesIO()
+            downloader = MediaIoBaseDownload(existing_stream, request)
+            done = False
+            while not done:
+                _, done = downloader.next_chunk()
+
+            existing_stream.seek(0)
+            try:
+                existing_chat_history = json.loads(existing_stream.read())
+            except json.JSONDecodeError:
+                existing_chat_history = []
+
+            updated_chat_history = existing_chat_history + chat_history
+
+            content = json.dumps(updated_chat_history, indent=2)
+            media = MediaIoBaseUpload(
+                io.BytesIO(content.encode()), mimetype="application/json"
+            )
             self.drive_service.files().update(
                 fileId=file_id, media_body=media
             ).execute()
         else:
-            metadata = {"name": filename, "parents": [self.folder_id]}
+            content = json.dumps(chat_history, indent=2)
+            media = MediaIoBaseUpload(
+                io.BytesIO(content.encode()), mimetype="application/json"
+            )
+            metadata = {"name": filename, "parents": [folder_id]}
             self.drive_service.files().create(body=metadata, media_body=media).execute()