Fix tool filtering with rename tool (#79)

Signed-off-by: Nathalie Jonathan <[email protected]>

Author: nathaliellenaa

src/tools/tool_filter.py

@@ -67,8 +67,11 @@ def process_tool_filter(
         tool_registry: The tool registry to filter.
     """
     try:
-        # Create case-insensitive lookup
-        tool_registry_lower = {k.lower(): k for k in tool_registry.keys()}
+        # Create display name lookup
+        display_name = {
+            tool_info.get('display_name', '').lower(): k
+            for k, tool_info in tool_registry.items()
+        }
 
         # Initialize collections
         category_to_tools = {}
@@ -120,7 +123,7 @@ def process_tool_filter(
         )
 
         # Get current tool names after allow_write filtering
-        current_tool_names = list(tool_registry.keys())
+        current_tool_names = [tool['display_name'] for tool in tool_registry.values()]
         disabled_tools_from_regex = process_regex_patterns(
             disabled_tools_regex_list, current_tool_names
         )
@@ -130,17 +133,13 @@ def process_tool_filter(
             # Validate and collect all disabled tools
             all_disabled_tools = set()
             all_disabled_tools.update(
-                validate_tools(disabled_tools_list, tool_registry_lower, 'disabled_tools')
+                validate_tools(disabled_tools_list, display_name, 'disabled_tools')
             )
             all_disabled_tools.update(
-                validate_tools(
-                    disabled_tools_from_categories, tool_registry_lower, 'disabled_categories'
-                )
+                validate_tools(disabled_tools_from_categories, display_name, 'disabled_categories')
             )
             all_disabled_tools.update(
-                validate_tools(
-                    disabled_tools_from_regex, tool_registry_lower, 'disabled_tools_regex'
-                )
+                validate_tools(disabled_tools_from_regex, display_name, 'disabled_tools_regex')
             )
 
             # Remove tools in the disabled list
@@ -150,8 +149,9 @@ def process_tool_filter(
 
         # Log results
         source = filter_path if filter_path else 'environment variables'
+        tool_display_names = [tool['display_name'] for tool in tool_registry.values()]
         logging.info(f'Applied tool filter from {source}')
-        logging.info(f'Available tools after filtering: {list(tool_registry.keys())}')
+        logging.info(f'Available tools after filtering: {tool_display_names}')
 
     except Exception as e:
         logging.error(f'Error processing tool filter: {str(e)}')
@@ -201,9 +201,10 @@ def get_tools(tool_registry: dict, mode: str = 'single', config_file_path: str =
         **{k: v for k, v in env_config.items() if not config_file_path},
     )
 
-    for name, info in TOOL_REGISTRY.items():
+    for name, info in tool_registry.items():
         # Create a copy to avoid modifying the original tool info
         tool_info = info.copy()
+        tool_name = tool_info['display_name']
 
         # If tool is not compatible with the current OpenSearch version, skip, don't enable
         if not is_tool_compatible(version, info):
@@ -218,6 +219,6 @@ def get_tools(tool_registry: dict, mode: str = 'single', config_file_path: str =
                 schema['properties'].pop(field, None)
         tool_info['input_schema'] = schema
 
-        enabled[name] = tool_info
+        enabled[tool_name] = tool_info
 
     return enabled

src/tools/utils.py

@@ -51,13 +51,15 @@ def load_yaml_config(filter_path):
         return None
 
 
-def validate_tools(tool_list, registry_lookup, source_name):
+def validate_tools(tool_list, display_lookup, source_name):
     """Validate tools against registry and return valid tools."""
     valid_tools = set()
     for tool in tool_list:
         tool_lower = tool.lower()
-        if tool_lower in registry_lookup:
-            valid_tools.add(tool_lower)
+        # Check if it matches tool display name
+        if tool_lower in display_lookup:
+            actual_tool = display_lookup[tool_lower]
+            valid_tools.add(actual_tool.lower())
         else:
             logging.warning(f"Ignoring invalid tool from '{source_name}': '{tool}'")
     return valid_tools

tests/tools/test_tool_filters.py

@@ -9,6 +9,7 @@
 # A dictionary for mocking TOOL_REGISTRY
 MOCK_TOOL_REGISTRY = {
     'ListIndexTool': {
+        'display_name': 'ListIndexTool',
         'description': 'List indices',
         'input_schema': {'type': 'object', 'properties': {'param1': {'type': 'string'}}},
         'function': MagicMock(),
@@ -17,6 +18,7 @@
         'max_version': '3.0.0',
     },
     'SearchIndexTool': {
+        'display_name': 'SearchIndexTool',
         'description': 'Search an index',
         'input_schema': {
             'type': 'object',
@@ -162,6 +164,7 @@ def test_get_tools_single_mode_handles_missing_properties(self, mock_patches):
         # Create tool with missing properties
         tool_without_properties = {
             'ListIndexTool': {
+                'display_name': 'ListIndexTool',
                 'description': 'List indices',
                 'input_schema': {'type': 'object', 'title': 'ListIndexArgs'},
                 'function': MagicMock(),
@@ -213,13 +216,15 @@ class TestProcessToolFilter:
     def setup_method(self):
         """Set up a fresh copy of the tool registry for each test."""
         self.tool_registry = {
-            'ListIndexTool': {'http_methods': 'GET'},
-            'SearchIndexTool': {'http_methods': 'GET, POST'},
-            'MsearchTool': {'http_methods': 'GET, POST'},
-            'ExplainTool': {'http_methods': 'GET, POST'},
-            'ClusterHealthTool': {'http_methods': 'GET'},
-            'IndicesCreateTool': {'http_methods': 'PUT'},
-            'IndicesStatsTool': {'http_methods': 'GET'},
+            'ListIndexTool': {'display_name': 'ListIndexTool', 'http_methods': 'GET'},
+            'SearchIndexTool': {'display_name': 'SearchIndexTool', 'http_methods': 'GET, POST'},
+            'MsearchTool': {'display_name': 'MsearchTool', 'http_methods': 'GET, POST'},
+            'ExplainTool': {'display_name': 'ExplainTool', 'http_methods': 'GET, POST'},
+            'ClusterHealthTool': {'display_name': 'ClusterHealthTool', 'http_methods': 'GET'},
+            'IndicesCreateTool': {'display_name': 'IndicesCreateTool', 'http_methods': 'PUT'},
+            'IndicesStatsTool': {'display_name': 'IndicesStatsTool', 'http_methods': 'GET'},
+            'CountTool': {'display_name': 'CustomCountTool', 'http_methods': 'GET'},
+            'ListModelTool': {'display_name': 'ModelListTool', 'http_methods': 'GET'},
         }
         self.category_to_tools = {
             'critical': ['SearchIndexTool', 'ExplainTool'],
@@ -269,3 +274,23 @@ def test_process_tool_filter_env(self, caplog):
         assert 'MsearchTool' in self.tool_registry
         assert 'SearchIndexTool' not in self.tool_registry  # In disabled_tools_regex
         assert 'ExplainTool' not in self.tool_registry  # In disabled_tools
+
+    def test_process_tool_filter_rename_tool(self):
+        """Test processing tool filtering with tool renaming feature"""
+        process_tool_filter(
+            tool_registry=self.tool_registry,
+            disabled_tools='CountTool',
+            disabled_tools_regex='list.*',
+            allow_write=True
+        )
+        assert 'CountTool' in self.tool_registry  # Renamed to CustomCountTool
+        assert 'ListModelTool' in self.tool_registry  # Renamed to ModelListTool
+
+        process_tool_filter(
+            tool_registry=self.tool_registry,
+            disabled_tools='CustomCountTool', 
+            disabled_tools_regex='model.*',
+            allow_write=True
+        )
+        assert 'CustomCountTool' not in self.tool_registry
+        assert 'ModelListTool' not in self.tool_registry