Databricks AI 开发套件 - 为AI驱动的开发赋能Databricks AI Dev Kit 是一个开源工具

Databricks AI Dev Kit

Databricks Certified Gold Project

概述

AI驱动开发（或称 "vibe coding"）在 Databricks 上从未如此简单。AI Dev Kit 为您的AI编程助手（如Claude Code、Cursor、Windsurf等）提供了构建 Databricks 应用所需的可信知识和工具，助其更快、更智能地完成工作。

架构图转存失败，建议直接上传图片文件

功能特性

全面的 Databricks 技能库：为AI助手提供关于 Spark 声明式管道、Unity Catalog、Databricks Jobs、MLflow、Model Serving、Databricks Apps 等的最佳实践和模式文档，确保生成的代码符合规范。
强大的 MCP 工具集成：通过 Model Context Protocol (MCP) 服务器，为AI助手提供可直接执行的 Databricks 操作，如执行 SQL、管理集群、操作 UC 对象、上传文件等。
一键式安装体验：提供统一的安装脚本，可轻松将技能和 MCP 服务器集成到 Claude Code、Cursor 等主流 AI 工具中，支持项目级或全局安装。
内置技能测试框架：包含一个 YAML 优先的评估框架，支持通过真实 Databricks 环境执行代码来验证技能的有效性，并通过“生成-审查-提升”工作流管理测试用例。
可嵌入的 Agent 服务：提供了一个可直接嵌入自定义应用的 Agent 服务，该服务集成了 Claude Agent SDK 和所有 Databricks MCP 工具，可作为“服务式代理”使用。
完整的 Web 应用示例：附带一个功能完备的 Builder App，展示了一个带有聊天界面、项目管理、对话历史的交互式应用，并演示了在 FastAPI 环境中运行 Claude Agent 的最佳实践（包括事件循环修复）。

安装指南

系统要求

Python 3.11 或更高版本
Node.js 18+ 和 npm (仅当运行 Builder App 时需要)
uv 包管理器（推荐，用于加速依赖安装）
一个 Databricks 工作区及有效的认证凭证

快速安装

您可以使用一行命令将 AI Dev Kit 安装到您当前的项目目录中：

bash <(curl -sL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/install.sh)

安装选项示例：

# 全局安装（例如对 Claude Code）
bash <(curl -sL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/install.sh) --global

# 指定 Databricks 配置文件并强制重新安装
bash <(curl -sL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/install.sh) --profile DEFAULT --force

# 仅为特定工具安装（如 Cursor 和 GitHub Copilot）
bash <(curl -sL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/install.sh) --tools cursor,copilot

# 仅安装技能库（跳过 MCP 服务器）
bash <(curl -sL https://raw.githubusercontent.com/databricks-solutions/ai-dev-kit/main/install.sh) --skills-only

配置认证

安装后，您需要配置 Databricks 认证。可以通过以下任一方式：

配置文件：在 ~/.databrickscfg 中设置您的凭证和默认配置文件。

环境变量：

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"
# 或
export DATABRICKS_CONFIG_PROFILE="your-profile-name"

开发环境设置（Builder App）

如果您想运行或开发附带的 Builder App：

进入应用目录：cd databricks-builder-app
运行安装脚本：./scripts/setup.sh
根据提示编辑 .env.local 文件，填入您的 Databricks 和数据库凭证。
启动开发服务器：./scripts/start_dev.sh

使用说明

在 AI 助手中使用

安装成功后，您的 AI 助手（如 Claude Code）将自动获得 Databricks 技能和工具。您可以直接向其提问：

示例提示词：

"帮我创建一个 bronze 层的数据管道，从 /Volumes/raw/orders 目录摄取 JSON 文件。"
"列出我工作区中所有正在运行的集群。"
"使用 Python SDK 创建一个新的 Databricks 作业，该作业每天运行一次。"
"帮我部署一个简单的 MLflow 模型到一个 Serving Endpoint。"

助手会利用技能库中的知识生成符合 Databricks 最佳实践的代码，并使用 MCP 工具直接与您的 Databricks 环境交互。

将 Agent 嵌入您的应用

您可以使用 server/services/agent.py 中的 stream_agent_response 函数，将 Databricks AI Agent 集成到您自己的 Python 应用中。

import asyncio
from server.services.agent import stream_agent_response
from databricks_tools_core.auth import set_databricks_auth, clear_databricks_auth

async def run_my_agent(user_message):
    # 设置用户认证上下文
    set_databricks_auth("workspace_url", "user_token")
    try:
        async for event in stream_agent_response(
            project_id="my-project",
            message=user_message,
            databricks_host="workspace_url",
            databricks_token="user_token",
        ):
            # 处理事件流，例如 text_delta, tool_use 等
            if event["type"] == "text_delta":
                print(event["text"], end="")
            elif event["type"] == "tool_use":
                print(f"\n[使用工具: {event['tool_name']}]")
    finally:
        # 清除认证上下文
        clear_databricks_auth()

# 运行
asyncio.run(run_my_agent("列出我的 SQL 仓库"))

核心代码

1. 安装脚本核心逻辑 (`install.sh`)

这是项目的统一安装器，负责解析参数、设置环境、并调用具体的安装逻辑。它支持通过命令行参数或环境变量进行配置。

#!/bin/bash
# Databricks AI Dev Kit - Unified Installer
#
# Usage: bash <(curl -sL ...) [OPTIONS]

set -e

# 默认配置
PROFILE="${DEVKIT_PROFILE:-DEFAULT}"
SCOPE="${DEVKIT_SCOPE:-project}"
FORCE="${DEVKIT_FORCE:-false}"
TOOLS="${DEVKIT_TOOLS:-}"

# 解析命令行参数
while [[ $# -gt 0 ]]; do
  case $1 in
    --global) SCOPE="global"; SCOPE_EXPLICIT=true; shift ;;
    --force) FORCE=true; shift ;;
    --profile) PROFILE="$2"; shift 2 ;;
    --tools) TOOLS="$2"; shift 2 ;;
    --skills-only) SKILLS_ONLY=true; shift ;;
    *) echo "Unknown option: $1"; exit 1 ;;
  esac
