Update app.py

app.py

@@ -1,3 +1,5 @@
+# Optimized app.py with lazy loading and preloading thread, fixed chatbot format and startup error handling
+
 import sys
 import os
 import gradio as gr
@@ -29,8 +31,7 @@ os.makedirs("/data/vllm_cache", exist_ok=True)
 # Lazy loading of heavy dependencies
 def lazy_load_agent():
     from txagent.txagent import TxAgent
-    
-    # Initialize agent with optimized settings
+
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -71,11 +72,11 @@ def process_file(file_path: str, file_type: str) -> str:
     try:
         h = file_hash(file_path)
         cache_path = f"/data/file_cache/{h}.json"
-        
+
         if os.path.exists(cache_path):
             with open(cache_path, "r", encoding="utf-8") as f:
                 return f.read()
-                
+
         if file_type == "pdf":
             content = extract_priority_pages(file_path)
             result = json.dumps({"filename": os.path.basename(file_path), "content": content})
@@ -91,7 +92,7 @@ def process_file(file_path: str, file_type: str) -> str:
         with open(cache_path, "w", encoding="utf-8") as f:
             f.write(result)
         return result
-        
+
     except Exception as e:
         return json.dumps({"error": str(e)})
 
@@ -100,7 +101,7 @@ def format_response(response: str) -> str:
     if "Based on the medical records provided" in response:
         parts = response.split("Based on the medical records provided")
         response = "Based on the medical records provided" + parts[-1]
-    
+
     replacements = {
         "1. **Missed Diagnoses**:": "### 🔍 Missed Diagnoses",
         "2. **Medication Conflicts**:": "\n### 💊 Medication Conflicts",
@@ -108,30 +109,27 @@ def format_response(response: str) -> str:
         "4. **Abnormal Results Needing Follow-up**:": "\n### ⚠️ Abnormal Results Needing Follow-up",
         "Overall, the patient's medical records": "\n### 📝 Overall Assessment"
     }
-    
+
     for old, new in replacements.items():
         response = response.replace(old, new)
-    
+
     return response
 
 def analyze_files(message: str, history: List, files: List):
     try:
-        # Wait for agent to load if not ready
         while agent is None:
             time.sleep(0.1)
-            
-        # Append user message to history in correct format
+
         history.append([message, None])
         yield history, None
-        
-        # Process files in parallel
+
         extracted_data = ""
         if files:
             with ThreadPoolExecutor(max_workers=4) as executor:
                 futures = [executor.submit(process_file, f.name, f.name.split(".")[-1].lower()) 
-                          for f in files if hasattr(f, 'name')]
+                           for f in files if hasattr(f, 'name')]
                 extracted_data = "\n".join(f.result() for f in as_completed(futures))
-        
+
         prompt = f"""Review these medical records:
 {extracted_data[:10000]}
 
@@ -142,7 +140,7 @@ Identify:
 4. Abnormal results needing follow-up
 
 Analysis:"""
-        
+
         response = ""
         for chunk in agent.run_gradio_chat(
             message=prompt,
@@ -155,88 +153,45 @@ Analysis:"""
                 response += chunk
             elif isinstance(chunk, list):
                 response += "".join(getattr(c, 'content', '') for c in chunk)
-            
+
             formatted = format_response(response)
             if formatted.strip():
                 history[-1][1] = formatted
                 yield history, None
-        
+
         final_output = format_response(response) or "No clear oversights identified."
         history[-1][1] = final_output
         yield history, None
-        
+
     except Exception as e:
         history[-1][1] = f"❌ Error: {str(e)}"
         yield history, None
 
-# Create optimized UI with better layout
-with gr.Blocks(title="Clinical Oversight Assistant", css="""
-    .gradio-container {
-        max-width: 1200px !important;
-        margin: auto;
-    }
-    .container {
-        max-width: 1200px !important;
-    }
-    .chatbot {
-        min-height: 500px;
-    }
-""") as demo:
+# UI definition
+with gr.Blocks(title="Clinical Oversight Assistant") as demo:
     gr.Markdown("""
-    <div style='text-align: center; margin-bottom: 20px;'>
-        <h1 style='margin-bottom: 10px;'>🩺 Clinical Oversight Assistant</h1>
+    <div style='text-align: center;'>
+        <h1>🩺 Clinical Oversight Assistant</h1>
         <p>Upload medical records to analyze for potential oversights in patient care</p>
     </div>
     """)
-    
+
     with gr.Row():
-        with gr.Column(scale=1, min_width=400):
-            file_upload = gr.File(
-                label="Upload Medical Records",
-                file_types=[".pdf", ".csv", ".xls", ".xlsx"],
-                file_count="multiple",
-                height=100
-            )
-            query = gr.Textbox(
-                label="Your Query",
-                placeholder="Ask about potential oversights...",
-                lines=3
-            )
+        with gr.Column(scale=1):
+            file_upload = gr.File(label="Upload Medical Records", file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
+            query = gr.Textbox(label="Your Query", placeholder="Ask about potential oversights...", lines=3)
             submit = gr.Button("Analyze", variant="primary")
-            
-            gr.Examples(
-                examples=[
-                    ["What potential diagnoses might have been missed?"],
-                    ["Are there any medication conflicts I should be aware of?"],
-                    ["What assessments appear incomplete in these records?"]
-                ],
-                inputs=query,
-                label="Example Queries"
-            )
-            
-        with gr.Column(scale=2, min_width=600):
-            chatbot = gr.Chatbot(
-                label="Analysis Results",
-                height=600,
-                bubble_full_width=False,
-                show_copy_button=True
-            )
-    
-    submit.click(
-        analyze_files,
-        inputs=[query, chatbot, file_upload],
-        outputs=[chatbot, gr.File(visible=False)]
-    )
-    
-    query.submit(
-        analyze_files,
-        inputs=[query, chatbot, file_upload],
-        outputs=[chatbot, gr.File(visible=False)]
-    )
+            gr.Examples([
+                ["What potential diagnoses might have been missed?"],
+                ["Are there any medication conflicts I should be aware of?"],
+                ["What assessments appear incomplete in these records?"]
+            ], inputs=query)
+
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(label="Analysis Results", height=600, type="messages")
+
+    submit.click(analyze_files, inputs=[query, chatbot, file_upload], outputs=[chatbot, gr.File(visible=False)])
+    query.submit(analyze_files, inputs=[query, chatbot, file_upload], outputs=[chatbot, gr.File(visible=False)])
 
 if __name__ == "__main__":
-    demo.queue(concurrency_count=1).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)