Update app.py

app.py

@@ -14,14 +14,14 @@ import logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# PyTorch cleanup
+# Cleanup
 def cleanup():
     if dist.is_initialized():
         logger.info("Cleaning up PyTorch distributed process group")
         dist.destroy_process_group()
 atexit.register(cleanup)
 
-# Directories
+# Cache dirs
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
 model_cache_dir = os.path.join(persistent_dir, "txagent_models")
@@ -105,25 +105,6 @@ Please analyze these clinical notes and provide:
 Provide a structured response with clear medical reasoning.
 """
 
-def validate_tool_file(tool_name: str, tool_path: str) -> bool:
-    try:
-        if not os.path.exists(tool_path):
-            logger.error(f"Missing tool file: {tool_path}")
-            return False
-        with open(tool_path, 'r') as f:
-            tool_data = json.load(f)
-        if isinstance(tool_data, list):
-            return all(isinstance(item, dict) and 'name' in item for item in tool_data)
-        elif isinstance(tool_data, dict):
-            if 'tools' in tool_data:
-                return all(isinstance(item, dict) and 'name' in item for item in tool_data['tools'])
-            return 'name' in tool_data
-        logger.error(f"Invalid format in tool: {tool_name}")
-        return False
-    except Exception as e:
-        logger.error(f"Error in {tool_name}: {e}")
-        return False
-
 def init_agent() -> TxAgent:
     new_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
     if not os.path.exists(new_tool_path):
@@ -133,6 +114,7 @@ def init_agent() -> TxAgent:
                 "description": "Default tool",
                 "tools": [{"name": "dummy_tool", "description": "test", "version": "1.0"}]
             }, f)
+
     tool_files = {
         'opentarget': '/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/opentarget_tools.json',
         'fda_drug_label': '/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/fda_drug_labeling_tools.json',
@@ -140,13 +122,35 @@ def init_agent() -> TxAgent:
         'monarch': '/home/user/.pyenv/versions/3.10.17/lib/python3.10/site-packages/tooluniverse/data/monarch_tools.json',
         'new_tool': new_tool_path
     }
-    valid_tools = {k: v for k, v in tool_files.items() if validate_tool_file(k, v)}
-    if not valid_tools:
-        raise ValueError("No valid tool files")
+
+    validated = {}
+    for name, path in tool_files.items():
+        try:
+            with open(path, 'r') as f:
+                data = json.load(f)
+            if isinstance(data, dict) and 'tools' in data:
+                tools = data['tools']
+            elif isinstance(data, list):
+                tools = data
+            elif isinstance(data, dict) and 'name' in data:
+                tools = [data]
+            else:
+                logger.warning(f"Skipping {name}: bad structure")
+                continue
+            if all(isinstance(t, dict) and 'name' in t for t in tools):
+                validated[name] = path
+            else:
+                logger.warning(f"Skipping {name}: items malformed")
+        except Exception as e:
+            logger.error(f"Invalid tool {name}: {e}")
+
+    if not validated:
+        raise ValueError("No valid tools to load")
+
     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=valid_tools,
+        tool_files_dict=validated,
         force_finish=True,
         enable_checker=True,
         step_rag_num=4,