done

# 核心安装流程：克隆仓库、复制技能、配置工具等
# ... (具体安装步骤)
echo "Databricks AI Dev Kit installation complete."

2. Agent 服务中的事件循环修复 (`server/services/agent.py`)

此代码解决了 claude-agent-sdk 在 FastAPI/uvicorn 环境中因事件循环冲突导致工具无法执行的问题。它通过在新线程中创建新的事件循环来运行 Agent，并确保认证上下文被正确复制。

def _run_agent_in_fresh_loop(message, options, result_queue, context):
    """在一个全新的线程和事件循环中运行agent，以解决事件循环冲突。"""
    # 设置新的事件循环
    asyncio.set_event_loop(asyncio.new_event_loop())
    loop = asyncio.get_event_loop()

    async def run():
        # 复制调用方的上下文变量（用于认证）
        for var, value in context.items():
            var.set(value)

        # 调用 Agent SDK 的查询方法
        async for event in query(message, options):
            result_queue.put(event)  # 将事件放回队列，供主循环处理

    try:
        loop.run_until_complete(run())
    finally:
        loop.close()

async def stream_agent_response(...):
    # ... 初始化
    # 创建上下文变量快照
    context = {
        _databricks_host_token_var: (host, token),
        # ... 其他上下文变量
    }

    result_queue = queue.Queue()
    # 在新线程中运行 Agent
    thread = threading.Thread(
        target=_run_agent_in_fresh_loop,
        args=(message, options, result_queue, context),
        daemon=True,
    )
    thread.start()

    # 在主事件循环中从队列获取事件并生成
    while thread.is_alive() or not result_queue.empty():
        try:
            event = result_queue.get(timeout=0.1)
            yield event
        except queue.Empty:
            await asyncio.sleep(0.05)

3. 技能测试框架的 YAML 数据模型 (`src/skill_test/dataset.py`)

该代码定义了技能测试框架中用于表示测试用例的数据模型，展示了如何通过 YAML 文件来驱动评估流程，并为未来集成 Unity Catalog 预留了接口。

from dataclasses import dataclass
from pathlib import Path
from typing import List, Dict, Any, Optional, Protocol
import yaml

@dataclass
class EvalRecord:
    """标准评估记录格式，与 MLflow 评估模式兼容。"""
    id: str
    inputs: Dict[str, Any]
    outputs: Optional[Dict[str, Any]] = None
    expectations: Optional[Dict[str, Any]] = None
    metadata: Optional[Dict[str, Any]] = None

    def to_eval_dict(self) -> Dict[str, Any]:
        """转换为 MLflow 评估所需的字典格式。"""
        result = {"inputs": self.inputs}
        if self.outputs:
            result["outputs"] = self.outputs
        if self.expectations:
            result["expectations"] = self.expectations
        return result

class DatasetSource(Protocol):
    """数据集源的协议，支持未来从 Unity Catalog 加载数据。"""
    def load(self) -> List[EvalRecord]: ...

@dataclass
class YAMLDatasetSource:
    """从 YAML 文件加载评估数据集 (Phase 1 实现)。"""
    yaml_path: Path

    def load(self) -> List[EvalRecord]:
        """从 ground_truth.yaml 文件加载记录。"""
        with open(self.yaml_path) as f:
            data = yaml.safe_load(f)

        records = []
        for case in data.get("test_cases", []):
            records.append(EvalRecord(...))
        return records

4. Builder App 的数据库迁移配置 (`alembic/env.py`)

此代码展示了如何在需要动态 OAuth 令牌的生产环境（Databricks Apps）中配置 Alembic 以连接到 Lakebase PostgreSQL。它使用 Databricks SDK 动态生成数据库凭证。

def get_url_and_connect_args():
    """从环境变量获取数据库 URL 和连接参数。"""
    global _resolved_hostaddr
    connect_args = {}

    url = os.environ.get('LAKEBASE_PG_URL')

    if not url:
        # 动态 OAuth 模式
        instance_name = os.environ.get('LAKEBASE_INSTANCE_NAME')
        # ... 检查 instance_name

        # 使用 Databricks SDK 生成令牌
        w = WorkspaceClient()
        instance = w.database.get_database_instance(name=instance_name)
        cred = w.database.generate_database_credential(...)

        # 构建包含令牌的 URL
        url = f'postgresql://{encoded_username}:{cred.token}@{host}:5432/{database_name}?sslmode=require'

        # ... 解析主机名以解决 macOS DNS 问题

    # 确保 URL 使用同步驱动进行迁移
    url = url.replace('postgresql+asyncpg://', 'postgresql://', 1)
    return url, connect_args
```FINISHED
xYXos8PemElNbrK1+4OeBBJy0jdEHLXpq0MBZcgHvTY=

Databricks AI 开发套件 - 为AI驱动的开发赋能