LoginSignup
4
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@isanakamishiro2(T Yamaya)

DatabricksのAIエージェントツールを様々なケースで作ってみる

Last updated at Posted at 2025-02-07

導入

Databricksは以前からAIエージェントの構築をサポートしていますが、個人的にその中でもAIエージェントが実行するツール管理は非常に重要だと考えています。
（それ以上に評価が大事だとも思ってますが）

DatabricksではUnity Catalog関数をツールとして利用できるため、アクセスコントロール含めてUnity Catalog上でツール管理ができます。

最近このあたりの機能も増えてるようなので、いろんなバリエーションでツール作成をしてみます。
基本的には以下ドキュメントのウォークスルーのような内容となります。

各種コードの実行はDatabrikcs on AWS上で実施しました。クラスタはサーバレスです。

準備

Pythonノートブックを作成し、関係するパッケージをインストールします。

# Databricks SDKの最新化
%pip install -U databricks-sdk

# Unity Catalog-AI
%pip install unitycatalog-ai[databricks]

# Unity CatalogのLangChain統合
%pip install unitycatalog-langchain[databricks]

# Databricks LangChain用パッケージ
%pip install -U databricks-langchain

%restart_python

その1: Python関数をツール化する

unitycatalog-aiパッケージを使うことで、Python関数を簡単にUnity Catalog関数(ツール)として登録することができます。
以下は、add_numbers関数という数字を二つ合算する関数をUnity Catalog関数として登録しています。(カタログ・スキーマの場所は適当です)

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

CATALOG = "training"
SCHEMA = "llm"

def add_numbers(number_1: float, number_2: float) -> float:
  """
  A function that accepts two floating point numbers adds them,
  and returns the resulting sum as a float.

  Args:
      number_1 (float): The first of the two numbers to add.
      number_2 (float): The second of the two numbers to add.

  Returns:
      float: The sum of the two input numbers.
  """
  return number_1 + number_2

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

カタログエクスプローラで確認すると以下のようになります。
利用可能なパッケージは制限がある（はず）ですが、Python関数をすぐにツール化できるので便利ですね。

image.png

その2: Pythonコードを実行するツール

system.aiカタログ・スキーマにはpython_execという関数が登録されています。
これはパラメータとして与えられたPythonコードを実行し、出力を結果として返すUnity Catalog関数です。

カタログエクスプローラで見ると、以下のように登録されています。

image.png

LLMで生成したPythonコードを実行するなど、非常に汎用性が高いツールですね。

ちなみにUnity Catalog関数なので、以下のようにSQL等でそのまま実行もできます。

%sql
SELECT system.ai.python_exec("""
import random
numbers = [random.random() for _ in range(10)]
print(numbers)
""") as a

その3: テーブルクエリ ツール

こちらのドキュメントより。

エージェント処理の中でテーブルからレコードを取り出すツールを使うケースは多いと思うのですが、以下のようなUnity Catalog関数を作成することで実現できます。

こちらはsamplesカタログにあるbakehouse.sales_customersテーブルから情報を取得する関数です。

%sql
CREATE OR REPLACE FUNCTION training.llm.lookup_customer_info(
    f_name STRING COMMENT 'First Name of the customer whose info to look up',
    l_name STRING COMMENT 'Last Name of the customer whose info to look up'
  )
  RETURNS STRING
  COMMENT 'Returns metadata about a particular customer given the customer name, including the customer email and ID. The customer ID can be used for other queries.'
  RETURN
    SELECT
      CONCAT('Customer ID: ', customerID, ', ', 'Customer Email: ', email_address)
    FROM
      samples.bakehouse.sales_customers
    WHERE
      first_name = f_name
      and last_name = l_name
    LIMIT 1;

Playgroundでこのツールを使って問い合わせをすると以下のようになります。

image.png

※　smithbrian@example.orgが取得されるのは、もともとそういうデータのため。正しくデータは取得できています。

その4: ベクトルインデックス ツール

こちらのドキュメントより。

Databricks Mosaic AI Vector Searchのインデックスを使った検索をUnity Catalog関数(ツール)として利用できます。
この例では以前作成したインデックスsample_documents_indexを指定してベクトル検索を行う関数を作成しました。

%sql
CREATE
OR REPLACE FUNCTION training.llm.sample_docs_vector_search (
  query STRING
  COMMENT 'オープンデータ文書を検索するためのクエリ'
) RETURNS TABLE
COMMENT '入力クエリに最も関連するテキスト文書を取得するために、日本のオープンデータ文書を検索します。' 
RETURN
SELECT
  chunk as page_content,
  id as chunk_id
FROM
  vector_search(
    -- ここにベクトル検索インデックス名を指定します
    index => 'training.llm.sample_documents_index',
    query => query,
    num_results => 5
  )

