Update app.py

app.py

@@ -1,57 +1,23 @@
 import gradio as gr
 import logging
-from txagent import TxAgent
-from tooluniverse import ToolUniverse
-from importlib.resources import files
+
+# Delay heavy imports until later to avoid multiprocessing conflicts
+tx_app = None  # Global agent instance
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-tx_app = None  # Global TxAgent instance
-
-def init_txagent():
-    logger.info("🔥 Initializing TxAgent...")
-
-    tool_files = {
-        "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
-        "fda_drug_label": str(files('tooluniverse.data').joinpath('fda_drug_labeling_tools.json')),
-        "special_tools": str(files('tooluniverse.data').joinpath('special_tools.json')),
-        "monarch": str(files('tooluniverse.data').joinpath('monarch_tools.json'))
-    }
-
-    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=tool_files,
-        enable_finish=True,
-        enable_rag=True,
-        enable_summary=False,
-        init_rag_num=0,
-        step_rag_num=10,
-        summary_mode='step',
-        summary_skip_last_k=0,
-        summary_context_length=None,
-        force_finish=True,
-        avoid_repeat=True,
-        seed=42,
-        enable_checker=True,
-        enable_chat=False,
-        additional_default_tools=["DirectResponse", "RequireClarification"]
-    )
-
-    agent.init_model()
-    logger.info("✅ TxAgent fully initialized")
-    return agent
-
+# ========== Dummy Response (will be replaced by real agent later) ==========
 def respond(message, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round):
     global tx_app
     if tx_app is None:
-        return chat_history + [("", "⚠️ Model not ready yet. Please wait a few seconds and try again.")]
+        return chat_history + [("", "⚠️ Model is still loading. Please wait a few seconds and try again.")]
 
     try:
-        if not isinstance(message, str) or len(message.strip()) <= 10:
-            return chat_history + [("", "Please provide a valid message longer than 10 characters.")]
-
+        if not isinstance(message, str) or len(message.strip()) < 10:
+            return chat_history + [("", "Please enter a longer message.")]
+        
+        # Convert chat format if needed
         if chat_history and isinstance(chat_history[0], dict):
             chat_history = [(h["role"], h["content"]) for h in chat_history if "role" in h and "content" in h]
 
@@ -65,7 +31,7 @@ def respond(message, chat_history, temperature, max_new_tokens, max_tokens, mult
             call_agent=multi_agent,
             conversation=conversation_state,
             max_round=max_round,
-            seed=42
+            seed=42,
         ):
             if isinstance(chunk, dict):
                 response += chunk.get("content", "")
@@ -77,15 +43,20 @@ def respond(message, chat_history, temperature, max_new_tokens, max_tokens, mult
             yield chat_history + [("user", message), ("assistant", response)]
 
     except Exception as e:
-        logger.error(f"Error in respond function: {str(e)}")
-        yield chat_history + [("", f"⚠️ Error: {str(e)}")]
+        logger.error(f"Respond error: {e}")
+        yield chat_history + [("", f"⚠️ Error: {e}")]
 
-# ✅ Top-level app object that HF Spaces can detect
+# ========== Gradio UI ==========
 with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
     gr.Markdown("# 🧠 TxAgent Biomedical Assistant")
 
     chatbot = gr.Chatbot(label="Conversation", height=600, type="messages")
-    msg = gr.Textbox(label="Your medical query", placeholder="Enter your biomedical question...", lines=3)
+
+    msg = gr.Textbox(
+        label="Your medical query",
+        placeholder="Enter your biomedical question...",
+        lines=3
+    )
 
     with gr.Row():
         temp = gr.Slider(0, 1, value=0.3, label="Temperature")
@@ -103,22 +74,53 @@ with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
         [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
         chatbot
     )
-
     clear.click(lambda: [], None, chatbot)
-
     msg.submit(
         respond,
         [msg, chatbot, temp, max_new_tokens, max_tokens, multi_agent, conversation_state, max_rounds],
         chatbot
     )
 
-    # ✅ hidden init trigger on page load
-    hidden_button = gr.Button(visible=False)
+    # === Hidden trigger to load model safely on app start ===
+    init_button = gr.Button(visible=False)
 
-    def initialize_agent():
+    def load_model():
         global tx_app
-        tx_app = init_txagent()
+        import torch
+        from txagent import TxAgent
+        from importlib.resources import files
+
+        logger.info("🔧 Loading full TxAgent model...")
+
+        tool_files = {
+            "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
+            "fda_drug_label": str(files('tooluniverse.data').joinpath('fda_drug_labeling_tools.json')),
+            "special_tools": str(files('tooluniverse.data').joinpath('special_tools.json')),
+            "monarch": str(files('tooluniverse.data').joinpath('monarch_tools.json'))
+        }
+
+        tx_app = 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=tool_files,
+            enable_finish=True,
+            enable_rag=True,
+            enable_summary=False,
+            init_rag_num=0,
+            step_rag_num=10,
+            summary_mode='step',
+            summary_skip_last_k=0,
+            summary_context_length=None,
+            force_finish=True,
+            avoid_repeat=True,
+            seed=42,
+            enable_checker=True,
+            enable_chat=False,
+            additional_default_tools=["DirectResponse", "RequireClarification"]
+        )
+
+        tx_app.init_model()
+        logger.info("✅ Model initialized successfully")
         return gr.update(visible=False)
 
-    app.load(hidden_button.click(fn=initialize_agent))
-
+    app.load(init_button.click(fn=load_model))