Update app2.py

app2.py

@@ -1953,6 +1953,9 @@ def legal_consultant_ui():
     if 'uploaded_document' not in st.session_state:
         st.session_state.uploaded_document = None
     
+    if 'consultant_chat_mode' not in st.session_state:  # New variable for consultant chat mode
+        st.session_state.consultant_chat_mode = "normal"
+    
     st.write('''
     Describe your legal situation or ask your legal question related to US law. 
     LexAI will provide information and guidance based on its understanding of the US legal system. 
@@ -1967,6 +1970,14 @@ def legal_consultant_ui():
         st.session_state.uploaded_document = extract_text_from_document(uploaded_file)
         st.success("Document uploaded successfully!")
     
+    # Chat mode toggle for Legal Consultant
+    if st.session_state.uploaded_document:
+        if st.button("Switch Chat Mode"):
+            st.session_state.consultant_chat_mode = "document" if st.session_state.consultant_chat_mode == "normal" else "normal"
+        
+        st.write(f"Current mode: {'Document-based' if st.session_state.consultant_chat_mode == 'document' else 'Normal'} chat")
+
+    # Display chat history (similar to Legal Chatbot)
     for message in st.session_state.consultant_chat_history:
         if isinstance(message, tuple):
             user_msg, bot_msg = message
@@ -1977,32 +1988,25 @@ def legal_consultant_ui():
                 web_results_msg = "Web Search Results:\n"
                 for result in message.get('results', []):
                     web_results_msg += f"[{result.get('title', 'No title')}]({result.get('link', '#')})\n{result.get('snippet', 'No snippet available.')}\n\n"
-                st.success(f"**Lex AI:** {web_results_msg}") 
+                st.success(f"**Lex AI:** {web_results_msg}")
 
     user_input = st.text_input("Your legal question:")
     
     if user_input and st.button("Send"):
         with st.spinner("Processing your question..."):
-            if st.session_state.uploaded_document:
-                full_input = f"Document context: {st.session_state.uploaded_document}\n\nUser question: {user_input}"
+            if st.session_state.consultant_chat_mode == "document" and st.session_state.uploaded_document:
+                ai_response = get_document_based_response(user_input, st.session_state.uploaded_document)  # Assuming you have this function
+                st.session_state.consultant_chat_history.append((user_input, ai_response))
             else:
-                full_input = user_input
-            
-            result = provide_legal_advice(full_input)
-
-            if "error" in result:
-                st.error(result["error"])
-            else:
-                st.session_state.consultant_chat_history.append((user_input, result["advice"]))
-                
-                # Generate web search query and perform search
-                search_query = generate_web_search_query(user_input)
-                web_results = search_web_duckduckgo(search_query)
+                ai_response = provide_legal_advice(user_input)
+                st.session_state.consultant_chat_history.append((user_input, ai_response))
                 
+                # Perform web search (same as Legal Chatbot)
+                web_results = search_web_duckduckgo(user_input)  # Using user_input directly here
                 st.session_state.consultant_chat_history.append({
                     'type': 'web_search',
                     'results': web_results
-                }) 
+                })
         
         st.rerun()