Optimized Tools._filter_custom_tools

[Saga4]

The optimized code replaces exception-based control flow with direct dictionary lookups, achieving a 201% speedup.

📄 202% (2.02x) speedup for Tools._filter_custom_tools in python/composio/core/models/tools.py

⏱️ Runtime : 1.14 milliseconds → 379 microseconds (best of 246 runs)

📝 Explanation and details

The optimized code replaces exception-based control flow with direct dictionary lookups, achieving a 201% speedup.

Key optimization: The original code used a try/except pattern to check if a tool exists in the custom tools registry:

try:
    _custom_tools.append(self._custom_tools[tool].info)
except KeyError:
    _tools.append(tool)

The optimized version eliminates exceptions by:

1. Caching the registry reference - custom_tools_registry = self._custom_tools.custom_tools_registry avoids repeated attribute lookups
2. Using dict.get() instead of KeyError handling - custom_tool = custom_tools_registry.get(tool) returns None for missing keys without raising exceptions

Why this is significantly faster:

• Exception handling overhead eliminated - The original code raised and caught KeyError exceptions for every non-custom tool (which appears to be common based on the test results)
• Reduced attribute access - The registry lookup self._custom_tools.custom_tools_registry happens once per method call instead of once per tool
• Fast-path dictionary operations - dict.get() is optimized for the common case of missing keys

Performance characteristics:
The optimization shows excellent scalability - larger improvements with more tools (326% faster for 500 standard tools, 240% faster for 1000 mixed tools). This suggests the original exception-based approach had O(n) overhead that scaled poorly, while the optimized version maintains efficient dictionary lookup performance.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 58 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import functools
import typing as t
from types import SimpleNamespace

# imports
import pytest  # used for our unit tests
from composio.core.models.tools import Tools


# Dummy classes for dependencies
class Tool:
    def __init__(self, name):
        self.name = name

class CustomTool:
    def __init__(self, slug, info):
        self.slug = slug
        self.info = info

CustomToolProtocol = t.Callable[..., t.Any]
CustomToolWithProxyProtocol = t.Callable[..., t.Any]

class HttpClient:
    pass

class FileHelper:
    def __init__(self, client, outdir=None):
        pass

class Resource:
    pass

TProvider = t.TypeVar("TProvider")
ToolkitVersionParam = t.Any

class CustomTools:
    """CustomTools class

    Used to manage custom tools created by users."""

    def __init__(self, client: HttpClient):
        self.client = client
        self.custom_tools_registry: dict[str, CustomTool] = {}

    def __getitem__(self, slug: str) -> CustomTool:
        """Get a custom tool by its slug."""
        return self.custom_tools_registry[slug]

    @t.overload
    def register(self, f: CustomToolProtocol) -> CustomTool: ...

    @t.overload
    def register(
        self, *, toolkit: t.Optional[str] = None
    ) -> t.Callable[[CustomToolWithProxyProtocol], CustomTool]: ...

    def register(
        self,
        f: t.Optional[CustomToolProtocol] = None,
        *,
        toolkit: t.Optional[str] = None,
    ) -> t.Union[
        CustomTool,
        t.Callable[[CustomToolProtocol], CustomTool],
        t.Callable[[CustomToolWithProxyProtocol], CustomTool],
    ]:
        """Register a custom tool."""
        if f is not None:
            return self._wrap_tool(f, toolkit)
        return functools.partial(self._wrap_tool, toolkit=toolkit)

    def _wrap_tool(self, f, toolkit):
        # Dummy implementation for testing
        return CustomTool(f.__name__, Tool(f.__name__))
from composio.core.models.tools import Tools

# ------------------ UNIT TESTS ------------------

@pytest.fixture
def tools_instance():
    """Fixture to provide a Tools instance with some custom tools registered."""
    client = HttpClient()
    provider = "dummy_provider"
    tools = Tools(client, provider)
    # Register some custom tools
    tools._custom_tools.custom_tools_registry = {
        "custom1": CustomTool("custom1", Tool("custom1_tool")),
        "custom2": CustomTool("custom2", Tool("custom2_tool")),
        "special-tool": CustomTool("special-tool", Tool("special_tool_info")),
    }
    return tools

# 1. Basic Test Cases

def test_empty_tools_list(tools_instance):
    """Test with an empty input list."""
    codeflash_output = tools_instance._filter_custom_tools([]); result = codeflash_output # 752ns -> 773ns (2.72% slower)

