From 2c5115b70d8e6b2618d56331ee7fd1ec0638ac65 Mon Sep 17 00:00:00 2001
From: Pranshav Lakhia <pranshav.lakhia@gmail.com>
Date: Sun, 13 Oct 2024 11:26:38 -0400
Subject: [PATCH] hopefully stops input while generation

---
 main.py | 68 ++++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 43 insertions(+), 25 deletions(-)

diff --git a/main.py b/main.py
index ebf39d2..fd40867 100644
--- a/main.py
+++ b/main.py
@@ -4,11 +4,13 @@ import datetime
 from pymongo import MongoClient
 from streamlit_cookies_manager import CookieManager
 
+
 # Inject the scroll control HTML/JS
 def load_html(file_name):
     with open(file_name, 'r', encoding='utf-8') as file:
         return file.read()
 
+
 # Initialize the Cookie Manager (no encryption)
 cookies = CookieManager(prefix="ktosiek/streamlit-cookies-manager/")
 if not cookies.ready():
@@ -25,7 +27,6 @@ if 'user_id' not in cookies:
     cookies.save()
 
 user_id = cookies.get('user_id')
-# st.write(f"Current User ID: {user_id}")
 
 # Check if the user exists in the database
 user_data = users_collection.find_one({"user_id": user_id})
@@ -42,7 +43,11 @@ if not user_data:
 else:
     # If user exists, load the chat history from the database
     st.session_state.messages = user_data["messages"]
-    # st.write("Returning user detected! Loaded chat history from the database.")
+
+# Initialize a processing flag in session state
+if 'processing' not in st.session_state:
+    st.session_state.processing = False
+
 
 # Function to send POST request to a local server
 def send_post_request(prompt):
@@ -55,6 +60,7 @@ def send_post_request(prompt):
     else:
         return json
 
+
 # Custom CSS to fix the input box at the bottom of the page
 st.markdown(
     """
@@ -106,31 +112,43 @@ if 'messages' not in st.session_state:
 for message in st.session_state.messages:
     st.chat_message(message['role']).markdown(message['content'])
 
-# Prompt input
-prompt = st.chat_input("Ask your question here")
+# Disable the input box if processing a prompt
+if st.session_state.processing:
+    st.chat_input("Processing...", disabled=True)
+else:
+    # Prompt input
+    prompt = st.chat_input("Ask your question here")
 
-if prompt:
-    # Display the user prompt in the chat UI
-    st.chat_message('user').markdown(prompt)
-    st.session_state.messages.append({'role': 'user', 'content': prompt})
+    if prompt:
+        # Set the processing flag to True to disable input
+        st.session_state.processing = True
 
-    # Get the bot's response via a POST request
-    response = send_post_request(prompt)
+        # Display the user prompt in the chat UI
+        st.chat_message('user').markdown(prompt)
+        st.session_state.messages.append({'role': 'user', 'content': prompt})
 
-    # Display the bot's response in the chat UI
-    st.chat_message('bot').markdown(response)
-    st.session_state.messages.append({'role': 'bot', 'content': response})
+        # Show spinner while generating the response
+        with st.spinner("Generating response..."):
+            # Get the bot's response via a POST request
+            response = send_post_request(prompt)
 
-    # Update the user's chat history in the database
-    users_collection.update_one(
-        {"user_id": user_id},
-        {"$push": {"messages": {'role': 'user', 'content': prompt}}},
-    )
-    users_collection.update_one(
-        {"user_id": user_id},
-        {"$push": {"messages": {'role': 'bot', 'content': response}}},
-    )
+        # Display the bot's response in the chat UI
+        st.chat_message('bot').markdown(response)
+        st.session_state.messages.append({'role': 'bot', 'content': response})
 
-    # Save updated history to cookies (optional, MongoDB handles persistence)
-    cookies['history'] = st.session_state.messages
-    cookies.save()
+        # Update the user's chat history in the database
+        users_collection.update_one(
+            {"user_id": user_id},
+            {"$push": {"messages": {'role': 'user', 'content': prompt}}},
+        )
+        users_collection.update_one(
+            {"user_id": user_id},
+            {"$push": {"messages": {'role': 'bot', 'content': response}}},
+        )
+
+        # Save updated history to cookies (optional, MongoDB handles persistence)
+        cookies['history'] = st.session_state.messages
+        cookies.save()
+
+        # Reset the processing flag to enable input again
+        st.session_state.processing = False