Update app.py

app.py

@@ -15,13 +15,13 @@ from PatientInfoExtractionEngine import PatientInfoExtractionEngine
 load_dotenv()
 APP_TITLE = "Risk Adjustment (HCC Chart Validation)"
 CSV_PATH = "hcc_mapping.csv"
-SAMPLE_PDF = "sample_patient_chart.pdf" # Place a sample PDF in the same folder
+SAMPLE_PDF = "sample_patient_chart.pdf"  # Place a sample PDF in the same folder
 
 
 # ---------- JSON to Markdown ----------
 def json_to_markdown(data) -> str:
     try:
-        if isinstance(data, dict) and "final_analysis" in data:
+        if isinstance(data, dict):
             file_name = data.get("file_name", "Unknown Patient")
             hcc_code = data.get("hcc_code", "N/A")
             model_version = data.get("model_version", "N/A")
@@ -35,20 +35,13 @@ def json_to_markdown(data) -> str:
             address = data.get("address", "")
             phone = data.get("phone", "")
             patient_identifier = data.get("patient_identifier", "")
-
-        elif isinstance(data, list):
-            file_name = "N/A"
-            hcc_code = "N/A"
-            model_version = "N/A"
-            analyses = data
-            patient_name = dob = age = gender = address = phone = patient_identifier = ""
         else:
             return "<div style='color:red; font-weight:bold;'>⚠️ Invalid data format for report.</div>"
 
         md = f"""
 <div style="border:2px solid #4CAF50; padding:15px; border-radius:10px; background:#f9fdf9;">
   <h2 style="color:#2e7d32;">📋 HCC Chart Validation Report </h2>
-  <p><b>🧾 Patient ID:</b> {file_name}</p>
+  <p><b>🧾 File Name:</b> {file_name}</p>
   <p><b>🏷️ HCC Code:</b> {hcc_code}</p>
   <p><b>⚙️ Model Version:</b> {model_version}</p>
 """
@@ -67,215 +60,184 @@ def json_to_markdown(data) -> str:
 
         md += "</div><br/>"
 
-        # Render analyses
-        for idx, diag in enumerate(analyses, 1):
-            md += f"""
+        # Render analyses if they exist
+        if analyses:
+            for idx, diag in enumerate(analyses, 1):
+                md += f"""
 <div style="border:1px solid #ccc; padding:12px; border-radius:8px; margin-bottom:12px;">
   <h3 style="color:#1565c0;">{idx}. {diag.get("diagnosis", "Unknown Diagnosis")}</h3>
   <p><b>ICD-10:</b> {diag.get("icd10", "N/A")}</p>
   <p><b>Reference:</b> <a href="{diag.get("reference","")}" target="_blank">{diag.get("reference","")}</a></p>
 """
-            explicit_ans = diag.get("answer_explicit", "N/A")
-            explicit_rat = diag.get("rationale_explicit", "")
-            implicit_ans = diag.get("answer_implicit", "N/A")
-            implicit_rat = diag.get("rationale_implicit", "")
+                explicit_ans = diag.get("answer_explicit", "N/A")
+                explicit_rat = diag.get("rationale_explicit", "")
+                implicit_ans = diag.get("answer_implicit", "N/A")
+                implicit_rat = diag.get("rationale_implicit", "")
 
-            if explicit_ans.lower() == "yes":
-                md += f"<p><b>Explicit:</b> {explicit_ans} — {explicit_rat}</p>"
-            else:
-                md += f"<p><b>Implicit:</b> {implicit_ans} — {implicit_rat}</p>"
+                if explicit_ans.lower() == "yes":
+                    md += f"<p><b>Explicit:</b> {explicit_ans} — {explicit_rat}</p>"
+                else:
+                    md += f"<p><b>Implicit:</b> {implicit_ans} — {implicit_rat}</p>"
 
-            md += f"""
+                md += f"""
   <p><b>Clinical Status:</b> {diag.get("clinical_status","N/A")}</p>
   <p><b>Status Rationale:</b> {diag.get("status_rationale","")}</p>
 """
