【databricks】Claude MCPをdatabricksで使ってみた2: Claudeとdatabricks LLMのgenieと接続する方法

databricks host ID                            genie space ID 
┌────────────────────────────────┐             ┌─────────┐ 
https://xxxxx.cloud.databricks.com/genie/rooms/xxxxxxxxxxx?o=xxxxxxxx

my-genie-mcp/
 ├ main.py 
 ├ dbapi.py
 ├ databricks_formatter.py
 ├ pyproject.toml

"""
Databricks Genie MCP Server
MCPを使用してGenieとの対話を可能にするサーバー
"""
from mcp.server.fastmcp import FastMCP
from dbapi import genie_conversation
from databricks_formatter import format_query_results

# MCPサーバーを初期化
# サーバー名を指定して新しいインスタンスを作成
mcp = FastMCP("Databricks Genie MCP Server")

@mcp.tool()
async def start_conversation(content: str) -> str:
  """
  Genie APIとの会話を開始する

  Args:
    content (str): Genieに送信するメッセージ内容

  Returns:
    str: Genieからのレスポンスを整形した結果
  """
  try:
     # Genie APIを呼び出して応答を取得
     result = await genie_conversation(content)
     # 応答結果を見やすい形式に整形
     return format_query_results(result)
  except Exception as e:
     # エラーが発生した場合はエラーメッセージを返す
     return result

# メインエントリーポイント
if __name__ == "__main__":
  # デフォルトでstdioトランスポートを使用してサーバーを起動
  # これによりコマンドラインからの対話が可能になる
  mcp.run()

"""
Databricks API通信モジュール
Genieとの対話に必要なAPI呼び出しを管理します
"""

from typing import Any, Dict, Optional
import os
import asyncio
import httpx
from dotenv import load_dotenv

# 環境変数を.envファイルから読み込み
load_dotenv()

# Databricks接続に必要な設定値
DATABRICKS_HOST = os.environ.get("DATABRICKS_HOST", "<your databricks host>")
DATABRICKS_TOKEN = os.environ.get("DATABRICKS_TOKEN", "<your databricks token>")
DATABRICKS_GENIE_SPACE_ID = os.environ.get("DATABRICKS_GENIE_SPACE_ID", "<your genie space ID>")

# APIエンドポイントの定義
GENIE_START_CONVERSATION_API = "/api/2.0/genie/spaces/{space_id}/start-conversation" # 会話開始用エンドポイント
GENIE_GET_MESSAGE_API = "/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}" # メッセージ取得用エンドポイント
STATEMENT_API = "/api/2.0/sql/statements/{statement_id}" # SQLステートメント実行結果取得用エンドポイント

async def make_databricks_request(
   method: str,
  endpoint: str,
  json_data: Optional[Dict[str, Any]] = None,
  params: Optional[Dict[str, Any]] = None
  ) -> Dict[str, Any]:
  """
  DatabricksのAPIにリクエストを送信する汎用関数

  Args:
    method (str): HTTPメソッド('get'または'post')
    endpoint (str): APIエンドポイントのパス
    json_data (Optional[Dict[str, Any]]): POSTリクエスト用のJSONデータ
    params (Optional[Dict[str, Any]]): GETリクエスト用のクエリパラメータ

  Returns:
    Dict[str, Any]: APIレスポンスのJSONデータ
  """
  url = f"{DATABRICKS_HOST}{endpoint}"

  # 認証ヘッダーの設定
  headers = {
    "Authorization": f"Bearer {DATABRICKS_TOKEN}",
  }

  async with httpx.AsyncClient() as client:
    try:
      # HTTPメソッドに応じてリクエストを実行
      if method.lower() == "get":
         response = await client.get(url, headers=headers, params=params, timeout=30.0)
      elif method.lower() == "post":
        response = await client.post(url, headers=headers, json=json_data, timeout=30.0)
      else:
        raise ValueError(f"サポートされていないHTTPメソッド: {method}")

      response.raise_for_status()
      return response.json()
    except httpx.HTTPStatusError as e:
      # HTTPエラーの詳細情報を取得して例外を発生
      error_message = f"HTTPエラー: {e.response.status_code}"
      try:
        error_detail = e.response.json()
        error_message += f" - {error_detail.get('message', '')}"
      except Exception:
        pass
     raise Exception(error_message)
   except Exception as e:
     raise Exception(f"Databricks APIリクエスト中にエラーが発生: {str(e)}")


