Update app.py

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
 import logging
 import os
+import multiprocessing
 
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 tx_app = None
-TOOL_CACHE_PATH = "/home/user/.cache/tool_embeddings_done"  # flag file for skip
+TOOL_CACHE_PATH = "/home/user/.cache/tool_embeddings_done"
 
 def respond(message, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round):
     global tx_app
@@ -44,7 +46,7 @@ def respond(message, chat_history, temperature, max_new_tokens, max_tokens, mult
         logger.error(f"Respond error: {e}")
         yield chat_history + [("", f"⚠️ Error: {e}")]
 
-# === Define Gradio interface ===
+# === Gradio UI ===
 with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
     gr.Markdown("# 🧠 TxAgent Biomedical Assistant")
 
@@ -74,15 +76,15 @@ with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
         chatbot
     )
 
-# === Safe model init block for vLLM + Hugging Face ===
+# === Safe model init ===
 if __name__ == "__main__":
-    import multiprocessing
     multiprocessing.set_start_method("spawn", force=True)
 
     from txagent import TxAgent
     from importlib.resources import files
+    from patched_tooluniverse import PatchedToolUniverse
 
-    logger.info("🔥 Initializing TxAgent inside __main__...")
+    logger.info("🔥 Initializing patched TxAgent with tool embedding skip")
 
     tool_files = {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
@@ -91,7 +93,6 @@ if __name__ == "__main__":
         "monarch": str(files('tooluniverse.data').joinpath('monarch_tools.json'))
     }
 
-    # Initialize agent
     tx_app = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
@@ -109,18 +110,14 @@ if __name__ == "__main__":
         seed=42,
         enable_checker=True,
         enable_chat=False,
-        additional_default_tools=["DirectResponse", "RequireClarification"]
+        additional_default_tools=["DirectResponse", "RequireClarification"],
+        tooluniverse_class=PatchedToolUniverse  # ✅ Custom patch!
     )
 
-    # ✅ Only do tool embedding the first time
     if not os.path.exists(TOOL_CACHE_PATH):
-        logger.info("🔧 First run: running full model + embedding")
-        tx_app.init_model()  # runs full setup
+        tx_app.init_model()
         os.makedirs(os.path.dirname(TOOL_CACHE_PATH), exist_ok=True)
         with open(TOOL_CACHE_PATH, "w") as f:
             f.write("done")
     else:
-        logger.info("⚡️ Skipping tool embedding (cached)...")
-        tx_app.init_model(skip_tool_embedding=True)  # assumes this param is supported
-
-    logger.info("✅ TxAgent is ready!")
+        tx_app.init_model(skip_tool_embedding=True)  # only if this param is supported