-            if "tests" in diag:
-                md += "<details><summary><b>🧪 Tests & Procedures</b></summary><ul>"
-                tests = diag["tests"]
-                if "vitals" in tests:
-                    md += "<li><b>Vitals:</b><ul>"
-                    for k, v in tests["vitals"].items():
-                        md += f"<li>{k}: {v}</li>"
-                    md += "</ul></li>"
-                if "procedures" in tests:
-                    md += "<li><b>Procedures:</b><ul>"
-                    for k, v in tests["procedures"].items():
-                        md += f"<li>{k}: {v}</li>"
-                    md += "</ul></li>"
-                if "lab_test" in tests:
-                    md += "<li><b>Lab Tests:</b><ul>"
-                    for k, v in tests["lab_test"].items():
-                        md += f"<li>{k}: {v}</li>"
-                    md += "</ul></li>"
-                md += "</ul></details>"
-
-            if "meat" in diag:
-                md += "<details><summary><b>🍖 MEAT Validation</b></summary><ul>"
-                for k, v in diag["meat"].items():
-                    emoji = "✅" if v else "❌"
-                    md += f"<li>{k.capitalize()}: {emoji}</li>"
-                md += "</ul>"
-                md += f"<p><b>MEAT Rationale:</b> {diag.get('meat_rationale','')}</p>"
-                md += "</details>"
-
-            if "comorbidities" in diag and diag["comorbidities"]:
-                md += "<details><summary><b>🩺 Comorbidities</b></summary><ul>"
-                for c in diag["comorbidities"]:
-                    emoji = "✅" if c.get("is_present") else "❌"
-                    md += f"<li>{emoji} <b>{c.get('condition')}</b><br/><i>{c.get('rationale')}</i></li>"
-                md += "</ul></details>"
-
-            md += "</div>"
-
+                if "tests" in diag:
+                    md += "<details><summary><b>🧪 Tests & Procedures</b></summary><ul>"
+                    tests = diag["tests"]
+                    if "vitals" in tests:
+                        md += "<li><b>Vitals:</b><ul>"
+                        for k, v in tests["vitals"].items(): md += f"<li>{k}: {v}</li>"
+                        md += "</ul></li>"
+                    if "procedures" in tests:
+                        md += "<li><b>Procedures:</b><ul>"
+                        for k, v in tests["procedures"].items(): md += f"<li>{k}: {v}</li>"
+                        md += "</ul></li>"
+                    if "lab_test" in tests:
+                        md += "<li><b>Lab Tests:</b><ul>"
+                        for k, v in tests["lab_test"].items(): md += f"<li>{k}: {v}</li>"
+                        md += "</ul></li>"
+                    md += "</ul></details>"
+
+                if "meat" in diag:
+                    md += "<details><summary><b>🍖 MEAT Validation</b></summary><ul>"
+                    for k, v in diag["meat"].items():
+                        emoji = "✅" if v else "❌"
+                        md += f"<li>{k.capitalize()}: {emoji}</li>"
+                    md += "</ul>"
+                    md += f"<p><b>MEAT Rationale:</b> {diag.get('meat_rationale','')}</p>"
+                    md += "</details>"
+
+                if "comorbidities" in diag and diag["comorbidities"]:
+                    md += "<details><summary><b>🩺 Comorbidities</b></summary><ul>"
+                    for c in diag["comorbidities"]:
+                        emoji = "✅" if c.get("is_present") else "❌"
+                        md += f"<li>{emoji} <b>{c.get('condition')}</b><br/><i>{c.get('rationale')}</i></li>"
+                    md += "</ul></details>"
+
+                md += "</div>"
         return md
-
     except Exception as e:
         return f"<div style='color:red; font-weight:bold;'>⚠️ Error rendering report: {e}</div>"
 
 
-# ---------- Processing Pipeline Functions ----------
-
-def extract_demographics_only(pdf_file):
-    """Step 1: Extracts and displays only the demographic information."""
-    if pdf_file is None:
-        return {
-            gr.Markdown(value="<div style='color:orange;'>⚠️ Please upload a patient chart PDF to begin.</div>"): gr.update(),
-            gr.State(): gr.update(),
-            gr.Button(): gr.update(visible=False)
-        }
-    
-    pdf_path = pdf_file.name
-    file_name = os.path.splitext(os.path.basename(pdf_path))[0]
-    
-    demographics_engine = PatientInfoExtractionEngine(pdf_path)
-    demographics_info = demographics_engine.run()
-    
-    # Store info for the next step
-    stored_data = {
-        "file_name": file_name,
-        "pdf_path": pdf_path,
-        "demographics_info": demographics_info
-    }
-    
-    # Create a partial report with just demographics
-    report_data = {
-        "file_name": file_name,
-        "hcc_code": "N/A",
-        "model_version": "N/A",
-        "final_analysis": [],
-        **demographics_info
-    }
-    
-    md_output = json_to_markdown(report_data)
-    
-    return {
-        gr.Markdown(value=md_output): gr.update(),
-        gr.State(value=stored_data): gr.update(),
-        gr.Button(value="✅ Demographics Confirmed. Run Full Validation.", visible=True): gr.update()
-    }
-
-def run_full_validation(stored_data, hcc_code, model_version, progress=gr.Progress()):
-    """Step 2: Runs the rest of the validation pipeline."""
+# ---------- Processing Pipeline with Gradio Progress ----------
+def process_pipeline(pdf_file, hcc_code, model_version, csv_path=CSV_PATH, output_folder="outputs", progress=gr.Progress()):
     try:
         start = time.time()
