Skip to content

datarobot_genai.drmcp.core.drtools_registry

drtools_registry

Registry loader for drtools functions decorated with tool_metadata.

register_drtools_function 

register_drtools_function(func: Callable, metadata: dict[str, Any]) -> None

Register a drtools tool function with the MCP server.

If metadata contains a feature_flag key, the named DR entitlement is evaluated for the static container account at registration time via the shared :func:is_tool_feature_enabled policy. When the flag is disabled — or the check raises for any reason (DR API unavailable, network failure, flag not provisioned) — registration is skipped (fail-closed) so the tool does not appear in list_tools. The gating policy lives in drtools so global-mcp's per-user registry reuses it; only the client differs.

Parameters:

Name Type Description Default
func Callable

The function to register

 required
metadata dict[str, Any]

Tool metadata (tags, name, description, feature_flag, etc.)

 required
Source code in datarobot_genai/drmcp/core/drtools_registry.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def register_drtools_function(func: Callable, metadata: dict[str, Any]) -> None:
    """Register a drtools tool function with the MCP server.

    If ``metadata`` contains a ``feature_flag`` key, the named DR entitlement
    is evaluated for the static container account at registration time via the
    shared :func:`is_tool_feature_enabled` policy. When the flag is disabled —
    or the check raises for any reason (DR API unavailable, network failure,
    flag not provisioned) — registration is skipped (fail-closed) so the tool
    does not appear in ``list_tools``. The gating policy lives in drtools so
    global-mcp's per-user registry reuses it; only the client differs.

    Args:
        func: The function to register
        metadata: Tool metadata (tags, name, description, feature_flag, etc.)
    """
    metadata = dict(metadata)
    feature_flag = metadata.pop("feature_flag", None)
    if not is_tool_feature_enabled(feature_flag, evaluator=_static_account_flag_enabled):
        logger.debug(
            "skipping registration of %s — feature flag %s disabled or unavailable",
            func.__name__,
            feature_flag,
        )
        return

    # Strip private @tool_metadata keys that must not reach the MCP client
    # (UI/gallery metadata and server-side registration hints).
    for key in DRTOOLS_PRIVATE_METADATA_KEYS:
        metadata.pop(key, None)

    # Apply the dr_mcp_tool decorator with the metadata
    dr_mcp_tool(tool_category=DataRobotMCPToolCategory.BUILT_IN_TOOL, **metadata)(func)

    logger.debug(f"Registered drtools tool: {func.__name__}")

load_drtools_registry 

load_drtools_registry(module_name: str) -> None

Load and register all tools from a drtools module.

Parameters:

Name Type Description Default
module_name str

Full module name (e.g., 'datarobot_genai.drtools.panels.tools')

 required
Source code in datarobot_genai/drmcp/core/drtools_registry.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def load_drtools_registry(module_name: str) -> None:
    """Load and register all tools from a drtools module.

    Args:
        module_name: Full module name (e.g., 'datarobot_genai.drtools.panels.tools')
    """
    try:
        # Import the module first to trigger the @tool_metadata decorators
        if module_name.startswith("datarobot_genai.drtools.core."):
            logger.debug(f"Skipping module: {module_name}")
            return

        logger.debug(f"Importing module: {module_name}")
        importlib.import_module(module_name)

        # Register each tool defined in this module
        tool_count = 0
        for func, metadata in get_registered_tools():
            if func.__module__ == module_name:
                register_drtools_function(func, metadata)
                tool_count += 1

        logger.debug(f"Registered {tool_count} tools from {module_name}")

    except ImportError as e:
        logger.debug(f"Could not import module {module_name}: {e}")
    except Exception as e:
        logger.error(f"Error loading drtools registry from {module_name}: {e}")