このツールを呼びだすエージェントを作ることで、簡単にAgentic RAGが実装できます。

その5: 外部サービス接続ツール

こちらのドキュメントより。

比較的最近リリースされた機能で、利用にあたってはプレビューメニューから、以下のオプションを有効化する必要があるようです。

image.png

HTTPリクエストを使用して外部APIを実行するUnity Catalog関数を作ることができます。
接続先をUnity Catalogの"接続(Connection)"として登録することで、Unity Catalogのアクセスコントロール等の機能を利用しつつ外部APIをエージェントから呼び出すことができます。

今回は例としてDatabricks APIを実行するツールを作成してみます。

まず、Databricks APIを利用するための"接続"を作成します。
カタログエクスプローラ上からも作成できるのですが、こんｋSQLで作成します。

%sql
CREATE CONNECTION databricks_api TYPE HTTP
OPTIONS (
  host 'https://xxx.databricks.com', --Workspace APIのホストURL
  port '443',
  base_path '/api/',
  bearer_token secret ('シークレットのスコープ','シークレットのキー') -- Databricks APIを実行する際に利用するAPIキーをシークレットから取得
)
COMMENT 'Databricks API connection.'

次に、この"接続"を使ってクラスタの一覧を取得するAPIを実行するUnity Catalog関数を作成します。

%sql
CREATE OR REPLACE FUNCTION training.llm.list_clusters()
RETURNS STRING
COMMENT 'List Databricks clusters in the workspace.'
RETURN (SELECT http_request(
  conn => 'databricks_api',
  method => 'GET',
  path => '2.1/clusters/list?page_size=1'
)).text

これで、Databricks APIを使ってクラスタ一覧を取得するツールが出来ました。

Unity Catalog関数のLangChainを使って動作確認をしてみます。

# Toolの実験
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit
from unitycatalog.ai.core.databricks import DatabricksFunctionClient
import json

client = DatabricksFunctionClient()
func_name = "training.llm.list_clusters"
toolkit = UCFunctionToolkit(function_names=[func_name], client=client)
tool = toolkit.tools[0]

json.loads(tool.invoke({}))
出力
{'format': 'SCALAR',
 'value': '{"clusters":[{"cluster_id":...

全ての出力結果を掲載していませんが、実行結果として1件のクラスタ情報が取得できます。

その6: 作ったツールを使ってみる

それでは、作ったツールを使って簡単なエージェント処理を作成します。

まず、Unity Catalog関数をLangChainのツールとして取得。

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

func_names = [
    f"{CATALOG}.{SCHEMA}.add_numbers",
    f"{CATALOG}.{SCHEMA}.lookup_customer_info",
    f"{CATALOG}.{SCHEMA}.sample_docs_vector_search",
    f"{CATALOG}.{SCHEMA}.list_clusters",
    "system.ai.python_exec",
]

# UC関数を使用してツールキットを作成
toolkit = UCFunctionToolkit(function_names=func_names, client=client)
tools = toolkit.tools

LangChainのLLMにツールをバインドしてクエリを実行。

from databricks_langchain import ChatDatabricks

# リトリーバーツールを選択したLangchain LLMにバインド
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")
llm_with_tools = llm.bind_tools(tools)

# ツール呼び出し機能をテストするためにLLMとチャット
llm_with_tools.invoke("150+1234443はいくつ？")
出力結果
AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_f8081074-7ca5-4b0f-acd9-a24fe5f84512', 'type': 'function', 'function': {'name': 'training__llm__add_numbers', 'arguments': '{ "number_1": 150, "number_2": 1234443 }'}}]}, response_metadata={'prompt_tokens': 1397, 'completion_tokens': 32, 'total_tokens': 1429}, id='run-2bcdce45-eefc-4734-9ebd-c22ee3b27a3a-0', tool_calls=[{'name': 'training__llm__add_numbers', 'args': {'number_1': 150, 'number_2': 1234443}, 'id': 'call_f8081074-7ca5-4b0f-acd9-a24fe5f84512', 'type': 'tool_call'}])

今回のクエリではadd_numbersのツールが選択されました。
あとは得られたパラメータを基にツールを実行し、その結果を基のクエリと合わせてLLMに渡すと単純なエージェント処理の完成です。（今回はそこまでやらない）

また、Databricks上で試す分にはPlaygroundを使うのが簡単です。

image.png

まとめ

Databricksのドキュメントを基に、いくつかのバリエーションでツール（＝Unity Catalog関数)を作成してみました。
かなり簡単にツール作成ができるようになっていると感じています。

AIエージェント、いろいろ作っていこう。

4
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
4
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?