-        # Start from step 2, since step 1 (demographics) is done
-        step = 1 
-        total_steps = 7 # 8 total steps minus the first one
-        
-        pdf_path = stored_data["pdf_path"]
-        file_name = stored_data["file_name"]
-        demographics_info = stored_data["demographics_info"]
+        step = 0
+        total_steps = 8  # Total number of steps in the pipeline
 
         def log(msg, current_step=0):
-            # The progress bar will show steps 2 through 8
-            # We add 1 to current_step for display purposes
-            display_step = current_step + 1
             elapsed = time.time() - start
             bar_html = '<div style="display: flex; width: 100%; gap: 2px; height: 12px; margin: 8px 0 5px 0;">'
-            for i in range(2, total_steps + 2):
-                if i < display_step: color = "#1e40af"
-                elif i == display_step: color = "#3b82f6"
+            for i in range(1, total_steps + 1):
+                if i < current_step: color = "#1e40af"
+                elif i == current_step: color = "#3b82f6"
                 else: color = "#e5e7eb"
                 bar_html += f'<div style="flex-grow: 1; background-color: {color}; border-radius: 3px;"></div>'
             bar_html += '</div>'
             return f"{msg}{bar_html}<small>⏳ Elapsed: {elapsed:.1f} sec</small>"
-        
+
+        if pdf_file is None:
+            yield log("⚠️ Please upload a patient chart PDF.", 0)
+            return
         hcc_code_str = str(hcc_code or "").strip()
         if not hcc_code_str:
-            yield "⚠️ Please enter a valid HCC Code before running validation."
+            yield log("⚠️ Please enter a valid HCC Code before running validation.", 0)
             return
 
-        # Initial Markdown with demographics
-        initial_md = json_to_markdown({
+        os.makedirs(output_folder, exist_ok=True)
+        pdf_path = pdf_file.name
+        file_name = os.path.splitext(os.path.basename(pdf_path))[0]
+        print(f"[PROCESSING] {file_name}")
+
+        # --- MODIFICATION: Extract and show demographics first ---
+        # Step 1: Extract Demographics
+        step += 1
+        progress((step, total_steps), desc="Extracting demographics")
+        yield log(f"🧍 Step {step}/{total_steps}: Extracting patient demographics...", step)
+        demographics_engine = PatientInfoExtractionEngine(pdf_path)
+        demographics_info = demographics_engine.run()
+        print(f"[DEMOGRAPHICS] Extracted: {demographics_info}")
+
+        # Create and yield the initial report with only demographics
+        initial_report_data = {
             "file_name": file_name,
             "hcc_code": hcc_code_str,
             "model_version": model_version,
             "final_analysis": [],
             **demographics_info
-        })
+        }
+        demographics_md = json_to_markdown(initial_report_data)
+        yield demographics_md
+        time.sleep(1) # Pause for a moment to show the demographics
+        # --- END OF MODIFICATION ---
+        
+        # Subsequent steps will append progress below the initial demographic display
         
         # Step 2: Diagnoses
         step += 1
-        progress((step-1, total_steps), desc="Extracting diagnoses")
-        yield initial_md + log(f"🔍 Step {step}/{total_steps+1}: Extracting possible HCC Diagnoses...", step)
-        diagnoses = HCCDiagnosisListEngine(hcc_code_str, model_version, CSV_PATH).run()
+        progress((step, total_steps), desc="Extracting diagnoses")
+        yield demographics_md + log(f"🔍 Step {step}/{total_steps}: Extracting possible HCC Diagnoses...", step)
+        diagnoses = HCCDiagnosisListEngine(hcc_code_str, model_version, csv_path).run()
         if not diagnoses:
