Update app.py

app.py

@@ -3,7 +3,6 @@ import logging
 import os
 import multiprocessing
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -81,10 +80,20 @@ if __name__ == "__main__":
     multiprocessing.set_start_method("spawn", force=True)
 
     from txagent import TxAgent
+    import tooluniverse.universe as universe_mod
+
+    # === Patch ToolUniverse.infer_tool_embeddings to prevent exit ===
+    original_infer = universe_mod.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
-    from patched_tooluniverse import PatchedToolUniverse
 
-    logger.info("🔥 Initializing patched TxAgent with tool embedding skip")
+    logger.info("🔥 Initializing patched TxAgent with embedding skip (manual patch)")
 
     tool_files = {
         "opentarget": str(files('tooluniverse.data').joinpath('opentarget_tools.json')),
@@ -110,8 +119,7 @@ if __name__ == "__main__":
         seed=42,
         enable_checker=True,
         enable_chat=False,
-        additional_default_tools=["DirectResponse", "RequireClarification"],
-        tooluniverse_class=PatchedToolUniverse  # ✅ Custom patch!
+        additional_default_tools=["DirectResponse", "RequireClarification"]
     )
 
     if not os.path.exists(TOOL_CACHE_PATH):
@@ -120,4 +128,4 @@ if __name__ == "__main__":
         with open(TOOL_CACHE_PATH, "w") as f:
             f.write("done")
     else:
-        tx_app.init_model(skip_tool_embedding=True)  # only if this param is supported
+        tx_app.init_model(skip_tool_embedding=True)