Update app.py

app.py

@@ -17,27 +17,18 @@ load_dotenv()
 # Set page configuration
 st.set_page_config(page_title="Legal Assistant", layout="wide")
 
-# 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'
-    ]
+# Create a unique key for the session state to help with resetting
+if 'reset_key' not in st.session_state:
+    st.session_state['reset_key'] = 0
+
+# Function to reset the entire session state
+def reset_session_state():
+    # Increment reset key to force a complete reset
+    st.session_state['reset_key'] += 1
     
-    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()
+    # Reset specific session state variables
+    st.session_state['last_response'] = None
+    st.session_state['current_question'] = ''
 
 # Title
 st.title("Legal Assistant")
@@ -49,6 +40,10 @@ st.sidebar.title("Chat History")
 groq_api_key = os.getenv('groqapi')
 os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
 
+# Initialize chat history if not exists
+if 'chat_history' not in st.session_state:
+    st.session_state['chat_history'] = []
+
 # LLM and Prompt Setup
 llm = ChatGroq(groq_api_key=groq_api_key, model_name="Llama3-8b-8192")
 prompt = ChatPromptTemplate.from_template(
@@ -82,35 +77,30 @@ def add_to_chat_history(question, answer):
         'answer': answer
     })
 
-# Main content
+# Main content area
 def main():
     # Clear chat button
     clear_button = st.button("Clear Chat")
     
-    # Handle clear chat
+    # Handle clear chat functionality
     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
-
-    # Text input with dynamic value management
-    current_question = st.text_input(
+        reset_session_state()
+
+    # Create a unique key for the text input to force reset
+    text_input_key = f'question_input_{st.session_state["reset_key"]}'
+    
+    # Text input with reset mechanism
+    prompt1 = st.text_input(
         "Enter Your Question From Documents", 
-        value="" if st.session_state.should_clear else st.session_state.current_question,
-        key="question_input"
+        key=text_input_key,
+        value=st.session_state.get('current_question', '')
     )
 
     # Process question if exists
-    if current_question:
+    if prompt1:
         try:
             # Store current question
-            st.session_state.current_question = current_question
+            st.session_state['current_question'] = prompt1
 
             # Create document and retrieval chains
             document_chain = create_stuff_documents_chain(llm, prompt)
@@ -119,21 +109,21 @@ def main():
             
             # Generate response
             start = time.process_time()
-            response = retrieval_chain.invoke({'input': current_question})
+            response = retrieval_chain.invoke({'input': prompt1})
             response_time = time.process_time() - start
             
             # Store and display response
-            st.session_state.last_response = response['answer']
+            st.session_state['last_response'] = response['answer']
             
             # Add to chat history
-            add_to_chat_history(current_question, response['answer'])
+            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)
+    if st.session_state.get('last_response'):
+        st.write(st.session_state['last_response'])
 
     # Sidebar content
     # Clear chat history button