-            yield initial_md + log(f"❌ No diagnoses found for HCC {hcc_code_str}.", step)
+            yield demographics_md + log(f"❌ No diagnoses found for HCC {hcc_code_str}.", step)
             return
 
         # Step 3: Chart checking
         step += 1
-        progress((step-1, total_steps), desc="Checking chart")
-        yield initial_md + log(f"📄 Step {step}/{total_steps+1}: Checking diagnoses in patient chart...", step)
+        progress((step, total_steps), desc="Checking chart")
+        yield demographics_md + log(f"📄 Step {step}/{total_steps}: Checking diagnoses in patient chart...", step)
         all_checked_results = ChartDiagnosisChecker(pdf_path).run(diagnoses)
         confirmed_diagnoses = [d for d in all_checked_results if d.get("answer_explicit", "").lower() == "yes" or d.get("answer_implicit", "").lower() == "yes"]
         if not confirmed_diagnoses:
-            yield initial_md + log(f"❌ No confirmed diagnoses for HCC {hcc_code_str} in {file_name}.", step)
+            yield demographics_md + log(f"❌ No confirmed diagnoses for HCC {hcc_code_str} in {file_name}.", step)
             return
 
         # Step 4: Tests
         step += 1
-        progress((step-1, total_steps), desc="Finding tests")
-        yield initial_md + log(f"🧪 Step {step}/{total_steps+1}: Finding relevant tests...", step)
+        progress((step, total_steps), desc="Finding tests")
+        yield demographics_md + log(f"🧪 Step {step}/{total_steps}: Finding relevant tests...", step)
         diagnoses_with_tests = TestFindingAgent(hcc_code=hcc_code_str, model_version=model_version).run(confirmed_diagnoses)
 
         # Step 5: Clinical Status
         step += 1
-        progress((step-1, total_steps), desc="Determining clinical status")
-        yield initial_md + log(f"⚕️ Step {step}/{total_steps+1}: Determining clinical status...", step)
+        progress((step, total_steps), desc="Determining clinical status")
+        yield demographics_md + log(f"⚕️ Step {step}/{total_steps}: Determining clinical status...", step)
         diagnoses_with_status = ClinicalStatusAgent().run(diagnoses_with_tests)
         active_diagnoses = [d for d in diagnoses_with_status if d.get("clinical_status") == "ACTIVE"]
 
         # Step 6: MEAT
         step += 1
-        progress((step-1, total_steps), desc="Validating MEAT")
+        progress((step, total_steps), desc="Validating MEAT")
         if active_diagnoses:
-            yield initial_md + log(f"🍖 Step {step}/{total_steps+1}: Validating MEAT...", step)
+            yield demographics_md + log(f"🍖 Step {step}/{total_steps}: Validating MEAT...", step)
             validated_meat_diagnoses = MEATValidatorAgent().run(active_diagnoses)
         else:
             validated_meat_diagnoses = []
-            yield initial_md + log("ℹ️ No ACTIVE diagnoses found. Skipping MEAT/Comorbidity.", step)
+            yield demographics_md + log("ℹ️ No ACTIVE diagnoses found. Skipping MEAT/Comorbidity.", step)
 
         # Step 7: Comorbidities
         step += 1
-        progress((step-1, total_steps), desc="Checking comorbidities")
+        progress((step, total_steps), desc="Checking comorbidities")
         diagnoses_passed_meat = [d for d in validated_meat_diagnoses if any(d.get("meat", {}).values())]
         if diagnoses_passed_meat:
-            yield initial_md + log(f"🤝 Step {step}/{total_steps+1}: Checking comorbidities...", step)
+            yield demographics_md + log(f"🤝 Step {step}/{total_steps}: Checking comorbidities...", step)
             comorbidity_results = ComorbidityCheckerAgent(pdf_path, hcc_code_str, model_version).run(diagnoses_passed_meat)
         else:
             comorbidity_results = []
 
         # Step 8: Final Report
         step += 1
-        progress((step-1, total_steps), desc="Generating report")
-        yield initial_md + log(f"✅ Step {step}/{total_steps+1}: Generating final report...", step)
+        progress((step, total_steps), desc="Generating report")
+        yield demographics_md + log(f"✅ Step {step}/{total_steps}: Generating final report...", step)
 