async def genie_conversation(content: str, space_id: Optional[str] = None) -> Dict[str, Any]:
  """
  Genieとの会話を実行し、結果が得られるまで待機する

  Args:
    content (str): Genieに送信するメッセージ内容
    space_id (Optional[str]): Genie Space ID(未指定時は環境変数から取得)

  Returns:
    Dict[str, Any]: 会話の実行結果
    """
  if not space_id:
   space_id = DATABRICKS_GENIE_SPACE_ID

  if not space_id:
    raise ValueError("Genie Space IDが必要です。環境変数DATABRICKS_GENIE_SPACE_IDを設定するか、パラメータとして指定してください。")

  # 会話開始リクエストの作成
  statement_data = {
    "content": content
  }

  # 会話を開始し、conversation_idとmessage_idを取得
  endpoint_url = GENIE_START_CONVERSATION_API.format(space_id=space_id)
  response = await make_databricks_request("post", endpoint_url, json_data=statement_data)
  message = response.get("message")
  conversation_id = message["conversation_id"]
  message_id = message["id"]

  if not conversation_id:
    raise Exception("レスポンスからconversation_IDの取得に失敗しました")

  # 会話生成完了をポーリング
  max_retries = 60 # 最大リトライ回数（10秒間隔で10分）
  retry_count = 0

  while retry_count < max_retries:
    # メッセージのステータスを確認
    message_status = await make_databricks_request(
     "get",
     GENIE_GET_MESSAGE_API.format(space_id=space_id, conversation_id=conversation_id, message_id=message_id)
    )

   status = message_status["status"]

   if status == "COMPLETED":
      # 完了した場合、クエリ結果を取得
      query_result_statement_id = message_status["attachments"][0]["query"]["statement_id"]

     statement_status = await make_databricks_request(
       "get",
       STATEMENT_API.format(statement_id=query_result_statement_id)
     )

     return statement_status
   elif status in ["FAILED", "CANCELED"]:
     error_message = message_status["status"]
     raise Exception(f"メッセージの取得に失敗: {error_message}")

   # 次のポーリングまで待機

    retry_count += 1

  return message["content"]

from typing import Any, Dict

def format_query_results(result: Dict[str, Any]) -> str:
  """Format query results into a readable string."""

  # Check if result is empty or doesn't have the expected structure
  if not result or 'manifest' not in result or 'result' not in result:
    return "No results or invalid result format."

  # Extract column names from the manifest
  column_names = []
  if 'manifest' in result and 'schema' in result['manifest'] and 'columns' in result['manifest']['schema']:
    columns = result['manifest']['schema']['columns']
    column_names = [col['name'] for col in columns] if columns else []

  # If no column names were found, return early
  if not column_names:
    return "No columns found in the result."

  # Extract rows from the result
  rows = []
  if 'result' in result and 'data_array' in result['result']:
    rows = result['result']['data_array']

    # If no rows were found, return just the column headers
    if not rows:
      # Format as a table
      output = []
      # Add header
      output.append(" | ".join(column_names))
      output.append("-" * (sum(len(name) + 3 for name in column_names) - 1))
      output.append("No data rows found.")
      return "\n".join(output)

  # Format as a table
  output = []

  # Add header
  output.append(" | ".join(column_names))
  output.append("-" * (sum(len(name) + 3 for name in column_names) - 1))
  # Add rows
  for row in rows:
    row_values = []
    for value in row:
      if value is None:
        row_values.append("NULL")
      else:
        row_values.append(str(value))
    output.append(" | ".join(row_values))

  return "\n".join(output)

[project]
name = "databricks_test"
version = "0.1.0"
description = "Databricks Genie Integration"
requires-python = ">=3.10"
dependencies = [
"httpx>=0.28.1",
"mcp[cli]>=1.3.0",
]

{
"mcpServers": {
"databricks_genie_test": {
"command": "uv",
"args": [
"--directory",
"<your directory to file>/my-genie-mcp",
"run",
"main.py"
],
"env": {
"DATABRICKS_HOST": "<your databricks host>",
"DATABRICKS_TOKEN": "<your databricks token>",
"DATABRICKS_GENIE_SPACE_ID": "<your genie space ID>"
}
}
}
}