Update app.py

app.py

@@ -1,10 +1,12 @@
 import gradio as gr
 import logging
+import os
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-tx_app = None  # global agent
+tx_app = None
+TOOL_CACHE_PATH = "/home/user/.cache/tool_embeddings_done"  # flag file for skip
 
 def respond(message, chat_history, temperature, max_new_tokens, max_tokens, multi_agent, conversation_state, max_round):
     global tx_app
@@ -42,7 +44,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 app at module level so Hugging Face Spaces can find it
+# === Define Gradio interface ===
 with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
     gr.Markdown("# 🧠 TxAgent Biomedical Assistant")
 
@@ -72,16 +74,15 @@ with gr.Blocks(title="TxAgent Biomedical Assistant") as app:
         chatbot
     )
 
-# 🔥 Safely initialize vLLM inside __main__
+# === Safe model init block for vLLM + Hugging Face ===
 if __name__ == "__main__":
     import multiprocessing
     multiprocessing.set_start_method("spawn", force=True)
 
-    import torch
     from txagent import TxAgent
     from importlib.resources import files
 
-    logger.info("🔥 Initializing TxAgent safely in __main__")
+    logger.info("🔥 Initializing TxAgent inside __main__...")
 
     tool_files = {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
@@ -90,6 +91,7 @@ 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",
@@ -110,5 +112,15 @@ if __name__ == "__main__":
         additional_default_tools=["DirectResponse", "RequireClarification"]
     )
 
-    tx_app.init_model()
-    logger.info("✅ TxAgent initialized.")
+    # ✅ 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
+        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!")