-        # Merge results
+        # Merge results for final output
         status_map = {d["diagnosis"]: d for d in diagnoses_with_status}
         meat_map = {d["diagnosis"]: d for d in validated_meat_diagnoses}
         comorbidity_map = {d["diagnosis"]: d for d in comorbidity_results}
-
         final_analysis = []
         for entry in all_checked_results:
             diag_name = entry["diagnosis"]
@@ -302,6 +264,7 @@ def run_full_validation(stored_data, hcc_code, model_version, progress=gr.Progre
         print(f"[ERROR] {e}")
         yield f"<div style='color:red; font-weight:bold;'>⚠️ Error: {e}</div>"
 
+
 # ---------- Gradio Theme and Helpers ----------
 simple_theme = gr.themes.Soft(
     primary_hue=gr.themes.colors.blue,
@@ -332,6 +295,7 @@ def pdf_to_iframe(file):
     except Exception as e:
         return f"<p style='color:red;'>Failed to display PDF: {e}</p>"
 
+
 # ---------- Gradio UI ----------
 with gr.Blocks(theme=simple_theme, title=APP_TITLE) as interface:
     gr.HTML(f"""
@@ -340,18 +304,12 @@ with gr.Blocks(theme=simple_theme, title=APP_TITLE) as interface:
             Upload a chart, set HCC + model version, and validate MEAT criteria.
         </p>
     """)
-    
-    # Store data between steps
-    stored_data = gr.State()
 
     with gr.Row():
         pdf_upload = gr.File(label="Upload Patient Chart (PDF)", file_types=[".pdf"], scale=1)
         hcc_code = gr.Textbox(label="HCC Code (e.g., 12)", placeholder="Enter HCC code", scale=1)
         model_version = gr.Dropdown(choices=["V24", "V28"], label="Model Version", value="V24", scale=1)
-    
-    with gr.Row():
-        run_demographics_btn = gr.Button("🔍 Extract Demographics", variant="primary", scale=1)
-        run_full_validation_btn = gr.Button("🚀 Run Full Validation", variant="primary", scale=1, visible=False)
+        run_btn = gr.Button("🚀 Run Validation", variant="primary", scale=1)
 
     with gr.Row():
         with gr.Column(scale=2):
@@ -359,43 +317,26 @@ with gr.Blocks(theme=simple_theme, title=APP_TITLE) as interface:
         with gr.Column(scale=2):
             output_md = gr.Markdown(
                 label="Validation Report",
-                value="<div style='border:2px solid #1e40af; border-radius:12px; padding:15px; background-color:#f0f9ff;'>📄 Upload a PDF and click <b>Extract Demographics</b> to start.</div>",
+                value="<div style='border:2px solid #1e40af; border-radius:12px; padding:15px; background-color:#f0f9ff;'>📄 Upload a PDF and click <b>Run Validation</b> to start.</div>",
             )
 
     pdf_upload.change(fn=pdf_to_iframe, inputs=pdf_upload, outputs=pdf_preview)
 
-    run_demographics_btn.click(
-        fn=extract_demographics_only,
-        inputs=[pdf_upload],
-        outputs=[output_md, stored_data, run_full_validation_btn],
-    )
-    
-    run_full_validation_btn.click(
-        fn=run_full_validation,
-        inputs=[stored_data, hcc_code, model_version],
+    run_btn.click(
+        fn=process_pipeline,
+        inputs=[pdf_upload, hcc_code, model_version],
         outputs=[output_md],
     )
-    
-    # Example loader needs to be adapted for the two-step process
-    def load_and_run_example(pdf_path):
-        sample_pdf = load_sample_pdf()
-        # Step 1
-        demographics_result = extract_demographics_only(sample_pdf)
-        # We need to yield updates to the UI, which is complex in gr.Examples
-        # For simplicity, we can just show the final result
-        # A better approach would be a custom example function that triggers both buttons.
-        # This is a simplified version for demonstration.
-        return "Examples are loaded. Please click the buttons to process."
 
     gr.Examples(
         examples=[[SAMPLE_PDF]],
         inputs=[pdf_upload],
-        # The output of an example click is now just a message.
         outputs=[output_md],
-        fn=load_and_run_example,
+        fn=lambda x: process_pipeline(load_sample_pdf(), hcc_code="12", model_version="V24"),
         cache_examples=False
     )
 
+
 if __name__ == "__main__":
     interface.queue().launch(
         server_name="0.0.0.0",