Update app.py

app.py

@@ -11,6 +11,7 @@ import shutil
 import re
 import psutil
 import subprocess
+import traceback
 
 # Persistent directory setup
 persistent_dir = "/data/hf_cache"
@@ -37,6 +38,7 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
+# Medical Keywords used in the PDF processing
 MEDICAL_KEYWORDS = {'diagnosis', 'assessment', 'plan', 'results', 'medications',
                     'allergies', 'summary', 'impression', 'findings', 'recommendations'}
 
@@ -55,14 +57,17 @@ def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
             for i, page in enumerate(pdf.pages[:3]):
                 text = page.extract_text() or ""
                 text_chunks.append(f"=== Page {i+1} ===\n{text.strip()}")
-            # Then include pages that mention one or more medical keywords
+            # Include further pages if they contain any medical keywords
             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):
                     text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
         return "\n\n".join(text_chunks)
     except Exception as e:
-        return f"PDF processing error: {str(e)}"
+        debug_msg = f"PDF processing error: {str(e)}"
+        print(debug_msg)
+        traceback.print_exc()
+        return debug_msg
 
 def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
@@ -93,7 +98,10 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             f.write(result)
         return result
     except Exception as e:
-        return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"})
+        error_msg = f"Error processing {os.path.basename(file_path)}: {str(e)}"
+        print(error_msg)
+        traceback.print_exc()
+        return json.dumps({"error": error_msg})
 
 def log_system_usage(tag=""):
     try:
@@ -109,41 +117,35 @@ def log_system_usage(tag=""):
             print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
     except Exception as e:
         print(f"[{tag}] GPU/CPU monitor failed: {e}")
+        traceback.print_exc()
 
 def init_agent():
-    print("🔁 Initializing model...")
-    log_system_usage("Before Load")
-    default_tool_path = os.path.abspath("data/new_tool.json")
-    target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
-    if not os.path.exists(target_tool_path):
-        shutil.copy(default_tool_path, target_tool_path)
-
-    agent = TxAgent(
-        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
-        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-        tool_files_dict={"new_tool": target_tool_path},
-        force_finish=True,
-        enable_checker=True,
-        step_rag_num=8,
-        seed=100,
-        additional_default_tools=[],
-    )
-    agent.init_model()
-    log_system_usage("After Load")
-    print("✅ Agent Ready")
-    return agent
-
-def clean_response(response: str) -> str:
-    """
-    Updated cleaner that removes the [TOOL_CALLS] tag and any JSON artifacts
-    while preserving the full analysis so that all identified oversights are displayed.
-    """
-    # Remove everything starting from the first [TOOL_CALLS] occurrence
-    cleaned = response.split("[TOOL_CALLS]")[0].strip()
-    # Remove any remaining JSON artifacts in case they appear
-    cleaned = re.sub(r'\{.*?\}', '', cleaned)
-    cleaned = re.sub(r'\[.*?\]', '', cleaned)
-    return cleaned
+    try:
+        print("🔁 Initializing model...")
+        log_system_usage("Before Load")
+        default_tool_path = os.path.abspath("data/new_tool.json")
+        target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
+        if not os.path.exists(target_tool_path):
+            shutil.copy(default_tool_path, target_tool_path)
+
+        agent = TxAgent(
+            model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
+            rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
+            tool_files_dict={"new_tool": target_tool_path},
+            force_finish=True,
+            enable_checker=True,
+            step_rag_num=8,
+            seed=100,
+            additional_default_tools=[],
+        )
+        agent.init_model()
+        log_system_usage("After Load")
+        print("✅ Agent Ready")
+        return agent
+    except Exception as e:
+        print("❌ Error initializing agent:", str(e))
+        traceback.print_exc()
+        raise e
 
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -155,20 +157,28 @@ def create_ui(agent):
         download_output = gr.File(label="Download Full Report")
 
         def analyze(message: str, history: list, files: list):
-            history.append({"role": "user", "content": message})
-            history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
-            yield 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 = [sanitize_utf8(f.result()) for f in as_completed(futures)]
-                    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:
+            try:
+                history.append({"role": "user", "content": message})
+                history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
+                yield 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:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
 3. Note incomplete assessments
@@ -180,10 +190,12 @@ Medical Records:
 ### Potential Oversights:
 """
 
-            try:
+                print("🔎 Generated prompt:")
+                print(prompt)
+
                 full_response = ""
                 finish_detected = False
-                
+
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -193,42 +205,51 @@ Medical Records:
                     call_agent=False,
                     conversation=[]
                 ):
-                    if chunk is None:
-                        continue
-                        
-                    if isinstance(chunk, str):
-                        full_response += chunk
-                    elif isinstance(chunk, list):
-                        chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-                        full_response += chunk_content
-                        if '"name": "Finish"' in chunk_content:
-                            finish_detected = True
-
-                    # Display intermediate response
-                    current_cleaned = clean_response(full_response)
-                    if current_cleaned:
-                        history[-1] = {"role": "assistant", "content": current_cleaned}
-                        yield history, None
-
-                # Final processing of the complete response
-                final_cleaned = clean_response(full_response)
-                
-                if not final_cleaned:
-                    final_cleaned = "⚠️ No clear oversights identified or model output was invalid."
-                
-                # Save report if a file was processed
+                    try:
+                        if chunk is None:
+                            continue
+                        if isinstance(chunk, str):
+                            print("DEBUG: Received string chunk")
+                            full_response += chunk
+                        elif isinstance(chunk, list):
+                            chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
+                            print("DEBUG: Received list chunk:", chunk_content)
+                            full_response += chunk_content
+                            if '"name": "Finish"' in chunk_content:
+                                finish_detected = True
+                        else:
+                            print("DEBUG: Received unknown type chunk", type(chunk))
+                    except Exception as e:
+                        print("❌ Error processing chunk:", str(e))
+                        traceback.print_exc()
+
+                    # Yield intermediate full (raw) response for debugging
+                    history[-1] = {"role": "assistant", "content": full_response}
+                    yield history, None
+
+                # Final processing
+                if not full_response:
+                    full_response = "⚠️ No clear oversights identified or model output was invalid."
+
+                # Save full report
                 report_path = None
                 if file_hash_value:
                     report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                    with open(report_path, "w", encoding="utf-8") as f:
-                        f.write(final_cleaned)
+                    try:
+                        with open(report_path, "w", encoding="utf-8") as f:
+                            f.write(full_response)
+                    except Exception as e:
+                        print("❌ Error saving report:", str(e))
+                        traceback.print_exc()
 
-                history[-1] = {"role": "assistant", "content": final_cleaned}
+                history[-1] = {"role": "assistant", "content": full_response}
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e:
-                print("❌ ERROR:", str(e))
-                history[-1] = {"role": "assistant", "content": f"❌ An error occurred: {str(e)}"}
+                error_message = f"❌ An error occurred in analyze: {str(e)}"
+                print(error_message)
+                traceback.print_exc()
+                history[-1] = {"role": "assistant", "content": error_message}
                 yield history, None
 
         send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
@@ -236,13 +257,17 @@ Medical Records:
     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
+        )
+    except Exception as e:
+        print("❌ Fatal error during launch:", str(e))
+        traceback.print_exc()