Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,13 +10,13 @@ from model import (
     extract_and_link_sources,
     parse_doctor_response
 )
-from report_agent import (
-    orchestrate_report_generation,
-    detect_report_request,
-    add_report_download_buttons,
-    MedicalReport
-)
 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
 
 
 # Set page config with dark theme
@@ -369,12 +369,6 @@ if 'use_rag' not in st.session_state:
     st.session_state.use_rag = True
 if 'processing' not in st.session_state:
     st.session_state.processing = False
-if 'report_generation' not in st.session_state:
-    st.session_state.report_generation = False
-if 'report_data' not in st.session_state:
-    st.session_state.report_data = MedicalReport().model_dump()
-if 'report_files' not in st.session_state:
-    st.session_state.report_files = None
 
 # Helper function to check if explanation has meaningful content
 def has_meaningful_content(text):
@@ -523,16 +517,277 @@ 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.
+"""
+    # Use our existing orchestrator_chat function
+    result, _, _ = orchestrator_chat(
+        [], # No chat history needed
+        conversation_text, 
+        use_rag=True,
+        is_system_prompt=True,
+        system_message=system_prompt
+    )
+    
+    # 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):
+        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": ""}
+    
+    # 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)
+            
+            # Offer download
+            st.download_button(
+                label="📥 Download Medical Report",
+                data=pdf_bytes,
+                file_name=f"medical_report_{datetime.now().strftime('%Y%m%d')}.pdf",
+                mime="application/pdf",
+                key="report_download"
+            )
+            
+            # Reset the report process for next time
+            st.session_state.report_step = 0
+            
+            # Display completion message
+            st.success("Your medical report is ready for download!")
+
 # Put the toggle in the sidebar - this is the most reliable approach
 with st.sidebar:
     st.header("Features")
     st.session_state.use_rag = st.toggle("Database Search", value=False, 
                                        help="Toggle to enable or disable medical database search")
-    st.session_state.report_generation = st.toggle("Report Generation", value=False,
-                                                help="Toggle to enable or disable medical report generation")
+    
+    # Add Generate Report button
+    if st.button("Generate Report", use_container_width=True):
+        st.session_state.show_report_form = True
+        # Reset report step for a new report
+        st.session_state.report_step = 1
+        st.rerun()
+    
+    # Show report form when button is clicked
+    if st.session_state.get('show_report_form', False):
+        st.subheader("Report Information")
+        generate_and_download_report()
+
+# Initialize new session state variables
+if 'show_report_form' not in st.session_state:
+    st.session_state.show_report_form = False
+if 'report_step' not in st.session_state:
+    st.session_state.report_step = 0
+if 'patient_info' not in st.session_state:
+    st.session_state.patient_info = {"name": "", "age": "", "gender": ""}
 
 # Main app area
-st.title("What are you Diagnosing Today?")
+st.title("How are you Feeling Today?")
 
 # Display chat history
 for message in st.session_state.history:
@@ -564,11 +819,6 @@ for message in st.session_state.history:
                                 st.markdown("---")
                                 st.markdown("**Legend:** 🔓 = Open Access (full text available)")
 
-# Display report download buttons if available
-if st.session_state.report_files and st.session_state.report_generation:
-    html_path, pdf_path, docx_path = st.session_state.report_files
-    add_report_download_buttons(html_path, pdf_path, docx_path)
-
 # Display typing indicator if processing
 if st.session_state.processing:
     with st.container():
@@ -603,77 +853,26 @@ if st.session_state.processing:
                                 if msg["role"] == "user"), None)
         
         if last_user_message:
-            # Check if this is a report generation request when toggle is on
-            if st.session_state.report_generation:
-                # Call the report generation orchestrator
-                response, report_files = orchestrate_report_generation(
-                    st.session_state.history,
-                    last_user_message,
-                    st.session_state.report_data
-                )
-                
-                # If we got a response (either a prompt for missing data or a confirmation)
-                if response:
-                    # Update session state with the report data
-                    st.session_state.report_data = response
-                    
-                    # If we got report files, store them in session state
-                    if report_files:
-                        st.session_state.report_files = report_files
-                        
-                        # Add assistant response to history with report confirmation
-                        st.session_state.history.append({
-                            "role": "assistant",
-                            "content": response
-                        })
-                    else:
-                        # This is a prompt for missing data
-                        st.session_state.history.append({
-                            "role": "assistant",
-                            "content": response
-                        })
-                else:
-                    # Process as normal conversation if it's not a report request
-                    reply, explanation, evidence = orchestrator_chat(
-                        [msg for msg in st.session_state.history if msg["role"] == "assistant" or 
-                        (msg["role"] == "user" and msg["content"] != last_user_message)],
-                        last_user_message,
-                        use_rag=st.session_state.use_rag,
-                        is_follow_up=len(st.session_state.history) > 1
-                    )
-                    
-                    # Clean up the response and explanation
-                    cleaned_reply = remove_reasoning_and_sources(reply)
-                    cleaned_explanation = clean_explanation(explanation) if explanation else ""
-                    
-                    # Add assistant response to history
-                    st.session_state.history.append({
-                        "role": "assistant", 
-                        "content": cleaned_reply,
-                        "explanation": cleaned_explanation,
-                        "evidence": evidence if evidence else []
-                    })
-            else:
-                # Normal conversation mode (report generation disabled)
-                reply, explanation, evidence = orchestrator_chat(
-                    [msg for msg in st.session_state.history if msg["role"] == "assistant" or 
-                     (msg["role"] == "user" and msg["content"] != last_user_message)],
-                    last_user_message,
-                    use_rag=st.session_state.use_rag,
-                    is_follow_up=len(st.session_state.history) > 1
-                )
-                
-                # Clean up the response and explanation
-                cleaned_reply = remove_reasoning_and_sources(reply)
-                cleaned_explanation = clean_explanation(explanation) if explanation else ""
-                
-                # Add assistant response to history
-                st.session_state.history.append({
-                    "role": "assistant", 
-                    "content": cleaned_reply,
-                    "explanation": cleaned_explanation,
-                    "evidence": evidence if evidence else []
-                })
+            # Process as normal conversation
+            reply, explanation, evidence = orchestrator_chat(
+                [msg for msg in st.session_state.history if msg["role"] == "assistant" or 
+                 (msg["role"] == "user" and msg["content"] != last_user_message)],
+                last_user_message,
+                use_rag=st.session_state.use_rag,
+                is_follow_up=len(st.session_state.history) > 1
+            )
+            
+            # Clean up the response and explanation
+            cleaned_reply = remove_reasoning_and_sources(reply)
+            cleaned_explanation = clean_explanation(explanation) if explanation else ""
+            
+            # Add assistant response to history
+            st.session_state.history.append({
+                "role": "assistant", 
+                "content": cleaned_reply,
+                "explanation": cleaned_explanation,
+                "evidence": evidence if evidence else []
+            })
     finally:
         # Set processing back to false regardless of success/failure
         st.session_state.processing = False