Databricks AI Playgroundの新機能！Vector Search IndexをAgentツールとして使ってみた

## はじめに

先日開催された[JEDAI in Osaka2025](https://jedai.connpass.com/event/349457/)にて、DatabricksのAI Playgroundを用いたAgent作成について発表させていただきました。

AI Playgroundは、Databricks上で大規模言語モデル（LLM）を試したり、Agentのプロトタイプを作成したりできるインタラクティブな環境です。
この度、AI Playgroundの新機能として、Agentが利用できるツールに**Vector Search Index**が追加されたようです。

参考: [Use Vector Search indexes to prototype tool-calling generative AI agents in AI Playground](https://docs.databricks.com/aws/en/release-notes/product/2025/april#use-vector-search-indexes-to-prototype-tool-calling-generative-ai-agents-in-ai-playground)

今回はこの新機能を試し、Databricksの公式ドキュメント（日本語版）の内容に基づいて質問に回答してくれるAgentを作成する手順をご紹介します。

なお、この記事を作成するにあたり、以下の記事を参考にさせていただきました。※というかほぼ二番煎じです...。
[Databricks AI Playground で Vector Search を Tool として利用できるようになったので試してみた](https://qiita.com/taka_yayoi/items/6c1fd00122033dd8ba0e)

## 公式doc Agentの作成手順

大まかな手順は以下の通りです。

1.  **Embeddingモデルのサービングエンドポイント作成**: Vector Search Indexで利用するEmbeddingモデルをAPI経由で利用可能にします。
2.  **Vector Search Indexの作成**: Databricksの公式ドキュメントを基に、Vector Search Indexを構築します。
3.  **AI PlaygroundでのAgent作成と動作確認**: 作成したVector Search Indexをツールとして利用するAgentをAI Playground上で設定し、動作を確認します。
4.  **Agentのデプロイ**: 作成したAgentをサービングエンドポイントとしてデプロイします。

各ステップについて説明します。

## Step1: Embeddingモデルのサービングエンドポイント作成

まず、Vector Search Indexを作成するために必要なEmbeddingモデルのサービングエンドポイントを用意します。

Databricks上でpay-per-token形式で提供されているEmbeddingモデルは、あまり日本語が得意ではありません。そこで今回は、Azure OpenAI Serviceの`text-embedding-3-large`モデルを外部モデルとしてDatabricksに登録し、利用することにしました。

![Pasted image 20250502041851.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/291acc34-c6d1-4def-92db-a356cf2771b1.png)


## Step2: Vector Search Indexの作成

次に、Databricksの公式ドキュメント（日本語版）をデータソースとしたVector Search Indexを作成します。
基本的な手順は、以下の記事を参考に進めますが、一部変更点があります。

参考: [【RAG】Databricksドキュメントを読み込ませてみる](https://qiita.com/taka_yayoi/items/f174599e4721e51e9e1d)

### dbdemosからサンプルノートブックのダウンロード

Databricksが提供するデモ用リソース管理ツール`dbdemos`を使用して、RAGアプリケーションのサンプルノートブックをダウンロードします。このノートブックには、公式ドキュメントを使ってVector Search Indexを作成するコードが含まれています。

```python
%pip install dbdemos
```

```python
import dbdemos

dbdemos.install('llm-rag-chatbot', catalog='<catalog_name>', schema='<schema_name>')
```

上記のコードを実行すると、Databricksワークスペース内に`llm-rag-chatbot`というフォルダが作成されます。

### 設定ファイルの修正

ダウンロードした`llm-rag-chatbot/config`ノートブックを開き、設定を一部修正します。

```python
VECTOR_SEARCH_ENDPOINT_NAME="<endpoint_name>"
# ★ sitemapを日本語のページに
DATABRICKS_SITEMAP_URL = "https://docs.databricks.com/aws/ja/sitemap.xml"
catalog = "<catalog_name>"
#email = spark.sql('select current_user() as user').collect()[0]['user']
#username = email.split('@')[0].replace('.', '_')
#dbName = db = f"dbdemos_rag_{username}"
dbName = db = "<schema_name>"
```

主な変更点は以下の通りです。

*   `DATABRICKS_SITEMAP_URL`: Databricks公式ドキュメントのサイトマップ（ウェブサイトの構造を示すXMLファイル）のURLを、**日本語版**のものに変更します。参考記事の時点からURLが変更されていますのでご注意ください。
*   `catalog`, `dbName`: ご自身のUnity Catalogのカタログ名とスキーマ名に合わせて設定します。
:::note warn
日本語版サイトマップのURLが正しく日本語記事を参照するようになっているため、参考記事にあったようなURL内の言語コード置換処理は不要になりました。
:::

### データ準備とインデックス作成ノートブックの実行

次に、`llm-rag-chatbot/02-simple-app/01-Data-Preparation-and-Index`ノートブックを実行します。

このノートブックでは、以下の処理が行われます。

1.  設定ファイルで指定したサイトマップURLに基づき、Databricks公式ドキュメントのコンテンツを取得します。
2.  取得したコンテンツをDelta Lakeテーブルに構造化して保存します。
3.  ドキュメントを適切な長さのチャンクに分割します。
4.  Step1で作成したEmbeddingモデルのサービングエンドポイントを使用して、各チャンクをベクトル化します。
5.  ベクトル化されたデータを用いてVector Search Indexを作成します。

:::note warn

*   **Embeddingモデルの指定**: ノートブック内でEmbeddingモデルを指定する箇所があります。Step1で作成した`text-embedding-3-large`のサービングエンドポイント名に変更してください。
*   **クラスター**: サーバーレスコンピュート環境で実行した場合、`pandas_udf`の処理中にメモリ容量に関するエラーが発生することがありました。そのため、今回は従来のクラスターを使用して実行しました。
*   **処理時間**: ドキュメントの量にもよりますが、処理完了までには20〜30分程度の時間がかかる場合があります。
:::

## Step3: AI PlaygroundでのAgent作成と動作確認

Vector Search Indexの準備ができましたので、AI PlaygroundでAgentを作成します。

1.  Databricksワークスペースの左メニューから「AI Playground」を開きます。
2.  使用するLLM（大規模言語モデル）を選択します。今回は`Claude 3.7 Sonnet`を使用しました。
3.  「ツール」セクションを開くと、「ベクトル検索」という新しいオプションが追加されています。
4.  「ベクトル検索を追加」をクリックし、Step2で作成したVector Search Indexを選択します。

![Pasted image 20250502045938.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/6d073a1a-874e-4e05-997f-247535d737f5.png)

これで設定は完了です。チャット入力欄にDatabricksに関する質問を入力してみましょう。Agentは、設定したVector Search Indexを利用して関連情報を検索し、回答を生成します。

![Pasted image 20250502050214.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/f1857524-483a-4ae6-922d-7a1b5acb0ab3.png)

「トレースを表示」をクリックすると、MLflow Tracing（AIアプリケーションの実行フローを追跡・可視化する機能）の結果を確認できます。

![Pasted image 20250502050428.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/57dad39a-6e57-4e78-9a79-d1d22c658227.png)
![Pasted image 20250502050455.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/815966f2-c066-40dd-b8a7-764aeab0811b.png)
![Pasted image 20250502050531.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/2020ff04-a97e-437f-871c-eb8874015de3.png)
![Pasted image 20250502050546.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/f1e9ee8f-ef55-47a1-baef-1a987b7cd9f1.png)


トレース結果を見ると、LLMがユーザーの質問に応じてVector Searchツールを呼び出し、関連情報を取得していることがわかります。また、一度の検索だけでなく、必要に応じて複数回の検索を行い、より精度の高い回答を生成しようとしている様子も確認できます。

## Step4: Agentのデプロイ

AI Playgroundで作成したAgentは、サービングエンドポイントとしてデプロイし、他のアプリケーションからAPI経由で利用することができます。
AI Playgroundの右上にある「エクスポート」ボタンをクリックすることでデプロイ用のノートブックが生成されます。


### 自動認証パススルーの設定

生成されたデプロイ用ノートブックを実行する前に、**自動認証パススルー**の設定を確認・修正する必要があります。

自動認証パススルーは、デプロイされたAgentが、Vector Search IndexやEmbeddingモデルのエンドポイントといった他のDatabricksリソースにアクセスする際に、認証を行うための仕組みです。
参考: [自動認証パススルーのシステム認証](https://learn.microsoft.com/ja-jp/azure/databricks/generative-ai/agent-framework/log-agent#-specify-resources-for-automatic-authentication-passthrough-system-authentication)

> 最も一般的な Databricks リソースの種類については、Databricks では、ログ記録中にエージェントのリソース依存関係を事前に宣言することをサポートし、推奨しています。 これにより、エージェントのデプロイ時 [自動認証パススルー] が可能になります。Databricks は、有効期間の短い資格情報を自動的にプロビジョニング、ローテーション、管理して、エージェント エンドポイント内からこれらのリソースの依存関係に安全にアクセスします。自動認証パススルーを有効にするには、次のコードに示すように、log_model() API の `resources` パラメーターを使用して依存リソースを指定します。

エクスポートされたノートブックのデフォルト状態では、`resources`パラメータにはAgentが使用するLLMのエンドポイントしか含まれていません。このままデプロイすると、AgentはVector Search Indexにアクセスするための認証ができず、エラーが発生します。

そこで、`mlflow.pyfunc.log_model()` を呼び出している箇所を探し、`resources`リストに以下のリソースを追加する必要があります。

*   **Vector Search Index**: Step2で作成したVector Search Index。
*   **Embeddingモデルのサービングエンドポイント**: Step1で作成したEmbeddingモデルのエンドポイント。

:::note warn
embeddingモデルのサービングエンドポイントの追加も必要な点に注意
:::

```python
# Determine Databricks resources to specify for automatic auth passthrough at deployment time

import mlflow
from agent import tools, LLM_ENDPOINT_NAME
from databricks_langchain import VectorSearchRetrieverTool
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint, DatabricksVectorSearchIndex
from unitycatalog.ai.langchain.toolkit import UnityCatalogTool

# TODO: Manually include underlying resources if needed. See the TODO in the markdown above for more information.
resources = [
    DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME),
    DatabricksServingEndpoint(endpoint_name='<emb_endpoint_name>'), # 追記
    DatabricksVectorSearchIndex(index_name="<vs_index_name>") # 追記
]

# ==== 中略 ====

with mlflow.start_run():
	logged_agent_info = mlflow.pyfunc.log_model(
	artifact_path="agent",
	python_model="agent.py",
	input_example=input_example,
	resources=resources,
	extra_pip_requirements=[
	"databricks-connect"
	]
)
```

`resources`リストを修正したら、ノートブック全体を実行します。これにより、AgentがMLflowモデルとして登録され、サービングエンドポイントとしてデプロイされます。

デプロイが完了すると、Databricksの「サービング」メニューからエンドポイントの状態を確認できます。また、デプロイされたAgentは、AI Playground上で「カスタムエージェント」として選択し、利用することも可能です。

![Pasted image 20250502111116.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/160102/e98c8431-2ef8-44f3-9616-fa87e18e6183.png)

## おわりに

これまで、Vector Search Indexを利用するには、データ取得ロジックを実装したUnity Catalog Function（UC関数）を別途定義する必要がありました。

今回のアップデートにより、作成済みのVector Search Indexをツールとして簡単に追加できるようになり、更にシームレスに機能を利用できるようになりました。

また、RAGの中身についても規定回数のRetrievalを行うのではなく、LLMの判断に基づいて複数回のretrievalを実行するフローとして実装されている点や、tracing機能によって検索クエリや、参照されたドキュメントチャンクを簡単に確認できる点も、非常に便利な機能だと感じました。

Databricksでは、Genie Agentや、今回のようにVector Search Indexのように自前のデータを活用する機能をツールとしてシームレスにAgentへ組み込めるようになってきており、Databricks上でAgent開発を行う大きな利点と感じています。