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

# 導入

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

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

https://docs.databricks.com/ja/generative-ai/agent-framework/ai-agents.html

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

https://docs.databricks.com/ja/generative-ai/agent-framework/agent-tool.html

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

# 準備

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

```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関数として登録しています。(カタログ・スキーマの場所は適当です)

```python
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](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/fa929cae-dc24-9559-b119-c5bad657b3f2.png)

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

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/8072a41b-2359-0b5e-dd19-4b7c9b31523b.png)

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

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

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

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

[こちらのドキュメント](https://docs.databricks.com/ja/generative-ai/agent-framework/structured-retrieval-tools.html)より。

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

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

```sql
%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](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/b2aabf56-1029-9c90-0015-965a84ed75d5.png)

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

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

[こちらのドキュメント](https://docs.databricks.com/ja/generative-ai/agent-framework/unstructured-retrieval-tools.html)より。

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

```sql
%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: 外部サービス接続ツール

[こちらのドキュメント](https://docs.databricks.com/ja/generative-ai/agent-framework/external-connection-tools.html)より。

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/0cb578c7-b1bd-b0a6-8d07-cce796e82450.png)

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

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

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

```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
%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を使って動作確認をしてみます。

```python
# 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({}))
```

```txt:出力
{'format': 'SCALAR',
 'value': '{"clusters":[{"cluster_id":...
```

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

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

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

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

```python
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にツールをバインドしてクエリを実行。

```python
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はいくつ？")
```

```txt:出力結果
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](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/fac7a985-17ef-c89e-7127-6be136fdf69e.png)

# まとめ

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

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