Databricks Unity CatalogとAnthropic SDKの連携

少し前に公開されていたUnity Catalog AIのライブラリを使ってみます。

OSS版UCのページの方が詳しいです。

PyPIからUnity Catalog AI Anthropicインテグレーションをインストールします。

%pip install unitycatalog-anthropic[databricks]
%pip install -U mlflow
dbutils.library.restartPython()

トレースをオンにします。

import mlflow
mlflow.anthropic.autolog()

APIキーを読み込みます。

import os
os.environ["ANTHROPIC_API_KEY"] = dbutils.secrets.get(scope="demo-token-takaaki.yayoi", key="anthropic_api_key")

Unity Catalogの関数クライアントのインスタンスを作成します。

from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

PythonでUnity Catalogの関数を作成します。

CATALOG = "takaakiyayoi_catalog"
SCHEMA = "agents"

func_name = f"{CATALOG}.{SCHEMA}.weather_function"

def weather_function(location: str) -> str:
  """
  指定された場所の現在の天気を摂氏で取得します。

  引数:
    location (str): 天気を取得する場所。
  戻り値:
    str: 指定された場所の現在の気温（摂氏）。
  """
  return f"The current temperature for {location} is 24.5 celsius"

client.create_python_function(
  func=weather_function,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True
)

関数が登録されました。

Unity Catalog関数を用いてツールキットのインスタンスを作成します。

from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit

# ツールキットのインスタンスを作成
toolkit = UCFunctionToolkit(function_names=[func_name], client=client)

Anthropic SDKでツールコールを使用します。

import anthropic

# AnthropicクライアントをAPIキーで初期化
anthropic_client = anthropic.Anthropic()

# ユーザーの質問
question = [{"role": "user", "content": "ニューヨーク市の天気はどうですか？"}]

# Anthropicへの初回呼び出しを行う
response = anthropic_client.messages.create(
  model="claude-opus-4-20250514",  # モデルを指定
  max_tokens=1024,  # 'max_tokens_to_sample'の代わりに'max_tokens'を使用
  tools=toolkit.tools,
  messages=question  # 会話履歴を提供
)

# レスポンスの内容を表示
display(response)

INFO:httpx:HTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
Message(id='msg_01EKmXH5u2DYn9FJH3R5Fw7V', content=[TextBlock(citations=None, text='ニューヨーク市の現在の天気を確認させていただきます。', type='text'), ToolUseBlock(id='toolu_01CFYABVFsv72qXa6ADiRnPp', input={'location': 'ニューヨーク市'}, name='takaakiyayoi_catalog__agents__weather_function', type='tool_use')], model='claude-opus-4-20250514', role='assistant', stop_reason='tool_use', stop_sequence=None, type='message', usage=Usage(cache_creation_input_tokens=0, cache_read_input_tokens=0, input_tokens=672, output_tokens=95, server_tool_use=None, service_tier='standard'))

ツールのレスポンスを作成します。Claudeモデルからのレスポンスには、ツールを呼び出す必要がある場合に、ツール要求メタデータのブロックが含まれます。

from unitycatalog.ai.anthropic.utils import generate_tool_call_messages

# UC関数を呼び出し、必要なフォーマット済みレスポンスを構築
tool_messages = generate_tool_call_messages(
  response=response,
  client=client,
  conversation_history=question
)

# Anthropicとの会話を続ける
tool_response = anthropic_client.messages.create(
  model="claude-opus-4-20250514",
  max_tokens=1024,
  tools=toolkit.tools,
  messages=tool_messages,
)

display(tool_response)

Message(id='msg_01SHxCo1WgVRdKMbPVtxMEk9', content=[TextBlock(citations=None, text='ニューヨーク市の現在の気温は24.5度（摂氏）です。過ごしやすい気温ですね。', type='text')], model='claude-opus-4-20250514', role='assistant', stop_reason='end_turn', stop_sequence=None, type='message', usage=Usage(cache_creation_input_tokens=0, cache_read_input_tokens=0, input_tokens=799, output_tokens=43, server_tool_use=None, service_tier='standard'))

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル