Add Ability to Disable Tools (#215)

# PR Description
This PR introduces a new environment variable `ARCADE_DISABLED_TOOLS`.
Tools that are added to this env var are not added to the worker's
`ToolCatalog`. In effect, they are disabled for the worker.

## How to use the `ARCADE_DISABLED_TOOLS` environment variable
* Each tool is separated by a comma.
* For each tool, specify the toolkit name in camel case and the tool
name in camel case.
* Do not include versions. (This is a simple implementation. We can add
disabling specific versions in the future if needed)
* Separate the toolkit name and the tool name with your environment's
tool name separator. By default, the tool name separator is `.`, but you
can override this with the `ARCADE_TOOL_NAME_SEPARATOR` environment
variable.

Correct: `export
ARCADE_DISABLED_TOOLS="Math.Add,Spotify.GetAvailableDevices,Math.Sqrt"`
Incorrect: `export
ARCADE_DISABLED_TOOLS="Math.Add@0.1.0,Spotify.get_available_devices,Sqrt`

arcade/arcade/cli/serve.py

@@ -100,11 +100,6 @@ def serve_default_worker(
     if not toolkits:
         raise RuntimeError("No toolkits found in Python environment.")
 
-    logger.info("Serving the following toolkits:")
-    for toolkit in toolkits:
-        num_tools = sum(len(tools) for tools in toolkit.tools.values())
-        logger.info(f"  - {toolkit.name} ({toolkit.package_name}): {num_tools} tools")
-
     worker_secret = os.environ.get("ARCADE_WORKER_SECRET")
     if not disable_auth and not worker_secret:
         logger.warning(
@@ -124,8 +119,19 @@ def serve_default_worker(
     worker = FastAPIWorker(
         app, secret=worker_secret, disable_auth=disable_auth, otel_meter=otel_handler.get_meter()
     )
+
+    toolkit_tool_counts = {}
     for toolkit in toolkits:
+        prev_tool_count = worker.catalog.get_tool_count()
         worker.register_toolkit(toolkit)
+        new_tool_count = worker.catalog.get_tool_count()
+        toolkit_tool_counts[f"{toolkit.name} ({toolkit.package_name})"] = (
+            new_tool_count - prev_tool_count
+        )
+
+    logger.info("Serving the following toolkits:")
+    for name, tool_count in toolkit_tool_counts.items():
+        logger.info(f"  - {name}: {tool_count} tools")
 
     logger.info("Starting FastAPI server...")
 

arcade/arcade/core/catalog.py

@@ -1,5 +1,8 @@
 import asyncio
 import inspect
+import logging
+import os
+import re
 import typing
 from collections.abc import Iterator
 from dataclasses import dataclass
@@ -49,6 +52,8 @@ from arcade.core.utils import (
     snake_to_pascal_case,
 )
 
+logger = logging.getLogger(__name__)
+
 InnerWireType = Literal["string", "integer", "number", "boolean", "json"]
 WireType = Union[InnerWireType, Literal["array"]]
 
@@ -112,6 +117,34 @@ class ToolCatalog(BaseModel):
 
     _tools: dict[FullyQualifiedName, MaterializedTool] = {}
 
+    _disabled_tools: set[str] = set()
+
+    def __init__(self, **data) -> None:  # type: ignore[no-untyped-def]
+        super().__init__(**data)
+        self._load_disabled_tools()
+
+    def _load_disabled_tools(self) -> None:
+        """Load disabled tools from the environment variable.
+
+        The ARCADE_DISABLED_TOOLS environment variable should contain a
+        comma-separated list of tools that are to be excluded from the
+        catalog.
+
+        The expected format for each disabled tool is:
+        - [CamelCaseToolkitName][TOOL_NAME_SEPARATOR][CamelCaseToolName]
+        """
+        disabled_tools = os.getenv("ARCADE_DISABLED_TOOLS", "").strip().split(",")
+        if not disabled_tools:
+            return
+
+        pattern = re.compile(rf"^[a-zA-Z]+{re.escape(TOOL_NAME_SEPARATOR)}[a-zA-Z]+$")
+
+        for tool in disabled_tools:
+            if not pattern.match(tool):
+                continue
+
+            self._disabled_tools.add(tool.lower())
+
     def add_tool(
         self,
         tool_func: Callable,
@@ -146,6 +179,10 @@ class ToolCatalog(BaseModel):
         if fully_qualified_name in self._tools:
             raise KeyError(f"Tool '{definition.name}' already exists in the catalog.")
 
+        if str(fully_qualified_name).lower() in self._disabled_tools:
+            logger.info(f"Tool '{fully_qualified_name!s}' is disabled and will not be cataloged.")
+            return
+
         self._tools[fully_qualified_name] = MaterializedTool(
             definition=definition,
             tool=tool_func,
@@ -271,6 +308,12 @@ class ToolCatalog(BaseModel):
 
         raise ValueError(f"Tool {name} not found.")
 
+    def get_tool_count(self) -> int:
+        """
+        Get the number of tools in the catalog.
+        """
+        return len(self._tools)
+
     @staticmethod
     def create_tool_definition(
         tool: Callable,

arcade/tests/core/test_catalog.py

@@ -136,3 +136,47 @@ def test_get_tool_by_name_with_invalid_version():
 
     with pytest.raises(ValueError):
         catalog.get_tool_by_name("SampleToolkit.SampleTool", version="2.0.0")
+
+
+def test_load_disabled_tools(monkeypatch):
+    disabled_tools = (
+        "SampleToolkitOne.SampleToolOne,"  # valid
+        + "SampleToolkitOne_SampleToolTwo,"  # invalid
+        + "SampleToolkitTwo.SampleToolThree,"  # valid
+        + "SampleToolkitTwo.SampleToolFour@0.0.1,"  # invalid
+        + "SampleToolkitThree_SampleToolFive@0.0.1,"  # invalid
+        + "SampleToolkitFour.sample_tool_six,"  # invalid
+        + "sample_toolkit5.SampleTool7,"  # invalid
+        + "sample_toolkit6.sample_tool_8"  # invalid
+    )
+    expected_disabled_tools = {
+        "sampletoolkitone.sampletoolone",
+        "sampletoolkittwo.sampletoolthree",
+    }
+
+    monkeypatch.setenv("ARCADE_DISABLED_TOOLS", disabled_tools)
+    catalog = ToolCatalog()
+
+    assert catalog._disabled_tools == expected_disabled_tools
+
+
+def test_add_tool_with_disabled_tool(monkeypatch):
+    monkeypatch.setenv("ARCADE_DISABLED_TOOLS", "SampleToolkitOne.SampleTool")
+    catalog = ToolCatalog()
+
+    catalog.add_tool(sample_tool, "SampleToolkitOne")
+    assert len(catalog._tools) == 0
+
+
+def test_add_tool_with_empty_string_disabled_tools(monkeypatch):
+    monkeypatch.setenv("ARCADE_DISABLED_TOOLS", "")
+    catalog = ToolCatalog()
+    catalog.add_tool(sample_tool, "SampleToolkitOne")
+    assert len(catalog._tools) == 1
+
+
+def test_add_tool_with_whitespace_disabled_tools(monkeypatch):
+    monkeypatch.setenv("ARCADE_DISABLED_TOOLS", "            ")
+    catalog = ToolCatalog()
+    catalog.add_tool(sample_tool, "SampleToolkitOne")
+    assert len(catalog._tools) == 1