Skip to content

datarobot_genai.drmcp.core.mcp_instance

mcp_instance

DataRobotMCP

Bases: FastMCP

Extended FastMCP that supports DataRobot specific features like deployments and prompts.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
 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
 79
 80
 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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
class DataRobotMCP(FastMCP):
    """Extended FastMCP that supports DataRobot specific features like deployments and prompts."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._deployments_map: dict[str, str] = {}
        self._prompts_map: dict[str, tuple[str, str]] = {}

    async def get_tools(self) -> dict[str, Tool]:
        """Compat wrapper: fastmcp 3.x renamed get_tools→list_tools and returns a list."""
        return {t.name: t for t in await self.list_tools()}

    async def get_prompts(self) -> dict[str, Prompt]:
        """Compat wrapper: fastmcp 3.x renamed get_prompts→list_prompts and returns a list."""
        return {p.name: p for p in await self.list_prompts()}

    async def get_resources(self) -> dict[str, Any]:
        """Compat wrapper: fastmcp 3.x renamed get_resources→list_resources and returns a list."""
        return {r.name: r for r in await self.list_resources()}

    async def notify_prompts_changed(self) -> None:
        """
        Notify connected clients that the prompt list has changed.

        This method attempts to send a prompts/list_changed notification to inform
        clients that they should refresh their prompt list.

        Note: In stateless HTTP mode (default for this server), notifications may not
        reach clients since each request is independent. This method still logs the
        change for auditing purposes and will work if the server is configured for
        stateful connections.

        See: https://github.com/modelcontextprotocol/python-sdk/issues/710
        """
        logger.info("Prompt list changed - attempting to notify connected clients")

        # Try to use FastMCP's built-in notification mechanism if in an MCP context
        try:
            context = get_context()
            context._queue_prompt_list_changed()
            logger.debug("Queued prompts_changed notification via MCP context")
        except RuntimeError:
            # No active MCP context - this is expected when called from REST API
            logger.debug(
                "No active MCP context for notification. "
                "In stateless mode, clients will see changes on next request."
            )

    async def get_deployment_mapping(self) -> dict[str, str]:
        """
        Get the list of deployment IDs for all registered dynamic tools.

        Returns
        -------
            Dictionary mapping deployment IDs to tool names.
        """
        return self._deployments_map.copy()

    async def set_deployment_mapping(self, deployment_id: str, tool_name: str) -> None:
        """
        Add or update the mapping of a deployment ID to a tool name.

        Args:
            deployment_id: The ID of the deployment.
            tool_name: The name of the tool associated with the deployment.
        """
        existing = self._deployments_map.get(deployment_id)
        if existing and existing != tool_name:
            logger.debug(
                f"Deployment ID {deployment_id} already mapped to {existing}, updating to "
                f"{tool_name}"
            )
            try:
                self.remove_tool(existing)
            except NotFoundError:
                logger.debug(f"Tool {existing} not found in registry, skipping removal")
        self._deployments_map[deployment_id] = tool_name

    async def remove_deployment_mapping(self, deployment_id: str) -> None:
        """
        Remove the mapping of a deployment ID to a tool name.

        Args:
            deployment_id: The ID of the deployment to remove.
        """
        removed = self._deployments_map.pop(deployment_id, None)
        if removed is not None:
            logger.debug(f"Removed deployment mapping for ID {deployment_id} with tool {removed}")
            try:
                self.remove_tool(removed)
            except NotFoundError:
                logger.debug(f"Tool {removed} not found in registry, skipping removal")

    async def get_prompt_mapping(self) -> dict[str, tuple[str, str]]:
        """
        Get the list of prompt ID for all registered dynamic prompts.

        Returns
        -------
            Dictionary mapping prompt template id to prompt template version id and name
        """
        return self._prompts_map.copy()

    async def set_prompt_mapping(
        self, prompt_template_id: str, prompt_template_version_id: str, prompt_name: str
    ) -> None:
        """
        Add or update the mapping of a deployment ID to a tool name.

        Args:
            prompt_template_id: The ID of the prompt template.
            prompt_template_version_id: The ID of the prompt template version.
            prompt_name: The prompt name associated with the prompt template id and version.
        """
        existing_prompt_template = self._prompts_map.get(prompt_template_id)

        if existing_prompt_template:
            existing_prompt_template_version_id, _ = existing_prompt_template

            logger.debug(
                f"Prompt template ID {prompt_template_id} "
                f"already mapped to {existing_prompt_template_version_id}. "
                f"Updating to version id = {prompt_template_version_id} and name = {prompt_name}"
            )
            await self.remove_prompt_mapping(
                prompt_template_id, existing_prompt_template_version_id
            )

        self._prompts_map[prompt_template_id] = (prompt_template_version_id, prompt_name)

    async def remove_prompt_mapping(
        self, prompt_template_id: str, prompt_template_version_id: str
    ) -> None:
        """
        Remove the mapping of a prompt_template ID to a version and prompt name.

        Args:
            prompt_template_id: The ID of the prompt template to remove.
            prompt_template_version_id: The ID of the prompt template version to remove.
        """
        if existing_prompt_template := self._prompts_map.get(prompt_template_id):
            existing_prompt_template_version_id, _ = existing_prompt_template
            if existing_prompt_template_version_id != prompt_template_version_id:
                logger.debug(
                    f"Found prompt template with id = {prompt_template_id} in registry, "
                    f"but with different version = {existing_prompt_template_version_id}, "
                    f"skipping removal."
                )
            else:
                prompts_d = await self.get_prompts()
                for prompt in prompts_d.values():
                    if (
                        prompt.meta is not None
                        and prompt.meta.get("prompt_template_id", "") == prompt_template_id
                        and prompt.meta.get("prompt_template_version_id", "")
                        == prompt_template_version_id
                    ):
                        try:
                            # remove_prompt_mapping now removes the prompt via the FastMCP 3.x API
                            self.local_provider.remove_prompt(prompt.name)
                        except KeyError:
                            logger.debug(
                                f"Prompt {prompt.name!r} not found in local provider, "
                                "skipping removal."
                            )

                self._prompts_map.pop(prompt_template_id, None)

                # Notify clients that the prompt list has changed
                await self.notify_prompts_changed()
        else:
            logger.debug(
                f"Do not found prompt template with id = {prompt_template_id} in registry, "
                f"skipping removal."
            )

get_tools async 

get_tools() -> dict[str, Tool]

Compat wrapper: fastmcp 3.x renamed get_tools→list_tools and returns a list.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
60
61
62
async def get_tools(self) -> dict[str, Tool]:
    """Compat wrapper: fastmcp 3.x renamed get_tools→list_tools and returns a list."""
    return {t.name: t for t in await self.list_tools()}

get_prompts async 

get_prompts() -> dict[str, Prompt]

Compat wrapper: fastmcp 3.x renamed get_prompts→list_prompts and returns a list.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
64
65
66
async def get_prompts(self) -> dict[str, Prompt]:
    """Compat wrapper: fastmcp 3.x renamed get_prompts→list_prompts and returns a list."""
    return {p.name: p for p in await self.list_prompts()}

get_resources async 

get_resources() -> dict[str, Any]

Compat wrapper: fastmcp 3.x renamed get_resources→list_resources and returns a list.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
68
69
70
async def get_resources(self) -> dict[str, Any]:
    """Compat wrapper: fastmcp 3.x renamed get_resources→list_resources and returns a list."""
    return {r.name: r for r in await self.list_resources()}

notify_prompts_changed async 

notify_prompts_changed() -> None

Notify connected clients that the prompt list has changed.

This method attempts to send a prompts/list_changed notification to inform clients that they should refresh their prompt list.

Note: In stateless HTTP mode (default for this server), notifications may not reach clients since each request is independent. This method still logs the change for auditing purposes and will work if the server is configured for stateful connections.

See: https://github.com/modelcontextprotocol/python-sdk/issues/710

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
async def notify_prompts_changed(self) -> None:
    """
    Notify connected clients that the prompt list has changed.

    This method attempts to send a prompts/list_changed notification to inform
    clients that they should refresh their prompt list.

    Note: In stateless HTTP mode (default for this server), notifications may not
    reach clients since each request is independent. This method still logs the
    change for auditing purposes and will work if the server is configured for
    stateful connections.

    See: https://github.com/modelcontextprotocol/python-sdk/issues/710
    """
    logger.info("Prompt list changed - attempting to notify connected clients")

    # Try to use FastMCP's built-in notification mechanism if in an MCP context
    try:
        context = get_context()
        context._queue_prompt_list_changed()
        logger.debug("Queued prompts_changed notification via MCP context")
    except RuntimeError:
        # No active MCP context - this is expected when called from REST API
        logger.debug(
            "No active MCP context for notification. "
            "In stateless mode, clients will see changes on next request."
        )

get_deployment_mapping async 

get_deployment_mapping() -> dict[str, str]

Get the list of deployment IDs for all registered dynamic tools.

Returns
Dictionary mapping deployment IDs to tool names.
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
100
101
102
103
104
105
106
107
108
async def get_deployment_mapping(self) -> dict[str, str]:
    """
    Get the list of deployment IDs for all registered dynamic tools.

    Returns
    -------
        Dictionary mapping deployment IDs to tool names.
    """
    return self._deployments_map.copy()

set_deployment_mapping async 

set_deployment_mapping(deployment_id: str, tool_name: str) -> None

Add or update the mapping of a deployment ID to a tool name.

Parameters:

Name Type Description Default
deployment_id str

The ID of the deployment.

 required
tool_name str

The name of the tool associated with the deployment.

 required
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
async def set_deployment_mapping(self, deployment_id: str, tool_name: str) -> None:
    """
    Add or update the mapping of a deployment ID to a tool name.

    Args:
        deployment_id: The ID of the deployment.
        tool_name: The name of the tool associated with the deployment.
    """
    existing = self._deployments_map.get(deployment_id)
    if existing and existing != tool_name:
        logger.debug(
            f"Deployment ID {deployment_id} already mapped to {existing}, updating to "
            f"{tool_name}"
        )
        try:
            self.remove_tool(existing)
        except NotFoundError:
            logger.debug(f"Tool {existing} not found in registry, skipping removal")
    self._deployments_map[deployment_id] = tool_name

remove_deployment_mapping async 

remove_deployment_mapping(deployment_id: str) -> None

Remove the mapping of a deployment ID to a tool name.

Parameters:

Name Type Description Default
deployment_id str

The ID of the deployment to remove.

 required
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
async def remove_deployment_mapping(self, deployment_id: str) -> None:
    """
    Remove the mapping of a deployment ID to a tool name.

    Args:
        deployment_id: The ID of the deployment to remove.
    """
    removed = self._deployments_map.pop(deployment_id, None)
    if removed is not None:
        logger.debug(f"Removed deployment mapping for ID {deployment_id} with tool {removed}")
        try:
            self.remove_tool(removed)
        except NotFoundError:
            logger.debug(f"Tool {removed} not found in registry, skipping removal")

get_prompt_mapping async 

get_prompt_mapping() -> dict[str, tuple[str, str]]

Get the list of prompt ID for all registered dynamic prompts.

Returns
Dictionary mapping prompt template id to prompt template version id and name
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
145
146
147
148
149
150
151
152
153
async def get_prompt_mapping(self) -> dict[str, tuple[str, str]]:
    """
    Get the list of prompt ID for all registered dynamic prompts.

    Returns
    -------
        Dictionary mapping prompt template id to prompt template version id and name
    """
    return self._prompts_map.copy()

set_prompt_mapping async 

set_prompt_mapping(prompt_template_id: str, prompt_template_version_id: str, prompt_name: str) -> None

Add or update the mapping of a deployment ID to a tool name.

Parameters:

Name Type Description Default
prompt_template_id str

The ID of the prompt template.

 required
prompt_template_version_id str

The ID of the prompt template version.

 required
prompt_name str

The prompt name associated with the prompt template id and version.

 required
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
async def set_prompt_mapping(
    self, prompt_template_id: str, prompt_template_version_id: str, prompt_name: str
) -> None:
    """
    Add or update the mapping of a deployment ID to a tool name.

    Args:
        prompt_template_id: The ID of the prompt template.
        prompt_template_version_id: The ID of the prompt template version.
        prompt_name: The prompt name associated with the prompt template id and version.
    """
    existing_prompt_template = self._prompts_map.get(prompt_template_id)

    if existing_prompt_template:
        existing_prompt_template_version_id, _ = existing_prompt_template

        logger.debug(
            f"Prompt template ID {prompt_template_id} "
            f"already mapped to {existing_prompt_template_version_id}. "
            f"Updating to version id = {prompt_template_version_id} and name = {prompt_name}"
        )
        await self.remove_prompt_mapping(
            prompt_template_id, existing_prompt_template_version_id
        )

    self._prompts_map[prompt_template_id] = (prompt_template_version_id, prompt_name)

remove_prompt_mapping async 

remove_prompt_mapping(prompt_template_id: str, prompt_template_version_id: str) -> None

Remove the mapping of a prompt_template ID to a version and prompt name.

Parameters:

Name Type Description Default
prompt_template_id str

The ID of the prompt template to remove.

 required
prompt_template_version_id str

The ID of the prompt template version to remove.

 required
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
async def remove_prompt_mapping(
    self, prompt_template_id: str, prompt_template_version_id: str
) -> None:
    """
    Remove the mapping of a prompt_template ID to a version and prompt name.

    Args:
        prompt_template_id: The ID of the prompt template to remove.
        prompt_template_version_id: The ID of the prompt template version to remove.
    """
    if existing_prompt_template := self._prompts_map.get(prompt_template_id):
        existing_prompt_template_version_id, _ = existing_prompt_template
        if existing_prompt_template_version_id != prompt_template_version_id:
            logger.debug(
                f"Found prompt template with id = {prompt_template_id} in registry, "
                f"but with different version = {existing_prompt_template_version_id}, "
                f"skipping removal."
            )
        else:
            prompts_d = await self.get_prompts()
            for prompt in prompts_d.values():
                if (
                    prompt.meta is not None
                    and prompt.meta.get("prompt_template_id", "") == prompt_template_id
                    and prompt.meta.get("prompt_template_version_id", "")
                    == prompt_template_version_id
                ):
                    try:
                        # remove_prompt_mapping now removes the prompt via the FastMCP 3.x API
                        self.local_provider.remove_prompt(prompt.name)
                    except KeyError:
                        logger.debug(
                            f"Prompt {prompt.name!r} not found in local provider, "
                            "skipping removal."
                        )

            self._prompts_map.pop(prompt_template_id, None)

            # Notify clients that the prompt list has changed
            await self.notify_prompts_changed()
    else:
        logger.debug(
            f"Do not found prompt template with id = {prompt_template_id} in registry, "
            f"skipping removal."
        )

ToolKwargs

Bases: TypedDict

Keyword arguments passed through to FastMCP's mcp.tool() decorator.

All parameters are optional and forwarded directly to FastMCP tool registration. See FastMCP documentation for full details on each parameter.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
class ToolKwargs(TypedDict, total=False):
    """Keyword arguments passed through to FastMCP's mcp.tool() decorator.

    All parameters are optional and forwarded directly to FastMCP tool registration.
    See FastMCP documentation for full details on each parameter.
    """

    name: str | None
    title: str | None
    description: str | None
    icons: list[Any] | None
    tags: set[str] | None
    output_schema: dict[str, Any] | None
    annotations: Any | None
    exclude_args: list[str] | None
    meta: dict[str, Any] | None

dr_core_mcp_tool 

dr_core_mcp_tool(**kwargs: Unpack[ToolKwargs]) -> Callable[[Callable[..., Any]], Callable[..., Any]]

Combine decorator that includes mcp.tool() and dr_mcp_extras().

All keyword arguments are passed through to FastMCP's mcp.tool() decorator. See ToolKwargs for available parameters.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
def dr_core_mcp_tool(
    **kwargs: Unpack[ToolKwargs],
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Combine decorator that includes mcp.tool() and dr_mcp_extras().

    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
    See ToolKwargs for available parameters.
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        instrumented = dr_mcp_extras()(func)
        mcp.tool(**kwargs)(instrumented)
        return instrumented

    return decorator

dr_mcp_tool 

dr_mcp_tool(tool_category: DataRobotMCPToolCategory = DataRobotMCPToolCategory.USER_TOOL, **mcp_tool_init_args: Unpack[ToolKwargs]) -> Callable[[Callable[..., Any]], Callable[..., Any]]

Combine decorator that includes mcp.tool() and dr_mcp_extras().

All keyword arguments are passed through to FastMCP's mcp.tool() decorator. See ToolKwargs for available parameters.

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
def dr_mcp_tool(
    tool_category: DataRobotMCPToolCategory = DataRobotMCPToolCategory.USER_TOOL,
    **mcp_tool_init_args: Unpack[ToolKwargs],
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Combine decorator that includes mcp.tool() and dr_mcp_extras().

    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
    See ToolKwargs for available parameters.
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        updated_kwargs = update_mcp_tool_init_args_with_tool_category(
            tool_category, **mcp_tool_init_args
        )
        # fastmcp 3.x removed 'enabled' from tool(); handle it separately
        enabled = updated_kwargs.pop("enabled", None)  # type: ignore[typeddict-item]
        # Apply the MCP decorators
        instrumented = dr_mcp_extras()(func)
        mcp.tool(**updated_kwargs)(instrumented)
        if enabled is False:
            tool_name = updated_kwargs.get("name") or func.__name__
            mcp.disable(names={tool_name}, components={"tool"})
        return instrumented

    return decorator

dr_mcp_integration_tool 

dr_mcp_integration_tool(**mcp_tool_init_args: Unpack[ToolKwargs]) -> Callable[[Callable[..., Any]], Callable[..., Any]]

Decorate mcp tool created as a wrapper of external service API (e.g., DataRobot Predictive AI, github API).

Source code in datarobot_genai/drmcp/core/mcp_instance.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
def dr_mcp_integration_tool(
    **mcp_tool_init_args: Unpack[ToolKwargs],
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Decorate mcp tool created as a wrapper of external service API (e.g., DataRobot Predictive
    AI, github API).
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        return dr_mcp_tool(
            tool_category=DataRobotMCPToolCategory.BUILT_IN_TOOL,
            **mcp_tool_init_args,
        )(func)

    return decorator

dr_mcp_extras 

dr_mcp_extras(type: str = 'tool') -> Callable[[Callable[..., Any]], Callable[..., Any]]

Combine decorator that includes log_execution and trace_execution().

Parameters:

Name Type Description Default
type str

default is "tool", other options are "prompt", "resource"

 'tool'
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
405
406
407
408
409
410
411
412
413
414
415
416
417
def dr_mcp_extras(
    type: str = "tool",
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Combine decorator that includes log_execution and trace_execution().

    Args:
        type: default is "tool", other options are "prompt", "resource"
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        return log_execution(trace_execution(trace_type=type)(func))

    return decorator

register_tools async 

register_tools(fn: AnyFunction, name: str | None = None, title: str | None = None, description: str | None = None, tags: set[str] | None = None, deployment_id: str | None = None, tool_category: DataRobotMCPToolCategory = DataRobotMCPToolCategory.USER_TOOL_DEPLOYMENT) -> Tool

Register new tools after server has started.

Parameters:

Name Type Description Default
fn AnyFunction

The function to register as a tool

 required
name str | None

Optional name for the tool (defaults to function name)

 None
title str | None

Optional human-readable title for the tool

 None
description str | None

Optional description of what the tool does

 None
tags set[str] | None

Optional set of tags to apply to the tool

 None
deployment_id str | None

Optional deployment ID associated with the tool

 None
tool_category DataRobotMCPToolCategory

Category of the tool. Its value is from DataRobotMCPToolCategory

 USER_TOOL_DEPLOYMENT
Returns
The registered Tool object
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
async def register_tools(
    fn: AnyFunction,
    name: str | None = None,
    title: str | None = None,
    description: str | None = None,
    tags: set[str] | None = None,
    deployment_id: str | None = None,
    tool_category: DataRobotMCPToolCategory = DataRobotMCPToolCategory.USER_TOOL_DEPLOYMENT,
) -> Tool:
    """
    Register new tools after server has started.

    Args:
        fn: The function to register as a tool
        name: Optional name for the tool (defaults to function name)
        title: Optional human-readable title for the tool
        description: Optional description of what the tool does
        tags: Optional set of tags to apply to the tool
        deployment_id: Optional deployment ID associated with the tool
        tool_category: Category of the tool. Its value is from DataRobotMCPToolCategory

    Returns
    -------
        The registered Tool object
    """
    tool_name = name or fn.__name__
    logger.info(f"Registering new tool: {tool_name}")

    wrapped_fn = dr_mcp_extras()(fn)

    # Create annotations only when additional metadata is required
    annotations: ToolAnnotations | None = None  # type: ignore[assignment]
    if deployment_id is not None:
        annotations = ToolAnnotations()  # type: ignore[call-arg]
        annotations.deployment_id = deployment_id  # type: ignore[attr-defined]

    tool = Tool.from_function(
        fn=wrapped_fn,
        name=tool_name,
        title=title,
        description=description,
        annotations=annotations,
        tags=tags,
        meta={"tool_category": tool_category.name},
    )

    # Register the tool
    registered_tool = mcp.add_tool(tool)

    # Map deployment ID to tool name if provided
    if deployment_id:
        await mcp.set_deployment_mapping(deployment_id, tool_name)

    await check_tool_registration_status_after_it_finishes(mcp, tool_name)

    return registered_tool

register_prompt async 

register_prompt(fn: AnyFunction, name: str | None = None, title: str | None = None, description: str | None = None, tags: set[str] | None = None, meta: dict[str, Any] | None = None, prompt_template: tuple[str, str] | None = None, prompt_category: DataRobotMCPPromptCategory = DataRobotMCPPromptCategory.USER_PROMPT_TEMPLATE_VERSION) -> Prompt

Register new prompt after server has started.

Parameters:

Name Type Description Default
fn AnyFunction

The function to register as a prompt

 required
name str | None

Optional name for the prompt (defaults to function name)

 None
title str | None

Optional human-readable title for the prompt

 None
description str | None

Optional description of what the prompt does

 None
tags set[str] | None

Optional set of tags to apply to the prompt

 None
meta dict[str, Any] | None

Optional dict of metadata to apply to the prompt

 None
prompt_template tuple[str, str] | None

Optional (id, version id) of the prompt template

 None
prompt_category DataRobotMCPPromptCategory

Category of prompt. Its value is from DataRobotMCPPromptCategory

 USER_PROMPT_TEMPLATE_VERSION
Returns
The registered Prompt object
Source code in datarobot_genai/drmcp/core/mcp_instance.py 
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
async def register_prompt(
    fn: AnyFunction,
    name: str | None = None,
    title: str | None = None,
    description: str | None = None,
    tags: set[str] | None = None,
    meta: dict[str, Any] | None = None,
    prompt_template: tuple[str, str] | None = None,
    prompt_category: DataRobotMCPPromptCategory = DataRobotMCPPromptCategory.USER_PROMPT_TEMPLATE_VERSION,  # noqa: E501
) -> Prompt:
    """
    Register new prompt after server has started.

    Args:
        fn: The function to register as a prompt
        name: Optional name for the prompt (defaults to function name)
        title: Optional human-readable title for the prompt
        description: Optional description of what the prompt does
        tags: Optional set of tags to apply to the prompt
        meta: Optional dict of metadata to apply to the prompt
        prompt_template: Optional (id, version id) of the prompt template
        prompt_category: Category of prompt. Its value is from DataRobotMCPPromptCategory

    Returns
    -------
        The registered Prompt object
    """
    prompt_name = name or fn.__name__
    logger.info(f"Registering new prompt: {prompt_name}")
    wrapped_fn = dr_mcp_extras(type="prompt")(fn)

    prompt_name_no_duplicate = await get_prompt_name_no_duplicate(mcp, prompt_name)

    meta = meta or {}
    meta["resource_category"] = prompt_category.name
    prompt = Prompt.from_function(
        fn=wrapped_fn,
        name=prompt_name_no_duplicate,
        title=title,
        description=description,
        tags=tags,
        meta=meta,
    )

    # Register the prompt
    if prompt_template:
        prompt_template_id, prompt_template_version_id = prompt_template
        await mcp.set_prompt_mapping(
            prompt_template_id, prompt_template_version_id, prompt_name_no_duplicate
        )

    registered_prompt = mcp.add_prompt(prompt)

    await check_prompt_registration_status_after_it_finishes(mcp, prompt_name_no_duplicate)

    # Notify clients that the prompt list has changed
    await mcp.notify_prompts_changed()

    return registered_prompt