AI导哈

数据查询代理

Databricks MCP代理示例展示了如何利用托管MCP服务器创建智能数据查询代理，集成Genie空间和Unity Catalog功能，提供基于LLM的智能响应。

  • by IanGagnonDB
  • Python
  • 本地部署
  • 1 星标
  • 更新于
在 GitHub 查看 IanGagnonDB/databricks-agent-mcp-genie

使用说明

_CHMCP代理示例

本项目演示了如何使用Genie的托管MCP（模型上下文协议）服务器在Rancher中创建代理。该代理与Rancher Genie空间和其他启用MCP的服务集成，以提供智能数据查询功能。

概述

MCP代理利用Rancher的托管MCP服务器创建一个会话式AI代理，该代理可以：

  • 通过Genie空间查询数据
  • 访问Unity目录功能
  • 与各种copula资源交互
  • 使用LLM端点提供智能响应

项目结构

  • driver.py -记录、注册和部署代理的主编排脚本
  • mcp_agent.py -使用MLF的ResponsesAgent框架实现核心代理
  • test_mcp_server.py -验证MCP服务器连接的测试脚本
  • requirements.txt -Python依赖关系
  • README.md -本文件

先决条件

可选：部署VM以访问专用工作区

如果您的copula工作区 不允许公共网络访问 （即，它仅配置为私有访问），您需要在工作区的虚拟网络（ExpressRoute）中有一个跳转框来访问copula和相关资源。

推荐方法（Azure）：

  1. 部署虚拟机（VM）：

    • 使用Azure门户或命令行界面（CLI）将Ubuntu 24.04虚拟机（VM）部署到与您的ViewModel工作区相同的ExpressRoute中。
    • 为虚拟机分配一个公共IP地址。
    • 允许在端口22（SSH）上对您的个人IP地址进行入站访问。

  2. 连接到VM：

    • SSH进入VM：
      ssh <your-username>@<public-ip-address>

  3. 设置Python环境：

    • 安装Python工具：
      sudo apt update
sudo apt install python3-pip python3-venv -y
    • 创建并激活虚拟环境：
      python3 -m venv venv
source venv/bin/activate

现在，您可以继续进行其余的设置，并从该VM中运行代理，确保与您的copula工作区的安全、私有连接。

Python环境设置

必须使用提供的正确设置Python环境 requirements.txt 文件：

pip install -r requirements.txt

所需的环境变量

在运行项目之前，必须定义以下环境变量：

# Databricks Configuration
export DATABRICKS_CLI_PROFILE="your-databricks-profile"
export EXPERIMENT_NAME="your-mlflow-experiment-name"

# Model and Endpoint Configuration
export LLM_ENDPOINT_NAME="your-llm-endpoint-name"
export GENIE_SPACE_ID="your-genie-space-id"

# Unity Catalog Configuration
export CATALOG_NAME="your-catalog-name"
export SCHEMA_NAME="your-schema-name"
export MODEL_NAME="your-model-name"

# Table Access (JSON array format)
export TABLE_NAMES='["table1", "table2", "table3"]'

copula设置

  1. copula命令行界面：确保 copula命令行界面 已安装并配置了适当的身份验证
  2. 工作区访问：您的个人资料必须能够访问目标的copula工作区
  3. 精灵空间：具有指定值的已配置Genie空间 GENIE_SPACE_ID
  4. LLM端点：访问LLM的服务端点（例如Claude、GPT模型）
  5. Unity目录：对指定目录和架构的适当权限

用法

1.测试MCP服务器连接

首先，验证您的MCP服务器连接是否正常工作：

python test_mcp_server.py

这将列出您配置的MCP服务器中的可用工具。

2.部署代理

运行主驱动程序脚本以记录、注册和部署您的代理：

python driver.py

此脚本将：

  1. 将代理模型记录到MLflow
  2. 在Unity目录中注册模型
  3. 将代理部署为服务端点

配置

MCP服务器URL

代理可以连接到多个MCP服务器。在中配置它们 mcp_agent.py:

MANAGED_MCP_SERVER_URLS = [
    f"{host}/api/2.0/mcp/genie/{GENIE_SPACE_ID}",
    # Add more managed MCP servers as needed
]

CUSTOM_MCP_SERVER_URLS = [
    # Add custom MCP servers hosted on Databricks Apps
]

系统提示

通过修改来自定义代理的行为 SYSTEM_PROMPT 在……里面 mcp_agent.py:

SYSTEM_PROMPT = "You are a helpful assistant specialized in data analysis."

建筑

代理遵循此架构：

  1. 请求处理：通过ResponsesAgent界面接收用户查询
  2. 工具发现：从配置的MCP服务器动态发现可用工具
  3. LLM集成：使用copula服务端点进行语言模型推理
  4. 工具执行：根据LLM决策执行MCP工具（如Genie查询）
  5. 响应生成：根据工具输出提供智能响应

安全与认证

  • 该代理对资源访问使用copula系统身份验证
  • 中的所有指定资源 driver.py 自动授予访问权限
  • MCP服务器通过copula工作区客户端进行身份验证

故障排除

常见问题

  1. 身份验证错误：验证您的 DATABRICKS_CLI_PROFILE 已正确配置
  2. 资源访问：确保环境变量中指定的所有表和端点都是可访问的
  3. MCP服务器连接：使用 test_mcp_server.py 验证服务器连接
  4. 环境变量：仔细检查是否设置了所有必需的环境变量

调试

  • 检查MLflow实验日志以了解模型注册问题
  • 通过copula工作区UI验证Genie空间的可访问性
  • 使用以下工具测试单个MCP工具 DatabricksMCPClient 直接

查询示例

部署后，您的代理可以处理以下查询：

  • “谁创造了最多的管道？”
  • “显示最新的数据质量指标”
  • “本月表现最好的模特是什么？”

代理将自动确定使用哪些工具，并根据您的数据提供智能响应。

相关 MCP