Update app.py

app.py

@@ -1,5 +1,3 @@
-# Full updated app.py with TOOL_CALLS rendered + copy.deepcopy fix for Gradio UI updates
-
 import sys
 import os
 import pandas as pd
@@ -16,7 +14,9 @@ import subprocess
 import traceback
 import torch
 import copy
+import time
 
+# Configure environment variables
 os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"
 if not torch.cuda.is_available():
     print("No GPU detected. Forcing CPU mode by setting CUDA_VISIBLE_DEVICES to an empty string.")
@@ -65,7 +65,7 @@ def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
                 text_chunks.append(f"=== Page {i+1} ===\n{text.strip()}")
             for i, page in enumerate(pdf.pages[3:max_pages], start=4):
                 page_text = page.extract_text() or ""
-                if any(re.search(rf'\\b{kw}\\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
+                if any(re.search(rf'\b{kw}\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
                     text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
         return "\n\n".join(text_chunks)
     except Exception as e:
@@ -153,35 +153,39 @@ def init_agent():
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
+        # Persistent conversation state
+        conversation_state = gr.State([])
         chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
         file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
         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")
 
-        def analyze(message: str, history: list, files: list):
-            try:
-                history.append({"role": "user", "content": message})
-                history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
-                yield copy.deepcopy(history), None
-
-                extracted = ""
-                file_hash_value = ""
-                if files:
-                    with ThreadPoolExecutor(max_workers=4) as executor:
-                        futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files]
-                        results = []
-                        for future in as_completed(futures):
-                            try:
-                                res = future.result()
-                                results.append(sanitize_utf8(res))
-                            except Exception as e:
-                                print("❌ Error in file processing:", str(e))
-                                traceback.print_exc()
-                    extracted = "\n".join(results)
-                    file_hash_value = file_hash(files[0].name)
-
-                prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
+        def analyze(message: str, state: list, files: list):
+            if state is None:
+                state = []
+            history = state
+            history.append({"role": "user", "content": message})
+            history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
+            yield copy.deepcopy(history), None, copy.deepcopy(history)
+
+            extracted = ""
+            file_hash_value = ""
+            if files:
+                with ThreadPoolExecutor(max_workers=4) as executor:
+                    futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files]
+                    results = []
+                    for future in as_completed(futures):
+                        try:
+                            res = future.result()
+                            results.append(sanitize_utf8(res))
+                        except Exception as e:
+                            print("❌ Error in file processing:", str(e))
+                            traceback.print_exc()
+                extracted = "\n".join(results)
+                file_hash_value = file_hash(files[0].name)
+
+            prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
 3. Note incomplete assessments
@@ -192,13 +196,14 @@ Medical Records:
 
 ### Potential Oversights:
 """
-                print("🔎 Generated prompt:")
-                print(prompt)
+            print("🔎 Generated prompt:")
+            print(prompt)
 
-                full_response = ""
-                response_chunks = []
-                tool_calls_rendered = []
+            full_response = ""
+            response_chunks = []
+            tool_calls_rendered = []
 
+            try:
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -229,7 +234,7 @@ Medical Records:
                     else:
                         history.append({"role": "assistant", "content": display_response})
 
-                    yield copy.deepcopy(history), None
+                    yield copy.deepcopy(history), None, copy.deepcopy(history)
 
                 full_response = re.sub(r"\[TOOL_CALLS\].*?\n*", "", full_response, flags=re.DOTALL).strip()
                 full_response = full_response.replace('[TxAgent]', '').strip()
@@ -245,15 +250,17 @@ Medical Records:
                 else:
                     history.append({"role": "assistant", "content": full_response})
 
-                yield copy.deepcopy(history), report_path if report_path and os.path.exists(report_path) else None
+                yield copy.deepcopy(history), report_path if report_path and os.path.exists(report_path) else None, copy.deepcopy(history)
 
             except Exception as e:
                 history.append({"role": "assistant", "content": f"❌ An error occurred in analyze: {str(e)}"})
                 traceback.print_exc()
-                yield copy.deepcopy(history), None
+                yield copy.deepcopy(history), None, copy.deepcopy(history)
 
-        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])
+        send_btn.click(analyze, inputs=[msg_input, conversation_state, file_upload],
+                       outputs=[chatbot, download_output, conversation_state])
+        msg_input.submit(analyze, inputs=[msg_input, conversation_state, file_upload],
+                         outputs=[chatbot, download_output, conversation_state])
     return demo
 
 if __name__ == "__main__":
@@ -266,8 +273,9 @@ if __name__ == "__main__":
             server_port=7860,
             show_error=True,
             allowed_paths=[report_dir],
-            share=False
+            share=False,
+            ssr=False  # Disable SSR to improve UI updates
         )
     except Exception as e:
         print("❌ Fatal error during launch:", str(e))
-        traceback.print_exc()
+        traceback.print_exc()