MLflow MCP 服务器：MLflow 的自然语言接口

本项目通过模型上下文协议 (MCP) 为 MLflow 提供了一个自然语言接口。它允许您使用简单的英语查询您的 MLflow 跟踪服务器，从而更轻松地管理和探索您的机器学习实验和模型。

概述

MLflow MCP 代理由两个主要组件组成：

1. MLflow MCP 服务器 (mlflow_server.py): 连接到您的 MLflow 跟踪服务器，并通过模型上下文协议 (MCP) 公开 MLflow 功能。

2. MLflow MCP 客户端 (mlflow_client.py): 提供一个自然语言接口，用于使用对话式 AI 助手与 MLflow MCP 服务器交互。

特性

• 自然语言查询: 使用简单的英语提问关于您的 MLflow 跟踪服务器的问题
• 模型注册表探索: 获取关于您注册模型的信息
• 实验跟踪: 列出并探索您的实验和运行
• 系统信息: 获取关于您的 MLflow 环境的状态和元数据

前提条件

• Python 3.8+
• MLflow 服务器正在运行 (默认: http://localhost:8080)
• OpenAI API 密钥用于 LLM

安装

1. 克隆此仓库:

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

2. 创建一个虚拟环境:

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

3. 安装所需的包:

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

4. 设置您的 OpenAI API 密钥:

   export OPENAI_API_KEY=your_key_here
   

5. (可选) 配置 MLflow 跟踪服务器 URI:

   export MLFLOW_TRACKING_URI=http://localhost:8080
   

用法

启动 MCP 服务器

首先，启动 MLflow MCP 服务器:

python mlflow_server.py

服务器连接到您的 MLflow 跟踪服务器，并通过 MCP 公开 MLflow 功能。

进行查询

一旦服务器运行，您可以使用客户端进行自然语言查询:

python mlflow_client.py "我在 MLflow 中注册了哪些模型？"

示例查询:

• "显示我在 MLflow 中所有注册的模型"
• "列出我所有的实验"
• "获取名为 'iris-classifier' 的模型的详细信息"
• "我的 MLflow 服务器的状态是什么？"

配置

您可以使用环境变量自定义行为:

• MLFLOW_TRACKING_URI: 您的 MLflow 跟踪服务器的 URI (默认: http://localhost:8080)
• OPENAI_API_KEY: 您的 OpenAI API 密钥
• MODEL_NAME: 要使用的 OpenAI 模型 (默认: gpt-3.5-turbo-0125)
• MLFLOW_SERVER_SCRIPT: MLflow MCP 服务器脚本的路径 (默认: mlflow_server.py)
• LOG_LEVEL: 日志级别 (默认: INFO)

MLflow MCP 服务器 (mlflow_server.py)

服务器连接到您的 MLflow 跟踪服务器，并通过 MCP 公开以下工具:

• list_models: 列出 MLflow 模型注册表中所有注册的模型
• list_experiments: 列出 MLflow 跟踪服务器中的所有实验
• get_model_details: 获取关于特定注册模型的详细信息
• get_system_info: 获取关于 MLflow 跟踪服务器和系统的信息

局限性

• 目前仅支持 MLflow 功能的一个子集
• 客户端需要互联网访问才能使用 OpenAI 模型
• 对于复杂的 MLflow 操作，错误处理可能受到限制

未来改进

• 增加对 MLflow 模型预测的支持
• 改进自然语言理解，以支持更复杂的查询
• 增加指标和参数的可视化功能
• 支持更多 MLflow 操作，如运行管理和工件处理

致谢

• 模型上下文协议 (MCP): 用于协议规范
• LangChain: 用于代理框架
• MLflow: 用于跟踪和模型注册表功能