diff --git a/arcade/arcade/cli/serve.py b/arcade/arcade/cli/serve.py index f07b6b15..d83d18d3 100644 --- a/arcade/arcade/cli/serve.py +++ b/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...") diff --git a/arcade/arcade/core/catalog.py b/arcade/arcade/core/catalog.py index d150f7bc..4602a930 100644 --- a/arcade/arcade/core/catalog.py +++ b/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, diff --git a/arcade/tests/core/test_catalog.py b/arcade/tests/core/test_catalog.py index 6254ce5f..45eb2524 100644 --- a/arcade/tests/core/test_catalog.py +++ b/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