Update app.py

app.py

@@ -214,7 +214,7 @@ def create_ui(agent):
         msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
         send_btn = gr.Button("Analyze", variant="primary")
         download_output = gr.File(label="Download Full Report")
-        progress_bar = gr.Progress()
+        progress_bar = gr.Progress()  # Keep for visual feedback
 
         prompt_template = """
 Analyze the patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown with findings grouped under tool-derived headings (e.g., 'Drugs'). For each finding, include clinical context, risks, and recommendations. Precede findings with a tool tag (e.g., [TOOL: get_abuse_info_by_drug_name]). Output only markdown bullet points under headings. If no issues, state "No issues identified".
@@ -225,14 +225,14 @@ Patient Record Excerpt (Chunk {0} of {1}):
 
         def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
             history.append({"role": "user", "content": message})
-            yield history, None, None
+            yield history, None
 
             extracted = ""
             file_hash_value = ""
             if files:
                 def update_extraction_progress(current, total):
                     progress(current / total, desc=f"Extracting text... Page {current}/{total}")
-                    return history, None, None
+                    return history, None
 
                 with ThreadPoolExecutor(max_workers=6) as executor:
                     futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files]
@@ -241,7 +241,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     file_hash_value = file_hash(files[0].name) if files else ""
 
             history.append({"role": "assistant", "content": "✅ Text extraction complete."})
-            yield history, None, None
+            yield history, None
 
             chunk_size = 6000
             chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
@@ -281,7 +281,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                         else:
                             combined_response += f"--- Analysis for Chunk {chunk_idx} ---\nNo oversights identified for this chunk.\n\n"
                         history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                        yield history, None, None
+                        yield history, None
 
                 if combined_response.strip() and not all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
                     history[-1]["content"] = combined_response.strip()
@@ -292,25 +292,29 @@ Patient Record Excerpt (Chunk {0} of {1}):
                 if report_path:
                     with open(report_path, "w", encoding="utf-8") as f:
                         f.write(combined_response)
-                yield history, report_path if report_path and os.path.exists(report_path) else None, None
+                yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e:
                 print("🚨 ERROR:", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield history, None, None
+                yield history, None
 
-        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, progress_bar])
-        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, progress_bar])
+        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
+        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
     return demo
 
 if __name__ == "__main__":
-    print("🚀 Launching app...")
-    agent = init_agent()
-    demo = create_ui(agent)
-    demo.queue(api_open=False).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        allowed_paths=[report_dir],
-        share=False
-    )
+    try:
+        print("🚀 Launching app...")
+        agent = init_agent()
+        demo = create_ui(agent)
+        demo.queue(api_open=False).launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True,
+            allowed_paths=[report_dir],
+            share=False
+        )
+    finally:
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()  # Clean up distributed resources