Update src/streamlit_app.py

src/streamlit_app.py

@@ -11,13 +11,14 @@ from model import (
     parse_doctor_response
 )
 import base64
-from reportlab.lib.pagesizes import A4
-from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
-from reportlab.lib.enums import TA_LEFT
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
-from reportlab.lib import colors
-import tempfile
 from sendgrid_service import SendGridService
+from report_generator import (
+    extract_medical_json,
+    build_medical_report,
+    format_conversation_history,
+    generate_and_download_report,
+    show_email_form
+)
 
 
 # Set page config with dark theme
@@ -528,315 +529,6 @@ assistant_logo_base64 = get_image_base64(assistant_logo_path)
 user_avatar = f"data:image/png;base64,{user_logo_base64}" if user_logo_base64 else "👤"
 assistant_avatar = f"data:image/png;base64,{assistant_logo_base64}" if assistant_logo_base64 else "🤖"
 
-# Add report generation functions
-def extract_medical_json(conversation_text: str) -> dict:
-    """
-    Extract medical report data from conversation text into structured JSON.
-    Uses the existing orchestrator_chat model.
-    """
-    system_prompt = """
-You are an expert medical report generator. Extract ALL info from the patient-assistant conversation into this JSON schema:
-{
-  "patient": {"name":"","age":"","gender":""},
-  "visit_date":"YYYY-MM-DD",
-  "chief_complaint":"",
-  "history_of_present_illness":"",
-  "past_medical_history":"",
-  "medications":"",
-  "allergies":"",
-  "examination":"",
-  "diagnosis":"",
-  "recommendations":"",
-  "reasoning":"",
-  "sources":""
-}
-Do NOT invent any data. Return ONLY the JSON object.
-"""
-    # Create a fake history with the system prompt as a system message
-    history = [
-        {"role": "system", "content": system_prompt}
-    ]
-    
-    # Use our existing orchestrator_chat function with valid parameters
-    result, _, _ = orchestrator_chat(
-        history=history,
-        query=conversation_text, 
-        use_rag=True,
-        is_follow_up=False
-    )
-    
-    # Extract JSON from response
-    try:
-        # Find JSON object in the response
-        json_match = re.search(r'({[\s\S]*})', result)
-        if json_match:
-            json_str = json_match.group(1)
-            return json.loads(json_str)
-        else:
-            # If no clear JSON format found, try to parse the whole response
-            return json.loads(result)
-    except json.JSONDecodeError:
-        # Fallback with basic structure
-        st.error("Failed to parse report data")
-        return {
-            "patient": {"name":"", "age":"", "gender":""},
-            "visit_date": datetime.now().strftime("%Y-%m-%d"),
-            "chief_complaint": "Error generating report"
-        }
-
-def build_medical_report(data: dict) -> bytes:
-    """
-    Generates a PDF from the extracted JSON and returns PDF bytes.
-    """
-    # Create a temporary file for the PDF
-    fd, temp_path = tempfile.mkstemp(suffix='.pdf')
-    os.close(fd)
-    
-    doc = SimpleDocTemplate(
-        temp_path, 
-        pagesize=A4,
-        rightMargin=40, 
-        leftMargin=40,
-        topMargin=60, 
-        bottomMargin=60
-    )
-    
-    styles = getSampleStyleSheet()
-    styles.add(ParagraphStyle(
-        name='SectionHeading',
-        fontSize=14, 
-        leading=16,
-        spaceAfter=8,
-        alignment=TA_LEFT,
-        textColor=colors.darkblue
-    ))
-    
-    styles.add(ParagraphStyle(
-        name='NormalLeft',
-        fontSize=11, 
-        leading=14,
-        spaceAfter=6,
-        alignment=TA_LEFT
-    ))
-    
-    elems = []
-
-    # Title
-    elems.append(Paragraph("Medical Consultation Report", styles['Title']))
-    elems.append(Spacer(1, 12))
-
-    # Patient info
-    patient = data.get('patient', {})
-    info = [
-        ['Name:', patient.get('name', '–')],
-        ['Age:', patient.get('age', '–')],
-        ['Gender:', patient.get('gender', '–')],
-        ['Date:', data.get('visit_date', datetime.now().strftime("%Y-%m-%d"))],
-    ]
-    
-    tbl = Table(info, colWidths=[80, 250])
-    tbl.setStyle(TableStyle([
-        ('BACKGROUND', (0, 0), (1, 0), colors.lightgrey),
-        ('BOX', (0, 0), (-1, -1), 0.5, colors.grey),
-        ('INNERGRID', (0, 0), (-1, -1), 0.5, colors.grey),
-        ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
-    ]))
-    
-    elems += [tbl, Spacer(1, 16)]
-
-    # Helper to add sections
-    def add_section(title, text):
-        # Convert text to string if it's a list or another non-string type
-        if isinstance(text, list):
-            text = '\n'.join(map(str, text))
-        elif text is not None and not isinstance(text, str):
-            text = str(text)
-            
-        elems.append(Paragraph(title, styles['SectionHeading']))
-        elems.append(Paragraph(text or '–', styles['NormalLeft']))
-        elems.append(Spacer(1, 12))
-
-    # Populate all sections
-    add_section("Chief Complaint", data.get('chief_complaint'))
-    add_section("History of Present Illness", data.get('history_of_present_illness'))
-    add_section("Past Medical History", data.get('past_medical_history'))
-    add_section("Medications", data.get('medications'))
-    add_section("Allergies", data.get('allergies'))
-    add_section("Examination Findings", data.get('examination'))
-    add_section("Diagnosis", data.get('diagnosis'))
-    add_section("Recommendations", data.get('recommendations'))
-    add_section("Reasoning", data.get('reasoning'))
-    add_section("Sources", data.get('sources'))
-
-    doc.build(elems)
-    
-    # Read PDF bytes
-    with open(temp_path, 'rb') as file:
-        pdf_bytes = file.read()
-    
-    # Clean up the temporary file
-    os.unlink(temp_path)
-    
-    return pdf_bytes
-
-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
-
-# Function to handle the report generation process
-def generate_and_download_report():
-    # 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
-    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"])
-        if st.button("Generate Report"):
-            st.session_state.patient_info["gender"] = gender
-            st.session_state.report_step = 4
-            st.rerun()
-
-    # Step 4: Generate report
-    elif st.session_state.report_step == 4:
-        with st.spinner("Generating medical report..."):
-            # Format conversation with patient info
-            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:
-                # Offer download
-                st.download_button(
-                    label="📥 Download Report",
-                    data=pdf_bytes,
-                    file_name=f"medical_report_{datetime.now().strftime('%Y%m%d')}.pdf",
-                    mime="application/pdf",
-                    key="report_download"
-                )
-            
-            with col2:
-                # Email button that shows email form when clicked
-                if st.button("📧 Email Report"):
-                    st.session_state.show_email_form = True
-                    st.rerun()
-            
-            # Show email form if the button was clicked
-            if st.session_state.get('show_email_form', False):
-                st.subheader("Send Report via Email")
-                
-                # Check SendGrid configuration
-                sendgrid_api_key = os.getenv('SENDGRID_API_KEY')
-                if not sendgrid_api_key:
-                    st.warning("Email service not configured. Please contact administrator.")
-                
-                # Email form
-                email = st.text_input("Recipient Email Address:")
-                
-                # Email send button
-                if st.button("Send"):
-                    if not email:
-                        st.error("Please enter an email address.")
-                    else:
-                        # Validate email
-                        sendgrid = SendGridService()
-                        if not sendgrid.validate_email(email):
-                            st.error("Please enter a valid email address.")
-                        else:
-                            # Send email
-                            with st.spinner("Sending email... This may take a few seconds."):
-                                success, message = sendgrid.send_report(
-                                    email,
-                                    st.session_state.pdf_data,
-                                    st.session_state.patient_info["name"]
-                                )
-                                
-                                if success:
-                                    st.success(message)
-                                    # Hide the form after successful send
-                                    st.session_state.show_email_form = False
-                                    st.rerun()
-                                else:
-                                    st.error(message)
-                
-                # Cancel button to hide email form
-                if st.button("Cancel"):
-                    st.session_state.show_email_form = False
-                    st.rerun()
-            
-            # 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()
-
 # Put the toggle in the sidebar - this is the most reliable approach
 with st.sidebar:
     st.header("Features")