Update app.py

app.py

@@ -3,12 +3,14 @@ import logging
 import os
 import multiprocessing
 
+# Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 tx_app = None
 TOOL_CACHE_PATH = "/home/user/.cache/tool_embeddings_done"
 
+# Chatbot response function
 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:
@@ -50,7 +52,7 @@ 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="Type here...", lines=3)
+    msg = gr.Textbox(label="Your medical query", placeholder="Type your biomedical question...", lines=3)
 
     with gr.Row():
         temp = gr.Slider(0, 1, value=0.3, label="Temperature")
@@ -75,25 +77,24 @@ with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
         chatbot
     )
 
-# === Safe model init ===
+# === Safe model initialization ===
 if __name__ == "__main__":
     multiprocessing.set_start_method("spawn", force=True)
 
+    import tooluniverse
     from txagent import TxAgent
-    import tooluniverse.universe as universe_mod
+    from importlib.resources import files
 
-    # === Patch ToolUniverse.infer_tool_embeddings to prevent exit ===
-    original_infer = universe_mod.ToolUniverse.infer_tool_embeddings
+    # ✅ Patch ToolUniverse to prevent exit() after embedding
+    original_infer = tooluniverse.ToolUniverse.infer_tool_embeddings
 
     def patched_infer(self, *args, **kwargs):
         original_infer(self, *args, **kwargs)
         print("✅ Patched: Skipping forced exit() after embedding.")
 
-    universe_mod.ToolUniverse.infer_tool_embeddings = patched_infer
-
-    from importlib.resources import files
+    tooluniverse.ToolUniverse.infer_tool_embeddings = patched_infer
 
-    logger.info("🔥 Initializing patched TxAgent with embedding skip (manual patch)")
+    logger.info("🔥 Initializing TxAgent...")
 
     tool_files = {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
@@ -122,6 +123,7 @@ if __name__ == "__main__":
         additional_default_tools=["DirectResponse", "RequireClarification"]
     )
 
+    # 🚀 Run full embedding once, then cache
     if not os.path.exists(TOOL_CACHE_PATH):
         tx_app.init_model()
         os.makedirs(os.path.dirname(TOOL_CACHE_PATH), exist_ok=True)
@@ -129,3 +131,5 @@ if __name__ == "__main__":
             f.write("done")
     else:
         tx_app.init_model(skip_tool_embedding=True)
+
+    logger.info("✅ TxAgent ready.")