Update app.py

app.py

@@ -17,13 +17,27 @@ load_dotenv()
 # Set page configuration
 st.set_page_config(page_title="Legal Assistant", layout="wide")
 
-# 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
-if 'last_question' not in st.session_state:
-    st.session_state.last_question = ""
+# Initialize session state
+def init_session_state():
+    # Initialize session state variables if they don't exist
+    session_vars = [
+        'chat_history',
+        'last_response',
+        'current_question',
+        'should_clear'
+    ]
+    
+    for var in session_vars:
+        if var not in st.session_state:
+            if var == 'chat_history':
+                st.session_state[var] = []
+            elif var == 'should_clear':
+                st.session_state[var] = False
+            else:
+                st.session_state[var] = None
+
+# Call initialization
+init_session_state()
 
 # Title
 st.title("Legal Assistant")
@@ -68,63 +82,71 @@ def add_to_chat_history(question, answer):
         'answer': answer
     })
 
-# Main input area with clear button
-col1, col2 = st.columns([3, 1])
-with col1:
-    # Text input with dynamic value management
-    prompt1 = st.text_input("Enter Your Question From Documents", 
-                             value=st.session_state.last_question, 
-                             key="user_question")
-
-with col2:
-    # Add clear chat button next to the input
-    st.text("") # Add some vertical space to align with input
+# Main content
+def main():
+    # Clear chat button
     clear_button = st.button("Clear Chat")
+    
+    # Handle clear chat
+    if clear_button:
+        # Set clear flag
+        st.session_state.should_clear = True
+        # Reset current question and last response
+        st.session_state.current_question = ""
+        st.session_state.last_response = None
+
+    # Reset clear flag after processing
+    if st.session_state.should_clear:
+        st.session_state.should_clear = False
 
-# Clear chat functionality
-if clear_button:
-    # Reset last question and response
-    st.session_state.last_question = ""
-    st.session_state.last_response = None
-
-# Process question and generate response
-if prompt1:
-    try:
-        # Store the current question
-        st.session_state.last_question = prompt1
-
-        # Create document and retrieval chains
-        document_chain = create_stuff_documents_chain(llm, prompt)
-        retriever = st.session_state.vectors.as_retriever()
-        retrieval_chain = create_retrieval_chain(retriever, document_chain)
-        
-        # Generate response
-        start = time.process_time()
-        response = retrieval_chain.invoke({'input': prompt1})
-        response_time = time.process_time() - start
-        
-        # Store and display response
-        st.session_state.last_response = response['answer']
+    # Text input with dynamic value management
+    current_question = st.text_input(
+        "Enter Your Question From Documents", 
+        value="" if st.session_state.should_clear else st.session_state.current_question,
+        key="question_input"
+    )
+
+    # Process question if exists
+    if current_question:
+        try:
+            # Store current question
+            st.session_state.current_question = current_question
+
+            # Create document and retrieval chains
+            document_chain = create_stuff_documents_chain(llm, prompt)
+            retriever = st.session_state.vectors.as_retriever()
+            retrieval_chain = create_retrieval_chain(retriever, document_chain)
+            
+            # Generate response
+            start = time.process_time()
+            response = retrieval_chain.invoke({'input': current_question})
+            response_time = time.process_time() - start
+            
+            # Store and display response
+            st.session_state.last_response = response['answer']
+            
+            # Add to chat history
+            add_to_chat_history(current_question, response['answer'])
         
-        # Add to chat history
-        add_to_chat_history(prompt1, response['answer'])
-    
-    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"):
-    st.session_state.chat_history = []
-
-# Display chat history
-st.sidebar.write("### Previous Questions")
-for idx, chat in enumerate(reversed(st.session_state.chat_history), 1):
-    # Expander for each chat history item
-    with st.sidebar.expander(f"Question {len(st.session_state.chat_history) - idx + 1}"):
-        st.write(f"**Question:** {chat['question']}")
-        st.write(f"**Answer:** {chat['answer']}")
+        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"):
+        st.session_state.chat_history = []
+
+    # Display chat history
+    st.sidebar.write("### Previous Questions")
+    for idx, chat in enumerate(reversed(st.session_state.chat_history), 1):
+        # Expander for each chat history item
+        with st.sidebar.expander(f"Question {len(st.session_state.chat_history) - idx + 1}"):
+            st.write(f"**Question:** {chat['question']}")
+            st.write(f"**Answer:** {chat['answer']}")
+
+# Run the main function
+main()