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