update

app.py

@@ -0,0 +1,125 @@
+import json
+import gradio as gr
+from mcp_mlflow_tools import (
+    set_tracking_uri,
+    get_system_info,
+    list_experiments,
+    create_experiment,
+    register_model,
+    search_runs,
+    list_registered_models,
+    get_model_info
+)
+
+def create_interface():
+    with gr.Blocks(title="MLflow MCP Service") as app:
+        gr.Markdown("# MLflow MCP Service")
+        gr.Markdown("A service that exposes MLflow functionality through a web interface and API endpoints.")
+        
+        with gr.Tab("Tracking & System Info"):
+            with gr.Group():
+                gr.Markdown("## Set Tracking URI")
+                uri_input = gr.Textbox(label="MLflow Tracking URI")
+                uri_output = gr.JSON(label="Result")
+                uri_button = gr.Button("Set URI")
+                uri_button.click(
+                    fn=set_tracking_uri,
+                    inputs=uri_input,
+                    outputs=uri_output
+                )
+            
+            with gr.Group():
+                gr.Markdown("## Get System Info")
+                sys_info_output = gr.JSON(label="System Information")
+                sys_info_button = gr.Button("Get Info")
+                sys_info_button.click(
+                    fn=get_system_info,
+                    inputs=[],
+                    outputs=sys_info_output
+                )
+        
+        with gr.Tab("Experiment Management"):
+            with gr.Group():
+                gr.Markdown("## List Experiments")
+                exp_list_output = gr.JSON(label="Experiments")
+                exp_list_button = gr.Button("List Experiments")
+                exp_list_button.click(
+                    fn=list_experiments,
+                    inputs=[],
+                    outputs=exp_list_output
+                )
+            
+            with gr.Group():
+                gr.Markdown("## Create Experiment")
+                exp_name_input = gr.Textbox(label="Experiment Name")
+                exp_tags_input = gr.Textbox(label="Tags (JSON format)", placeholder='{"key": "value"}')
+                exp_create_output = gr.JSON(label="Result")
+                exp_create_button = gr.Button("Create Experiment")
+                
+                def create_exp_with_tags(name, tags_str):
+                    try:
+                        tags = json.loads(tags_str) if tags_str else None
+                    except json.JSONDecodeError:
+                        return {"error": True, "message": "Invalid JSON format for tags"}
+                    return create_experiment(name, tags)
+                
+                exp_create_button.click(
+                    fn=create_exp_with_tags,
+                    inputs=[exp_name_input, exp_tags_input],
+                    outputs=exp_create_output
+                )
+        
+        with gr.Tab("Model Registry"):
+            with gr.Group():
+                gr.Markdown("## Register Model")
+                reg_run_id = gr.Textbox(label="Run ID")
+                reg_artifact_path = gr.Textbox(label="Artifact Path")
+                reg_model_name = gr.Textbox(label="Model Name")
+                reg_output = gr.JSON(label="Result")
+                reg_button = gr.Button("Register Model")
+                reg_button.click(
+                    fn=register_model,
+                    inputs=[reg_run_id, reg_artifact_path, reg_model_name],
+                    outputs=reg_output
+                )
+            
+            with gr.Group():
+                gr.Markdown("## List Registered Models")
+                list_models_output = gr.JSON(label="Models")
+                list_models_button = gr.Button("List Models")
+                list_models_button.click(
+                    fn=list_registered_models,
+                    inputs=[],
+                    outputs=list_models_output
+                )
+            
+            with gr.Group():
+                gr.Markdown("## Get Model Info")
+                model_info_name = gr.Textbox(label="Model Name")
+                model_info_output = gr.JSON(label="Model Information")
+                model_info_button = gr.Button("Get Info")
+                model_info_button.click(
+                    fn=get_model_info,
+                    inputs=model_info_name,
+                    outputs=model_info_output
+                )
+        
+        with gr.Tab("Run Search"):
+            with gr.Group():
+                gr.Markdown("## Search Runs")
+                search_exp_id = gr.Textbox(label="Experiment ID")
+                search_filter = gr.Textbox(label="Filter String")
+                search_max_results = gr.Number(label="Max Results", value=100, precision=0)
+                search_output = gr.JSON(label="Search Results")
+                search_button = gr.Button("Search")
+                search_button.click(
+                    fn=search_runs,
+                    inputs=[search_exp_id, search_filter, search_max_results],
+                    outputs=search_output
+                )
+    
+    return app
+
+if __name__ == "__main__":
+    app = create_interface()
+    app.launch(mcp_server=True) 

mcp_mlflow_tools.py

