Update app.py

app.py

@@ -11,16 +11,11 @@ import shutil
 import re
 import psutil
 import subprocess
-import traceback
-import torch
-
-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.")
-    os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
+# Persistent directory
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
+
 model_cache_dir = os.path.join(persistent_dir, "txagent_models")
 tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
 file_cache_dir = os.path.join(persistent_dir, "cache")
@@ -65,10 +60,7 @@ def extract_priority_pages(file_path: str, max_pages: int = 20) -> str:
                     text_chunks.append(f"=== Page {i} ===\n{page_text.strip()}")
         return "\n\n".join(text_chunks)
     except Exception as e:
-        debug_msg = f"PDF processing error: {str(e)}"
-        print(debug_msg)
-        traceback.print_exc()
-        return debug_msg
+        return f"PDF processing error: {str(e)}"
 
 def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
@@ -99,10 +91,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             f.write(result)
         return result
     except Exception as 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})
+        return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"})
 
 def log_system_usage(tag=""):
     try:
@@ -118,35 +107,29 @@ 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():
-    try:
-        print("\U0001F501 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
+    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 create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -158,40 +141,30 @@ def create_ui(agent):
         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 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:
-                                results.append(sanitize_utf8(future.result()))
-                            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)
-
-                max_content_length = 8000
-                prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
+            history = history + [{"role": "user", "content": message},
+                                 {"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:
 1. List potential missed diagnoses
 2. Flag any medication conflicts
 3. Note incomplete assessments
 4. Highlight abnormal results needing follow-up
 Medical Records:
-{extracted[:max_content_length]}
+{extracted[:12000]}
 ### Potential Oversights:
 """
-
-                full_response = ""
-                response_chunks = []
-
+            response = ""
+            try:
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -199,68 +172,26 @@ Medical Records:
                     max_new_tokens=2048,
                     max_token=4096,
                     call_agent=False,
-                    conversation=[]
+                    conversation=[],
                 ):
-                    try:
-                        chunk_content = ""
-                        if isinstance(chunk, str):
-                            chunk_content = chunk
-                        elif hasattr(chunk, 'content'):
-                            chunk_content = chunk.content
-                        elif isinstance(chunk, list):
-                            chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-
-                        if not chunk_content:
-                            continue
-
-                        response_chunks.append(chunk_content)
-                        full_response = "".join(response_chunks)
-
-                        display_response = re.split(r"\\[TOOL_CALLS\\].*?$", full_response, flags=re.DOTALL)[0].strip()
-                        display_response = display_response.replace('[TxAgent]', '').strip()
-
-                        if len(history) > 1 and history[-2]["role"] == "assistant" and history[-2]["content"] == display_response:
-                            pass
-                        else:
-                            if len(history) > 0 and history[-1]["role"] == "assistant":
-                                history[-1]["content"] = display_response
-                            else:
-                                history.append({"role": "assistant", "content": display_response})
-
-                        yield history, None
-                    except Exception as e:
-                        print("❌ Error processing chunk:", str(e))
-                        traceback.print_exc()
+                    if chunk is None:
                         continue
+                    if isinstance(chunk, str):
+                        response += chunk
+                    elif isinstance(chunk, list):
+                        response += "".join([c.content for c in chunk if hasattr(c, 'content') and c.content])
 
-                if not full_response:
-                    full_response = "⚠️ No clear oversights identified or model output was invalid."
-                else:
-                    full_response = re.split(r"\\[TOOL_CALLS\\].*?$", full_response, flags=re.DOTALL)[0].strip()
-                    full_response = full_response.replace('[TxAgent]', '').strip()
-
-                report_path = None
-                if file_hash_value:
-                    report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
-                    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()
-
-                if len(history) > 0 and history[-1]["role"] == "assistant":
-                    history[-1]["content"] = full_response
-                else:
-                    history.append({"role": "assistant", "content": full_response})
+                cleaned = response.split("[TOOL_CALLS]")[0].strip()
+                if not cleaned:
+                    cleaned = "No clear oversights identified. Recommend comprehensive review."
 
+                history[-1] = {"role": "assistant", "content": cleaned}
+                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e:
-                error_message = f"❌ An error occurred in analyze: {str(e)}"
-                print(error_message)
-                traceback.print_exc()
-                history.append({"role": "assistant", "content": error_message})
+                print("🚨 ERROR:", e)
+                history[-1] = {"role": "assistant", "content": f"❌ Error occurred: {str(e)}"}
                 yield history, None
 
         send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
@@ -268,17 +199,13 @@ Medical Records:
     return demo
 
 if __name__ == "__main__":
-    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()
+    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
+    )