mlflow--mlflow
wehub-resource-sync
75 行
2.0 KiB
Python
75 行
2.0 KiB
Python
"""
|
|
Decorator for exposing MLflow CLI commands as MCP tools.
|
|
|
|
Usage:
|
|
from mlflow.mcp.decorator import mlflow_mcp
|
|
|
|
@commands.command("search")
|
|
@mlflow_mcp(tool_name="search_traces")
|
|
@click.option(...)
|
|
def search_traces(...):
|
|
...
|
|
|
|
The decorator attaches MCP metadata to the Click command, which is then
|
|
used by the MCP server to register the tool with the specified name.
|
|
"""
|
|
|
|
from typing import Callable, TypeVar
|
|
|
|
import click
|
|
|
|
# Attribute name used to store MCP metadata on Click commands
|
|
MCP_METADATA_ATTR = "_mlflow_mcp_metadata"
|
|
|
|
F = TypeVar("F", bound=Callable)
|
|
|
|
|
|
def mlflow_mcp(tool_name: str) -> Callable[[F], F]:
|
|
"""
|
|
Decorator to expose a Click command as an MCP tool with a curated name.
|
|
|
|
Args:
|
|
tool_name: The name to use for the MCP tool. This should be a clear,
|
|
agent-friendly name that describes what the tool does.
|
|
Convention: action_entity (e.g., "search_traces", "get_experiment")
|
|
|
|
Example:
|
|
@commands.command("search")
|
|
@mlflow_mcp(tool_name="search_traces")
|
|
def search(...):
|
|
'''Search for traces in the specified experiment.'''
|
|
...
|
|
|
|
The decorator stores metadata on the function that the MCP server reads
|
|
when registering tools. Commands without this decorator are not exposed
|
|
as MCP tools.
|
|
"""
|
|
|
|
def decorator(fn: F) -> F:
|
|
# Store MCP metadata on the function
|
|
setattr(fn, MCP_METADATA_ATTR, {"tool_name": tool_name})
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def get_mcp_tool_name(cmd: click.Command) -> str | None:
|
|
"""
|
|
Get the MCP tool name from a Click command, if it has been decorated.
|
|
|
|
Args:
|
|
cmd: The Click command to check.
|
|
|
|
Returns:
|
|
The MCP tool name if the command has been decorated with @mlflow_mcp,
|
|
None otherwise.
|
|
"""
|
|
if cmd.callback is None:
|
|
return None
|
|
|
|
metadata = getattr(cmd.callback, MCP_METADATA_ATTR, None)
|
|
if metadata is None:
|
|
return None
|
|
|
|
return metadata.get("tool_name")
|