Update app.py

app.py

@@ -17,9 +17,11 @@ load_dotenv()
 # Set page configuration
 st.set_page_config(page_title="Legal Assistant", layout="wide")
 
-# Initialize session state for chat history if not exists
+# Initialize session state for chat history and response
 if 'chat_history' not in st.session_state:
     st.session_state.chat_history = []
+if 'last_response' not in st.session_state:
+    st.session_state.last_response = None
 
 # Title
 st.title("Legal Assistant")
@@ -57,9 +59,6 @@ def vector_embedding():
 # Perform vector embedding
 vector_embedding()
 
-# Main input area
-prompt1 = st.text_input("Enter Your Question From Documents")
-
 # Function to add to chat history
 def add_to_chat_history(question, answer):
     st.session_state.chat_history.append({
@@ -67,6 +66,29 @@ def add_to_chat_history(question, answer):
         'answer': answer
     })
 
+# Clear chat function
+def clear_chat():
+    # Clear the last response
+    st.session_state.last_response = None
+    # Optional: You can also clear the specific key used for the text input
+    st.session_state.prompt1 = ""
+
+# Main input area with clear button
+col1, col2 = st.columns([3, 1])
+with col1:
+    # Use a key to manage the text input state
+    prompt1 = st.text_input("Enter Your Question From Documents", key="prompt1")
+
+with col2:
+    # Add clear chat button next to the input
+    st.text("") # Add some vertical space to align with input
+    clear_button = st.button("Clear Chat")
+
+# Clear chat functionality
+if clear_button:
+    clear_chat()
+    st.experimental_rerun()
+
 # Process question and generate response
 if prompt1:
     try:
@@ -80,8 +102,8 @@ if prompt1:
         response = retrieval_chain.invoke({'input': prompt1})
         response_time = time.process_time() - start
         
-        # Display response
-        st.write(response['answer'])
+        # Store and display response
+        st.session_state.last_response = response['answer']
         
         # Add to chat history
         add_to_chat_history(prompt1, response['answer'])
@@ -89,6 +111,10 @@ if prompt1:
     except Exception as e:
         st.error(f"An error occurred: {e}")
 
+# Display the last response if exists
+if st.session_state.last_response:
+    st.write(st.session_state.last_response)
+
 # Sidebar content
 # Clear chat history button
 if st.sidebar.button("Clear Chat History"):