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

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

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 仓库"))

#!/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."

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)

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

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