def test_only_standard_tools(tools_instance):
    """Test with a list containing only standard (non-custom) tool names."""
    input_tools = ["std1", "std2", "std3"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 3.31μs -> 1.37μs (142% faster)

def test_only_custom_tools(tools_instance):
    """Test with a list containing only custom tool names."""
    input_tools = ["custom1", "custom2"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 1.60μs -> 1.43μs (12.1% faster)
    # Should return empty standard tools, and both custom tool infos
    custom_infos = [tools_instance._custom_tools["custom1"].info,
                    tools_instance._custom_tools["custom2"].info]

def test_mixed_tools(tools_instance):
    """Test with a mix of standard and custom tool names."""
    input_tools = ["custom1", "std1", "custom2", "std2"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 3.53μs -> 1.64μs (115% faster)
    custom_infos = [tools_instance._custom_tools["custom1"].info,
                    tools_instance._custom_tools["custom2"].info]

def test_duplicate_names(tools_instance):
    """Test with duplicate tool names in the input list."""
    input_tools = ["custom1", "custom1", "std1", "std1", "custom2"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 3.49μs -> 1.77μs (97.4% faster)
    custom_infos = [tools_instance._custom_tools["custom1"].info,
                    tools_instance._custom_tools["custom1"].info,
                    tools_instance._custom_tools["custom2"].info]

# 2. Edge Test Cases

def test_non_existent_tool_names(tools_instance):
    """Test with tool names that are neither standard nor custom (should be treated as standard)."""
    input_tools = ["unknown1", "unknown2"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 2.92μs -> 1.14μs (156% faster)

def test_special_characters_in_names(tools_instance):
    """Test with tool names containing special characters."""
    input_tools = ["special-tool", "weird@tool", "custom2", "std!tool"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 3.54μs -> 1.72μs (105% faster)
    custom_infos = [tools_instance._custom_tools["special-tool"].info,
                    tools_instance._custom_tools["custom2"].info]

def test_case_sensitivity(tools_instance):
    """Test case sensitivity: custom tool names are case-sensitive."""
    input_tools = ["Custom1", "custom1", "CUSTOM2", "custom2"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 3.53μs -> 1.70μs (108% faster)
    # Only exact matches should be considered custom
    custom_infos = [tools_instance._custom_tools["custom1"].info,
                    tools_instance._custom_tools["custom2"].info]

def test_empty_string_and_whitespace_names(tools_instance):
    """Test with empty string and whitespace-only tool names."""
    input_tools = ["", " ", "custom1", "std1"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 3.54μs -> 1.60μs (121% faster)
    custom_infos = [tools_instance._custom_tools["custom1"].info]

def test_long_tool_names(tools_instance):
    """Test with very long tool names."""
    long_name = "x" * 256
    input_tools = [long_name, "custom1"]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 2.67μs -> 1.33μs (101% faster)
    custom_infos = [tools_instance._custom_tools["custom1"].info]

def test_all_custom_tools_are_filtered(tools_instance):
    """Test that all custom tools in the registry are filtered if present in input."""
    input_tools = list(tools_instance._custom_tools.custom_tools_registry.keys())
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 1.71μs -> 1.49μs (14.2% faster)
    custom_infos = [tools_instance._custom_tools[key].info for key in input_tools]

# 3. Large Scale Test Cases

def test_large_number_of_standard_tools(tools_instance):
    """Test with a large list of only standard tools."""
    input_tools = [f"stdtool_{i}" for i in range(500)]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 164μs -> 38.6μs (326% faster)

def test_large_number_of_custom_tools(tools_instance):
    """Test with a large list of only custom tools."""
    # Add more custom tools to the registry
    for i in range(100):
        slug = f"custom_{i}"
        tools_instance._custom_tools.custom_tools_registry[slug] = CustomTool(slug, Tool(f"tool_{i}"))
    input_tools = [f"custom_{i}" for i in range(100)]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 13.7μs -> 11.0μs (24.7% faster)
    expected_custom_infos = [tools_instance._custom_tools[f"custom_{i}"].info for i in range(100)]

def test_large_mixed_tools(tools_instance):
    """Test with a large list mixing standard and custom tools."""
    # Add more custom tools to the registry
    for i in range(200):
        slug = f"custom_{i}"
        tools_instance._custom_tools.custom_tools_registry[slug] = CustomTool(slug, Tool(f"tool_{i}"))
    input_tools = []
    expected_custom_infos = []
    expected_standard = []
    # Mix: every even index is custom, odd is standard
    for i in range(400):
        if i % 2 == 0:
            slug = f"custom_{i//2}"
            input_tools.append(slug)
            expected_custom_infos.append(tools_instance._custom_tools[slug].info)
        else:
            std_name = f"stdtool_{i//2}"
            input_tools.append(std_name)
            expected_standard.append(std_name)
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 92.0μs -> 32.7μs (181% faster)

def test_performance_large_input(tools_instance):
    """Performance: Test with the maximum allowed input size (1000 elements)."""
    # Add 500 custom tools
    for i in range(500):
        slug = f"custom_{i}"
        tools_instance._custom_tools.custom_tools_registry[slug] = CustomTool(slug, Tool(f"tool_{i}"))
    # Mix 500 custom and 500 standard
    input_tools = [f"custom_{i}" if i % 2 == 0 else f"stdtool_{i//2}" for i in range(1000)]
    expected_custom_infos = [tools_instance._custom_tools[f"custom_{i//2}"].info for i in range(0, 1000, 2)]
    expected_standard = [f"stdtool_{i//2}" for i in range(1, 1000, 2)]
    codeflash_output = tools_instance._filter_custom_tools(input_tools); result = codeflash_output # 290μs -> 85.4μs (240% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

import functools
import typing as t

# imports
import pytest  # used for our unit tests
from composio.core.models.tools import Tools


# Minimal stubs for missing classes/types
class Tool:
    def __init__(self, name: str):
        self.name = name

class CustomTool:
    def __init__(self, slug: str, info: Tool):
        self.slug = slug
        self.info = info

class HttpClient:
    pass

class FileHelper:
    def __init__(self, client, outdir=None):
        pass

class Resource:
    pass

TProvider = t.TypeVar("TProvider")
ToolkitVersionParam = t.Any

# Protocol stubs for CustomTools
CustomToolProtocol = t.Callable[..., t.Any]
CustomToolWithProxyProtocol = t.Callable[..., t.Any]

class CustomTools:
    """CustomTools class

    Used to manage custom tools created by users."""

    def __init__(self, client: HttpClient):
        self.client = client
        self.custom_tools_registry: dict[str, CustomTool] = {}

    def __getitem__(self, slug: str) -> CustomTool:
        """Get a custom tool by its slug."""
        return self.custom_tools_registry[slug]

    @t.overload
    def register(self, f: CustomToolProtocol) -> CustomTool: ...

    @t.overload
    def register(
        self, *, toolkit: t.Optional[str] = None
    ) -> t.Callable[[CustomToolWithProxyProtocol], CustomTool]: ...

    def register(
        self,
        f: t.Optional[CustomToolProtocol] = None,
        *,
        toolkit: t.Optional[str] = None,
    ) -> t.Union[
        CustomTool,
        t.Callable[[CustomToolProtocol], CustomTool],
        t.Callable[[CustomToolWithProxyProtocol], CustomTool],
    ]:
        """Register a custom tool."""
        if f is not None:
            return self._wrap_tool(f, toolkit)
        return functools.partial(self._wrap_tool, toolkit=toolkit)
    
    def _wrap_tool(self, f, toolkit):
        # Dummy implementation for testing
        return CustomTool(slug=f.__name__, info=Tool(name=f.__name__))
from composio.core.models.tools import Tools

# ================== UNIT TESTS ==================

@pytest.fixture
def tools_with_customs():
    """
    Fixture to create a Tools instance with some custom tools registered.
    """
    client = HttpClient()
    tools = Tools(client, provider="dummy_provider")
    # Register some custom tools
    ct1 = CustomTool("custom1", Tool("Custom Tool 1"))
    ct2 = CustomTool("custom2", Tool("Custom Tool 2"))
    tools._custom_tools.custom_tools_registry["custom1"] = ct1
    tools._custom_tools.custom_tools_registry["custom2"] = ct2
    return tools

# ------------------- Basic Test Cases -------------------

def test_all_standard_tools(tools_with_customs):
    """
    Test with only standard tools (no custom tools).
    Should return all input as standard, no custom tools.
    """
    input_tools = ["toolA", "toolB", "toolC"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 3.24μs -> 1.34μs (142% faster)

def test_some_custom_tools(tools_with_customs):
    """
    Test with a mix of standard and custom tools.
    Should separate custom and standard tools correctly.
    """
    input_tools = ["toolA", "custom1", "toolB", "custom2"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 3.46μs -> 1.67μs (107% faster)

def test_all_custom_tools(tools_with_customs):
    """
    Test with only custom tools.
    Should return empty standard, all custom tools.
    """
    input_tools = ["custom1", "custom2"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 1.56μs -> 1.37μs (14.2% faster)

def test_empty_input(tools_with_customs):
    """
    Test with empty input list.
    Should return two empty lists.
    """
    std, custom = tools_with_customs._filter_custom_tools([]) # 769ns -> 806ns (4.59% slower)

def test_duplicate_tools(tools_with_customs):
    """
    Test with duplicate tool names in input.
    Should preserve duplicates in output lists.
    """
    input_tools = ["toolA", "custom1", "toolA", "custom2", "custom1"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 3.56μs -> 1.77μs (101% faster)
    # custom1 appears twice, should have two Tool objects with name "Custom Tool 1"
    custom_names = [t.name for t in custom]

# ------------------- Edge Test Cases -------------------

def test_nonexistent_custom_tool(tools_with_customs):
    """
    Test with a tool name that is not in the custom tools registry.
    Should treat as standard tool.
    """
    input_tools = ["custom1", "unknown_custom"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 2.68μs -> 1.37μs (95.3% faster)

def test_tool_name_case_sensitivity(tools_with_customs):
    """
    Test that tool name matching is case-sensitive.
    """
    input_tools = ["custom1", "Custom1", "CUSTOM1"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 3.17μs -> 1.41μs (124% faster)

def test_tool_name_empty_string(tools_with_customs):
    """
    Test with empty string as a tool name.
    Should treat as standard tool if not registered as custom.
    """
    input_tools = [""]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 2.24μs -> 1.03μs (118% faster)



def test_custom_tool_with_empty_slug(tools_with_customs):
    """
    Register a custom tool with an empty string as slug.
    Should be recognized as custom tool.
    """
    ct_empty = CustomTool("", Tool("Empty Slug Tool"))
    tools_with_customs._custom_tools.custom_tools_registry[""] = ct_empty
    input_tools = ["", "toolA"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 2.67μs -> 1.32μs (102% faster)

def test_custom_tool_with_special_characters(tools_with_customs):
    """
    Register custom tools with special characters and test lookup.
    """
    ct_special = CustomTool("tool@#%", Tool("Special Tool"))
    tools_with_customs._custom_tools.custom_tools_registry["tool@#%"] = ct_special
    input_tools = ["tool@#%", "toolA"]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 2.72μs -> 1.28μs (113% faster)

# ------------------- Large Scale Test Cases -------------------

def test_large_number_of_standard_tools(tools_with_customs):
    """
    Test with a large number of standard tools (no customs).
    Should return all as standard, none as custom.
    """
    input_tools = [f"tool{i}" for i in range(500)]
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 163μs -> 36.0μs (354% faster)

def test_large_number_of_custom_tools():
    """
    Test with a large number of custom tools.
    Should return all as custom, none as standard.
    """
    client = HttpClient()
    tools = Tools(client, provider="dummy_provider")
    # Register 500 custom tools
    for i in range(500):
        slug = f"custom{i}"
        tools._custom_tools.custom_tools_registry[slug] = CustomTool(slug, Tool(f"Custom Tool {i}"))
    input_tools = [f"custom{i}" for i in range(500)]
    std, custom = tools._filter_custom_tools(input_tools) # 58.9μs -> 46.8μs (25.9% faster)

def test_large_mixed_tools():
    """
    Test with a large mix of standard and custom tools.
    """
    client = HttpClient()
    tools = Tools(client, provider="dummy_provider")
    # Register 250 custom tools
    for i in range(250):
        slug = f"custom{i}"
        tools._custom_tools.custom_tools_registry[slug] = CustomTool(slug, Tool(f"Custom Tool {i}"))
    # Input: 500 tools, alternate between standard and custom
    input_tools = []
    for i in range(500):
        if i % 2 == 0:
            input_tools.append(f"custom{i//2}")  # custom0, custom1, ...
        else:
            input_tools.append(f"tool{i//2}")    # tool0, tool1, ...
    std, custom = tools._filter_custom_tools(input_tools) # 120μs -> 46.3μs (160% faster)

def test_performance_large_input(tools_with_customs):
    """
    Performance test: Ensure function completes in reasonable time with large input.
    """
    # Register 100 custom tools
    for i in range(100):
        slug = f"custom{i}"
        tools_with_customs._custom_tools.custom_tools_registry[slug] = CustomTool(slug, Tool(f"Custom Tool {i}"))
    # Mix 500 standard and 100 custom tools, shuffled
    input_tools = [f"tool{i}" for i in range(500)] + [f"custom{i}" for i in range(100)]
    import random
    random.shuffle(input_tools)
    std, custom = tools_with_customs._filter_custom_tools(input_tools) # 182μs -> 52.6μs (247% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

More optimizations for composio : https://github.com/aseembits93/composio/pulls, Looking to hear back from the team.

[Equlnox]

Thanks for the improvement @Saga4

Can we instead update the get function of CustomTools itself to do a custom_tools_registry.get(tool) instead ?