Update src/report_generator.py

src/report_generator.py

@@ -16,6 +16,7 @@ import requests
 
 from model import orchestrator_chat
 from utils import format_conversation_history
+from session_state import get_full_history
 
 
 # Function to fetch current date and time from an API
@@ -514,8 +515,8 @@ def generate_and_download_report():
     elif st.session_state.report_step == 2:
         with st.spinner("Generating report..."):
             try:
-                # Get full conversation history from local session state
-                full_history = st.session_state.get('history', [])
+                # Get full conversation history from database
+                full_history = get_full_history()
                 
                 # Format the conversation history for the report
                 formatted_conversation = format_conversation_history(
@@ -561,15 +562,22 @@ def generate_and_download_report():
                 if st.session_state.show_email_form:
                     show_email_form()
                 
-                # After displaying download/email options, reset report_step to prevent re-generation
-                # UNLESS an email form is active, in which case, it will be reset by show_email_form or cancel.
-                if not st.session_state.get('show_email_form', False):
-                    st.session_state.report_step = 0 # Reset to neutral state
+                # Option to start over
+                if st.button("New Report", use_container_width=True):
+                    # Reset report step
+                    st.session_state.report_step = 1
+                    # Clear PDF data
+                    if 'pdf_data' in st.session_state:
+                        del st.session_state.pdf_data
+                    # Reset patient info
+                    st.session_state.patient_info = {"name": "", "age": "", "gender": ""}
+                    # Hide email form if shown
+                    st.session_state.show_email_form = False
+                    # Force refresh
+                    st.rerun()
                     
             except Exception as e:
                 st.error(f"Error generating report: {str(e)}")
-                # Ensure report_step is reset even on error to allow trying again from step 1
-                st.session_state.report_step = 1 
                 st.button("Try Again", on_click=lambda: setattr(st.session_state, "report_step", 1))
 
 
@@ -609,7 +617,6 @@ def show_email_form():
                         st.success(message)
                         # Hide the form after successful send
                         st.session_state.show_email_form = False
-                        st.session_state.report_step = 0 # Reset report step
                         st.rerun()
                     else:
                         st.error(message)
@@ -617,5 +624,4 @@ def show_email_form():
     # Cancel button to hide email form
     if st.button("Cancel"):
         st.session_state.show_email_form = False
-        st.session_state.report_step = 0 # Reset report step
         st.rerun()