@@ -0,0 +1,380 @@
+import logging
+import mlflow
+from datetime import datetime
+from typing import Dict, List, Optional, Literal
+from mlflow.tracking import MlflowClient
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def _format_timestamp(ts: int) -> str:
+    """Convert MLflow timestamp (milliseconds since epoch) to readable string."""
+    dt = datetime.fromtimestamp(ts / 1000.0)
+    return dt.strftime("%Y-%m-%d %H:%M:%S")
+
+def set_tracking_uri(uri: str) -> Dict:
+    """Set MLflow tracking URI and verify connection."""
+    if not uri:
+        return {"error": True, "message": "URI cannot be empty"}
+    
+    try:
+        logger.info(f"Setting MLflow tracking URI to {uri}")
+        mlflow.set_tracking_uri(uri)
+        return get_system_info()
+    except Exception as e:
+        return {"error": True, "message": f"Failed to set URI: {str(e)}"}
+
+def get_system_info() -> Dict:
+    """Get MLflow system information."""
+    try:
+        client = MlflowClient()
+        return {
+            "mlflow_version": mlflow.__version__,
+            "tracking_uri": mlflow.get_tracking_uri(),
+            "registry_uri": mlflow.get_registry_uri(),
+            "artifact_uri": mlflow.get_artifact_uri(),
+            "python_version": mlflow.__version__,
+            "server_time": _format_timestamp(int(datetime.now().timestamp() * 1000)),
+            "experiment_count": len(mlflow.search_experiments()),
+            "model_count": len(client.search_registered_models())
+        }
+    except Exception as e:
+        return {"error": True, "message": f"Failed to fetch system info: {str(e)}"}
+
+def list_experiments(name_contains: Optional[str] = "", max_results: Optional[int] = 100) -> Dict:
+    """
+    List all experiments in the MLflow tracking server, with optional filtering.
+    Includes run count for each experiment.
+    
+    Args:
+        name_contains: Optional filter to only include experiments whose names contain this string (case-insensitive).
+        max_results: Maximum number of results to return (default: 100). None means no limit after filtering.
+                     A negative value will result in an empty list.
+    
+    Returns:
+        A dictionary containing the total count of returned experiments and a list of their details.
+        Format: {"total_experiments": count, "experiments": [exp_details, ...]}
+        Returns {"error": True, "message": ...} on failure.
+    """
+    logger.info(f"Fetching experiments (filter: '{name_contains}', max_results: {max_results})")
+    
+    try:
+        client = MlflowClient()
+        all_mlflow_experiments: List[mlflow.entities.Experiment] = client.search_experiments()
+        filtered_experiments: List[mlflow.entities.Experiment]
+        processed_name_filter = name_contains.strip().lower() if name_contains else ""
+        
+        if processed_name_filter:
+            filtered_experiments = [
+                exp for exp in all_mlflow_experiments 
+                if processed_name_filter in exp.name.lower()
+            ]
+        else:
+            filtered_experiments = all_mlflow_experiments
+        
+        # Apply max_results limit
+        limited_experiments: List[mlflow.entities.Experiment]
+        if max_results is not None:
+            if max_results < 0:
+                limited_experiments = []
+            else:
+                limited_experiments = filtered_experiments[:max_results]
+        else: # max_results is None, return all filtered experiments
+            limited_experiments = filtered_experiments
+            
+        experiments_info = []
+        
+        for exp in limited_experiments:
+            creation_time_str = None
+            if hasattr(exp, "creation_time") and exp.creation_time is not None:
+                creation_time_str = _format_timestamp(exp.creation_time)
+            
+            tags_dict = {}
+            if hasattr(exp, "tags") and exp.tags:
+                tags_dict = dict(exp.tags) # exp.tags is already a dict {key: value}
+
+            exp_detail = {
+                "experiment_id": exp.experiment_id,
+                "name": exp.name,
+                "artifact_location": exp.artifact_location,
+                "lifecycle_stage": exp.lifecycle_stage,
+                "creation_time": creation_time_str,
+                "tags": tags_dict
+            }
+            
+            run_count_val: any # Can be int or str
+            try:
+                # Check if any runs exist for this experiment (counts active and deleted)
+                probe_runs = client.search_runs(
+                    experiment_ids=[exp.experiment_id], 
+                    max_results=1, 
+                    run_view_type=mlflow.entities.ViewType.ALL
+                )
+                if probe_runs:
+                    # If runs exist, get a more accurate count up to a practical limit
+                    all_runs_for_count = client.search_runs(
+                        experiment_ids=[exp.experiment_id], 
+                        max_results=50000, # Practical limit for counting
+                        run_view_type=mlflow.entities.ViewType.ALL
+                    )
+                    run_count_val = len(all_runs_for_count)
+                else:
+                    run_count_val = 0
+            except Exception as e_runs:
+                logger.warning(f"Error getting run count for experiment '{exp.name}' (ID: {exp.experiment_id}): {str(e_runs)}")
+                run_count_val = "Error getting count"
+            
+            exp_detail["run_count"] = run_count_val
+            experiments_info.append(exp_detail)
+        
+        result = {
+            "total_experiments": len(experiments_info),
+            "experiments": experiments_info
+        }
+        
+        return result
+    
+    except Exception as e:
+        error_msg = f"Error listing experiments: {str(e)}"
+        logger.error(error_msg, exc_info=True)
+        return {"error": True, "message": error_msg}
+
+def create_experiment(name: str, tags: Optional[Dict[str, str]] = None) -> Dict:
+    """Create a new experiment."""
+    if not name:
+        return {"error": True, "message": "Experiment name cannot be empty"}
+    
+    try:
+        experiment_id = mlflow.create_experiment(name=name, tags=tags or {})
+        return {
+            "experiment_id": experiment_id,
+            "message": "Created experiment"
+        }
+    except Exception as e:
+        return {"error": True, "message": f"Failed to create experiment: {str(e)}"}
+
+def search_runs(
+    experiment_id: str,
+    filter_string: str,
+    order_string: Optional[str] = None,
+    max_results: int = 100
+) -> Dict:
+    """
+    Search runs in a given experiment, with filtering and ordering.
+
+    Args:
+        experiment_id: The ID of the experiment to search runs in.
+        filter_string: A filter query string used to search for runs. 
+                       It follows the MLflow search filter syntax.
+                       Examples:
+                         - "metrics.accuracy > 0.95"
+                         - "params.learning_rate = '0.001'"
+                         - "tags.environment = 'production'"
+                         - "attributes.status = 'FINISHED'"
+                         - "metrics.loss < 0.2 AND params.optimizer = 'Adam'"
+                       If an empty string is provided, no filtering is applied by this string.
+                       Multiple conditions can be combined using 'AND' or 'OR'.
+        order_string: An optional string to define the order of the results.
+                      It should be a single string composed of a metric, parameter, or attribute
+                      followed by 'ASC' (ascending) or 'DESC' (descending).
+                      Examples:
+                        - "metrics.validation_loss ASC"
+                        - "params.num_epochs DESC"
+                        - "attributes.start_time DESC"
+                      If None or an empty string, results are ordered by MLflow's default (usually start_time DESC).
+        max_results: Maximum number of runs to return (default: 100).
+    
+    Returns:
+        A dictionary containing a list of runs matching the criteria or an error message.
+        Format: {"runs": [run_details, ...]} or {"error": True, "message": ...}
+    """
+    # Validate experiment_id (must be non-empty)
+    if not experiment_id:
+        return {"error": True, "message": "Experiment ID cannot be empty"}
+    
+    # Validate max_results
+    if max_results <= 0:
+        return {"error": True, "message": "max_results must be a positive integer"}
+
+    # Ensure filter_string is not None, default to empty if it is (for mlflow.search_runs)
+    current_filter_string = filter_string if filter_string is not None else ""
+
+    found_runs: List[mlflow.entities.Run] # Type hint for the list of Run objects
+    try:
+        logger.info(f"Searching runs in experiment '{experiment_id}' with filter '{current_filter_string}', order by '{order_string}', max_results '{max_results}'")
+        
+        order_by_list = [order_string] if order_string and order_string.strip() else None
+
+        found_runs = mlflow.search_runs(
+            experiment_ids=[str(experiment_id)], # Ensure experiment_id is a string
+            filter_string=current_filter_string,
+            max_results=max_results,
+            order_by=order_by_list,
+            output_format="list"  # Get a list of Run objects instead of DataFrame
+        )
+    except Exception as e_search:
+        logger.error(f"MLflow search_runs API call failed for experiment_id '{experiment_id}': {str(e_search)}", exc_info=True)
+        return {"error": True, "message": f"MLflow search_runs API call failed: {str(e_search)}"}
+
+    processed_runs_info = []
+    if not found_runs:
+        logger.info(f"No runs found for experiment_id '{experiment_id}' with the given criteria.")
+        return {"runs": []} 
+
+    for run_obj in found_runs:
+        run_id_for_log = run_obj.info.run_id if run_obj.info else "N/A"
+        try:
+            start_time_ms = run_obj.info.start_time
+            end_time_ms = run_obj.info.end_time
+
+            run_details = {
+                "run_id": run_obj.info.run_id,
+                "status": run_obj.info.status,
+                "start_time": _format_timestamp(start_time_ms) if start_time_ms is not None else None,
+                "end_time": _format_timestamp(end_time_ms) if end_time_ms is not None else None,
+                "params": dict(run_obj.data.params),
+                "metrics": dict(run_obj.data.metrics),
+                "tags": dict(run_obj.data.tags)
+            }
+            processed_runs_info.append(run_details)
+        except Exception as e_process_run:
+            logger.warning(
+                f"Failed to process data for run_id '{run_id_for_log}' in experiment '{experiment_id}'. Error: {str(e_process_run)}. Skipping this run.",
+                exc_info=True
+            )
+            continue # Skip to the next run
+        
+    return {"runs": processed_runs_info}
+
+def list_registered_models() -> Dict:
+    """List all registered models."""
+    try:
+        logger.info("Listing registered models")
+        client = MlflowClient()
+        models = client.search_registered_models()
+        
+        return {
+            "models": [
+                {
+                    "name": model.name,
+                    "creation_timestamp": _format_timestamp(model.creation_timestamp),
+                    "last_updated_timestamp": _format_timestamp(model.last_updated_timestamp),
+                    "description": model.description or "",
+                    "tags": {tag.key: tag.value for tag in model.tags} if hasattr(model, "tags") else {},
+                    "latest_versions": [mv.version for mv in model.latest_versions]
+                }
+                for model in models
+            ]
+        }
+    except Exception as e:
+        return {"error": True, "message": f"Failed to list registered models: {str(e)}"}
+    
+def get_model_info(model_name: str) -> Dict:
+    """Get detailed information about a registered model."""
+    if not model_name:
+        return {"error": True, "message": "Model name cannot be empty"}
+    
+    try:
+        logger.info(f"Fetching info for model '{model_name}'")
+        client = MlflowClient()
+        model = client.get_registered_model(name=model_name)
+        
+        model_info = {
+            "name": model.name,
+            "creation_timestamp": _format_timestamp(model.creation_timestamp),
+            "last_updated_timestamp": _format_timestamp(model.last_updated_timestamp),
+            "description": model.description or "",
+            "tags": {tag.key: tag.value for tag in model.tags} if hasattr(model, "tags") else {},
+            "versions": []
+        }
+        
+        for mv in model.latest_versions:
+            run_id = mv.run_id
+            version_dict = {
+                "version": mv.version,
+                "current_stage": mv.current_stage,
+                "creation_timestamp": _format_timestamp(mv.creation_timestamp),
+                "last_updated_timestamp": _format_timestamp(mv.last_updated_timestamp),
+                "run": {}
+            }
+            
+            run = client.get_run(run_id)
+            version_dict["run"] = {
+                "status": run.info.status,
+                "start_time": _format_timestamp(run.info.start_time),
+                "end_time": _format_timestamp(run.info.end_time) if run.info.end_time else None,
+                "metrics": run.data.metrics
+            }
+            model_info["versions"].append(version_dict)
+        
+        return {"model": model_info}
+    except Exception as e:
+        return {"error": True, "message": f"Failed to fetch model info: {str(e)}"} 
+
+def register_model(
+    run_id: str, 
+    model_name: str, 
+    description: Optional[str] = None,
+    tags: Optional[Dict[str, str]] = None
+) -> Dict:
+    """
+    Register a model from a run, with optional description and tags.
+    Assumes artifact path is 'model'.
+    """
+    if not all([run_id, model_name]):
+        return {"error": True, "message": "Run ID and model name must be non-empty"}
+    
+    # Prepare description and tags
+    final_description = description + "Model registered by LLM through MCP service."
+    
+    final_tags = {
+        "registered_by": "mcp-llm-service",
+        "registration_timestamp": datetime.now().isoformat()
+    }
+    if tags:
+        final_tags.update(tags)
+
+    try:
+        logger.info(f"Registering model '{model_name}' from run '{run_id}/model' with description and tags.")
+        model_uri = f"runs:/{run_id}/model"
+        result = mlflow.register_model(
+            model_uri=model_uri, 
+            name=model_name,
+            tags=final_tags
+        )
+
+        client = MlflowClient()
+        client.update_model_version(
+            name=model_name,
+            version=result.version,
+            description=final_description
+        )
+        
+        return {
+            "model_name": model_name,
+            "version": result.version,
+            "description": final_description,
+            "tags": final_tags,
+            "message": "Registered successfully"
+        }
+    except Exception as e:
+        return {"error": True, "message": f"Registration failed: {str(e)}"}
+
+
+
+
+    
+
+set_tracking_uri("http://127.0.0.1:5000")
+#result = search_runs(1, "", "metrics.rmse ASC", 3)
+# result = list_registered_models()
+# print(result)
+# print(get_model_info("RandomForestBestModel"))
+
+print(register_model(
+    "a217f7261dbf4b0f8dbc575dad0e2f67", 
+    "RandomForestSecondBestModel",
+    "This is a test model."
+))
+# run_id: str, artifact_path: str, model_name: str

requirements.txt

@@ -0,0 +1,2 @@
+mlflow>=2.22.0
+gradio[mcp]>=5.32.0