Mlflowagent

🚀 MLflow MCP Server: Machine Learning Workflow with Natural Language Interface

This project provides a natural language interface for MLflow, connecting to the MLflow tracking server via the Model Context Protocol (MCP). It allows you to query your MLflow tracking server using plain English, making it easier to manage and explore your machine learning experiments and models.

🚀 Quick Start

Prerequisites

• Python 3.8+
• A running MLflow tracking server (default: http://localhost:8080)
• An OpenAI API key for the LLM

Installation

1. Clone this repository:

   git clone https://github.com/iRahulPandey/mlflowMCPServer.git
   cd mlflowMCPServer
   

2. Create a virtual environment:

   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate
   

3. Install the required packages:

   pip install mcp[cli] langchain-mcp-adapters langchain-openai langgraph mlflow
   

4. Set your OpenAI API key:

   export OPENAI_API_KEY=your_key_here
   

5. (Optional) Configure the MLflow tracking server URI:

   export MLFLOW_TRACKING_URI=http://localhost:8080
   

Usage

Start the MCP Server

First, start the MLflow MCP server:

python mlflow_server.py

This server connects to your MLflow tracking server and exposes MLflow functionality via MCP.

Execute Queries

Once the server is running, you can use the client to perform natural language queries:

python mlflow_client.py "What models do I have registered in MLflow?"

Example queries:

• "Show all the models I have registered."
• "List all my experiments."
• "Get the details of the model named 'iris-classifier'."
• "How is the status of my MLflow server?"

✨ Features

• Natural Language Queries: Ask questions about your MLflow tracking server in plain English.
• Model Registry Exploration: Obtain information about your registered models.
• Experiment Tracking: List and explore your experiments and runs.
• System Information: Get the status and metadata of the MLflow environment.

📚 Documentation

Overview

The MLflow MCP proxy consists of two main components:

1. MLflow MCP Server (mlflow_server.py): Connects to your MLflow tracking server and exposes MLflow functionality via the Model Context Protocol (MCP).
2. MLflow MCP Client (mlflow_client.py): Provides a natural language interface and uses a conversational AI assistant to interact with the MLflow MCP server.

Configuration

You can customize the behavior using environment variables:

• MLFLOW_TRACKING_URI: The URI of your MLflow tracking server (default: http://localhost:8080).
• OPENAI_API_KEY: The key for interacting with the OpenAI API.
• PYTHONPATH: Ensure that you can import mlflow and other dependencies.

MLflow MCP Server Script

The MLflow MCP server (mlflow_server.py) supports the following functions:

1. Model Registry Exploration:
   • List all registered models.
   • Get detailed information about a specific model.
2. Experiment Tracking:
   • List all experiments.
   • Get detailed information about a specific experiment.
3. System Information:
   • Get the status and metadata of the MLflow server.

🔧 Technical Details

Limitations

• Currently, only basic query functions are supported.
• Advanced analysis functions are still under development.

Future Improvements

• Increase support for more MLflow features.
• Improve the complexity and accuracy of queries.
• Support a multi - language interface.

📄 License

This project thanks the following projects for their contributions:

• Model Context Protocol (MCP): mlflow
• LangChain: langchain
• MLflow: mlflow

Through this project, we hope to provide a more intuitive and efficient way of interaction for MLflow users.