1  """
 2  Decorator for exposing MLflow CLI commands as MCP tools.
 3  
 4  Usage:
 5      from mlflow.mcp.decorator import mlflow_mcp
 6  
 7      @commands.command("search")
 8      @mlflow_mcp(tool_name="search_traces")
 9      @click.option(...)
10      def search_traces(...):
11          ...
12  
13  The decorator attaches MCP metadata to the Click command, which is then
14  used by the MCP server to register the tool with the specified name.
15  """
16  
17  from typing import Callable, TypeVar
18  
19  import click
20  
21  # Attribute name used to store MCP metadata on Click commands
22  MCP_METADATA_ATTR = "_mlflow_mcp_metadata"
23  
24  F = TypeVar("F", bound=Callable)
25  
26  
27  def mlflow_mcp(tool_name: str) -> Callable[[F], F]:
28      """
29      Decorator to expose a Click command as an MCP tool with a curated name.
30  
31      Args:
32          tool_name: The name to use for the MCP tool. This should be a clear,
33              agent-friendly name that describes what the tool does.
34              Convention: action_entity (e.g., "search_traces", "get_experiment")
35  
36      Example:
37          @commands.command("search")
38          @mlflow_mcp(tool_name="search_traces")
39          def search(...):
40              '''Search for traces in the specified experiment.'''
41              ...
42  
43      The decorator stores metadata on the function that the MCP server reads
44      when registering tools. Commands without this decorator are not exposed
45      as MCP tools.
46      """
47  
48      def decorator(fn: F) -> F:
49          # Store MCP metadata on the function
50          setattr(fn, MCP_METADATA_ATTR, {"tool_name": tool_name})
51          return fn
52  
53      return decorator
54  
55  
56  def get_mcp_tool_name(cmd: click.Command) -> str | None:
57      """
58      Get the MCP tool name from a Click command, if it has been decorated.
59  
60      Args:
61          cmd: The Click command to check.
62  
63      Returns:
64          The MCP tool name if the command has been decorated with @mlflow_mcp,
65          None otherwise.
66      """
67      if cmd.callback is None:
68          return None
69  
70      metadata = getattr(cmd.callback, MCP_METADATA_ATTR, None)
71      if metadata is None:
72          return None
73  
74      return metadata.get("tool_name")