Update app.py

app.py

@@ -131,12 +131,16 @@ def json_to_markdown(data) -> str:
 
 def process_pipeline(pdf_file, hcc_code, model_version, csv_path=CSV_PATH, output_folder="outputs"):
     try:
+        start = time.time()
+
         if pdf_file is None:
-            return "<div style='color:orange; font-weight:bold;'>⚠️ Please upload a patient chart PDF.</div>"
+            yield "<div style='color:orange; font-weight:bold;'>⚠️ Please upload a patient chart PDF.</div>"
+            return
 
         hcc_code_str = str(hcc_code or "").strip()
         if not hcc_code_str:
-            return "<div style='color:orange; font-weight:bold;'>⚠️ Please enter a valid HCC Code before running validation.</div>"
+            yield "<div style='color:orange; font-weight:bold;'>⚠️ Please enter a valid HCC Code before running validation.</div>"
+            return
 
         os.makedirs(output_folder, exist_ok=True)
         pdf_path = pdf_file.name
@@ -144,15 +148,20 @@ def process_pipeline(pdf_file, hcc_code, model_version, csv_path=CSV_PATH, outpu
         print(f"\n[PROCESSING] {patient_name}")
 
         # Step 1: Get all possible HCC Diagnoses
+        yield "🔍 Step 1: Extracting possible HCC Diagnoses..."
         diagnoses = HCCDiagnosisListEngine(hcc_code_str, model_version, csv_path).run()
+        yield f"⏳ Elapsed: {time.time() - start:.1f} sec"
+
         if not diagnoses:
-            print(f"[STOP] No diagnoses found for HCC {hcc_code_str}.")
-            return f"<b>❌ No diagnoses found for HCC {hcc_code_str}.</b>"
+            msg = f"<b>❌ No diagnoses found for HCC {hcc_code_str}.</b>"
+            yield msg
+            return
 
-        # Step 2: Check which diagnoses are present in the patient chart
+        # Step 2: Check patient chart
+        yield "📄 Step 2: Checking diagnoses in patient chart..."
         all_checked_results = ChartDiagnosisChecker(pdf_path).run(diagnoses)
+        yield f"⏳ Elapsed: {time.time() - start:.1f} sec"
 
-        # --- First Checkpoint: Confirmed Diagnoses (explicit OR implicit = yes) ---
         confirmed_diagnoses = [
             d for d in all_checked_results
             if d.get("answer_explicit", "").lower() == "yes"
@@ -160,7 +169,7 @@ def process_pipeline(pdf_file, hcc_code, model_version, csv_path=CSV_PATH, outpu
         ]
 
         if not confirmed_diagnoses:
-            print(f"[STOP] No confirmed diagnoses found in chart for HCC {hcc_code_str}.")
+            msg = f"<b>❌ No confirmed diagnoses for HCC {hcc_code_str} in {patient_name}. Report saved.</b>"
             output_data = {
                 "patient_id": patient_name,
                 "hcc_code": hcc_code_str,
@@ -171,35 +180,38 @@ def process_pipeline(pdf_file, hcc_code, model_version, csv_path=CSV_PATH, outpu
             with open(outpath, "w", encoding="utf-8") as f:
                 json.dump(output_data, f, indent=2, ensure_ascii=False)
             print(f"[DONE] Saved basic chart check report to {outpath}")
-            return f"<b>❌ No confirmed diagnoses for HCC {hcc_code_str} in {patient_name}. Report saved.</b>"
+            yield msg
+            return
 
-        print(f"\n[INFO] Found {len(confirmed_diagnoses)} confirmed diagnoses. Proceeding with detailed analysis...")
-
-        # Step 3: Run TestFindingAgent
+        # Step 3: Tests
+        yield "🧪 Step 3: Finding relevant tests..."
         diagnoses_with_tests = TestFindingAgent(hcc_code=hcc_code_str, model_version=model_version).run(confirmed_diagnoses)
+        yield f"⏳ Elapsed: {time.time() - start:.1f} sec"
 
-        # Step 4: Determine Clinical Status
+        # Step 4: Clinical Status
+        yield "⚕️ Step 4: Determining clinical status..."
         diagnoses_with_status = ClinicalStatusAgent().run(diagnoses_with_tests)
+        yield f"⏳ Elapsed: {time.time() - start:.1f} sec"
 
-        # --- Second Checkpoint: Filter ACTIVE diagnoses ---
         active_diagnoses = [d for d in diagnoses_with_status if d.get("clinical_status") == "ACTIVE"]
 
         validated_meat_diagnoses = []
         if active_diagnoses:
-            print(f"[INFO] Found {len(active_diagnoses)} ACTIVE diagnoses. Validating MEAT...")
+            yield "🍖 Step 5: Validating MEAT..."
             validated_meat_diagnoses = MEATValidatorAgent().run(active_diagnoses)
+            yield f"⏳ Elapsed: {time.time() - start:.1f} sec"
         else:
-            print("[INFO] No ACTIVE diagnoses. Skipping MEAT/Comorbidity.")
+            yield "ℹ️ No ACTIVE diagnoses found. Skipping MEAT/Comorbidity."
 
-        # --- Third Checkpoint: Diagnoses passing MEAT ---
         diagnoses_passed_meat = [d for d in validated_meat_diagnoses if any(d.get("meat", {}).values())]
 
         comorbidity_results = []
         if diagnoses_passed_meat:
-            print(f"[INFO] {len(diagnoses_passed_meat)} diagnoses passed MEAT. Checking comorbidities...")
+            yield "🤝 Step 6: Checking comorbidities..."
             comorbidity_results = ComorbidityCheckerAgent(pdf_path, hcc_code_str, model_version).run(diagnoses_passed_meat)
+            yield f"⏳ Elapsed: {time.time() - start:.1f} sec"
 
-        # --- Merge results ---
+        # Merge results
         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}
@@ -217,31 +229,29 @@ def process_pipeline(pdf_file, hcc_code, model_version, csv_path=CSV_PATH, outpu
             else:
                 final_analysis.append(entry)
 
-        # --- Final Filter: Keep only YES answers ---
         filtered_final_analysis = [
             e for e in final_analysis
             if e.get("answer_explicit", "").lower() == "yes"
             or e.get("answer_implicit", "").lower() == "yes"
         ]
 
-        # Save report
         output_data = {
             "patient_id": patient_name,
             "hcc_code": hcc_code_str,
             "model_version": model_version,
             "final_analysis": filtered_final_analysis
         }
-        # outpath = os.path.join(output_folder, f"{patient_name}.json")
-        # with open(outpath, "w", encoding="utf-8") as f:
-        #     json.dump(output_data, f, indent=2, ensure_ascii=False)
-        # print(f"[DONE] Saved final report to {outpath}")
 
-        # Return styled markdown
-        return json_to_markdown(output_data)
+        # Final Report
+        yield "✅ Step 7: Generating final report..."
+        elapsed = time.time() - start
+        final_report = json_to_markdown(output_data) + f"<br/><div style='color:green;'>✅ Completed in {elapsed:.1f} sec</div>"
+        yield final_report
 
     except Exception as e:
         print(f"[ERROR] {e}")
-        return f"<div style='color:red; font-weight:bold;'>⚠️ Error: {e}</div>"
+        yield f"<div style='color:red; font-weight:bold;'>⚠️ Error: {e}</div>"
+