Update src/report_generator.py

src/report_generator.py

@@ -14,43 +14,8 @@ from reportlab.lib.units import inch, cm
 from reportlab.platypus import PageBreak, KeepTogether
 
 from model import orchestrator_chat
-
-
-# Function to format conversation history for report generation
-def format_conversation_history(history, patient_info=None):
-    """
-    Format the conversation history into a string suitable for LLM processing.
-    Optionally adds patient info at the beginning.
-    """
-    formatted_text = "# Medical Consultation\n\n"
-    
-    # Add patient info if provided
-    if patient_info:
-        formatted_text += "## Patient Information\n"
-        formatted_text += f"* Name: {patient_info.get('name', '')}\n"
-        formatted_text += f"* Age: {patient_info.get('age', '')}\n"
-        formatted_text += f"* Gender: {patient_info.get('gender', '')}\n\n"
-    
-    formatted_text += "## Conversation Transcript\n\n"
-    
-    for message in history:
-        role = message.get("role", "").strip()
-        content = message.get("content", "").strip()
-        
-        if not content:
-            continue  # Skip empty messages
-        
-        if role.lower() == "user":
-            formatted_text += f"PATIENT: {content}\n\n"
-        elif role.lower() == "assistant":
-            formatted_text += f"ASSISTANT: {content}\n\n"
-            # Include explanations which often contain diagnostic reasoning
-            if "explanation" in message and message["explanation"]:
-                explanation = message.get("explanation", "").strip()
-                if explanation:
-                    formatted_text += f"REASONING: {explanation}\n\n"
-    
-    return formatted_text
+from utils import format_conversation_history
+from session_state import get_full_history
 
 
 # Function to encode images to base64 (similar to app.py)
@@ -461,91 +426,83 @@ def build_medical_report(data: dict) -> bytes:
 
 
 def generate_and_download_report():
-    """Handle the report generation process with multi-step form UI."""
-    # Store the current conversation step
-    if 'report_step' not in st.session_state:
-        st.session_state.report_step = 1
-        st.session_state.patient_info = {"name": "", "age": "", "gender": ""}
-        st.session_state.pdf_data = None  # Store PDF data for email
-
-    # Step 1: Collect patient name
+    """
+    Generate a medical report from the conversation history.
+    Shows a form for entering patient info and then allows downloading the report.
+    """
+    # Step 1: Collect patient info
     if st.session_state.report_step == 1:
-        name = st.text_input("Patient Name:")
-        if st.button("Next"):
-            st.session_state.patient_info["name"] = name
-            st.session_state.report_step = 2
-            st.rerun()
-
-    # Step 2: Collect age
-    elif st.session_state.report_step == 2:
-        age = st.text_input("Patient Age:")
-        if st.button("Next"):
-            st.session_state.patient_info["age"] = age
-            st.session_state.report_step = 3
-            st.rerun()
-
-    # Step 3: Collect gender
-    elif st.session_state.report_step == 3:
-        gender = st.selectbox("Patient Gender:", ["Male", "Female", "Other", "Prefer not to say"])
+        st.text_input("Patient Name", key="patient_name", 
+                     value=st.session_state.patient_info.get("name", ""))
+        st.text_input("Age", key="patient_age", 
+                     value=st.session_state.patient_info.get("age", ""))
+        st.selectbox("Gender", ["Male", "Female", "Other"], key="patient_gender",
+                    index=["Male", "Female", "Other"].index(st.session_state.patient_info.get("gender", "Male")))
+        
         if st.button("Generate Report"):
-            st.session_state.patient_info["gender"] = gender
-            st.session_state.report_step = 4
+            # Save patient info to session state
+            st.session_state.patient_info = {
+                "name": st.session_state.patient_name,
+                "age": st.session_state.patient_age,
+                "gender": st.session_state.patient_gender
+            }
+            st.session_state.report_step = 2
             st.rerun()
-
-    # Step 4: Generate report
-    elif st.session_state.report_step == 4:
-        with st.spinner("Generating medical report..."):
-            # Format conversation with patient info - using the local format_conversation_history
-            conversation_text = format_conversation_history(
-                st.session_state.history, 
-                st.session_state.patient_info
-            )
-            
-            # Extract structured data
-            report_json = extract_medical_json(conversation_text)
-            
-            # Override with collected patient info
-            report_json["patient"] = st.session_state.patient_info
-            report_json["visit_date"] = datetime.now().strftime("%Y-%m-%d")
-            
-            # Generate PDF
-            pdf_bytes = build_medical_report(report_json)
-            st.session_state.pdf_data = pdf_bytes  # Store PDF data for email
             
-            # Display success message
-            st.success("Your medical report is ready!")
-            
-            # Create two columns for the download and email buttons
-            col1, col2 = st.columns(2)
-            
-            with col1:
-                # Format the filename to include patient name like in email
-                patient_name = st.session_state.patient_info["name"]
-                safe_name = patient_name.replace(' ', '_').replace(',', '').replace('.', '')
-                # Offer download with consistent filename format
+    # Step 2: Generate and display report
+    elif st.session_state.report_step == 2:
+        with st.spinner("Generating report..."):
+            try:
+                # Get full conversation history from database
+                full_history = get_full_history()
+                
+                # Format the conversation history for the report
+                formatted_conversation = format_conversation_history(
+                    full_history, 
+                    st.session_state.patient_info
+                )
+                
+                # Extract structured data from the conversation
+                medical_data = extract_medical_json(formatted_conversation)
+                
+                # Update with patient info from the form
+                medical_data["patient"].update(st.session_state.patient_info)
+                
+                # Set the visit date to today if not present
+                if not medical_data.get("visit_date"):
+                    medical_data["visit_date"] = datetime.now().strftime("%Y-%m-%d")
+                
+                # Generate PDF
+                pdf_data = build_medical_report(medical_data)
+                
+                # Store PDF data in session state
+                st.session_state.pdf_data = pdf_data
+                
+                # Create download button
                 st.download_button(
                     label="Download Report",
-                    data=pdf_bytes,
-                    file_name=f"medical_report_daease_{safe_name}.pdf",
+                    data=pdf_data,
+                    file_name=f"medical_report_{st.session_state.patient_info['name'].replace(' ', '_')}.pdf",
                     mime="application/pdf",
-                    key="report_download"
+                    use_container_width=True
                 )
-            
-            with col2:
-                # Email button that shows email form when clicked
-                if st.button("Email Report"):
+                
+                # Show email option
+                if st.button("Email Report", use_container_width=True):
                     st.session_state.show_email_form = True
+                
+                # Show email form if requested
+                if st.session_state.show_email_form:
+                    show_email_form()
+                
+                # Option to start over
+                if st.button("New Report", use_container_width=True):
+                    st.session_state.report_step = 1
                     st.rerun()
-            
-            # Show email form if the button was clicked
-            if st.session_state.get('show_email_form', False):
-                show_email_form()
-            
-            # Button to start over
-            if st.button("Generate Another Report"):
-                st.session_state.report_step = 1
-                st.session_state.show_email_form = False
-                st.rerun()
+                    
+            except Exception as e:
+                st.error(f"Error generating report: {str(e)}")
+                st.button("Try Again", on_click=lambda: setattr(st.session_state, "report_step", 1))
 
 
 def show_email_form():