はじめに
先日開催されたJEDAI in Osaka2025にて、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
今回はこの新機能を試し、Databricksの公式ドキュメント（日本語版）の内容に基づいて質問に回答してくれるAgentを作成する手順をご紹介します。
なお、この記事を作成するにあたり、以下の記事を参考にさせていただきました。※というかほぼ二番煎じです...。
Databricks AI Playground で Vector Search を Tool として利用できるようになったので試してみた
公式doc Agentの作成手順
大まかな手順は以下の通りです。
- Embeddingモデルのサービングエンドポイント作成: Vector Search Indexで利用するEmbeddingモデルをAPI経由で利用可能にします。
- Vector Search Indexの作成: Databricksの公式ドキュメントを基に、Vector Search Indexを構築します。
- AI PlaygroundでのAgent作成と動作確認: 作成したVector Search Indexをツールとして利用するAgentをAI Playground上で設定し、動作を確認します。
- Agentのデプロイ: 作成したAgentをサービングエンドポイントとしてデプロイします。
各ステップについて説明します。
Step1: Embeddingモデルのサービングエンドポイント作成
まず、Vector Search Indexを作成するために必要なEmbeddingモデルのサービングエンドポイントを用意します。
Databricks上でpay-per-token形式で提供されているEmbeddingモデルは、あまり日本語が得意ではありません。そこで今回は、Azure OpenAI Serviceのtext-embedding-3-largeモデルを外部モデルとしてDatabricksに登録し、利用することにしました。
Step2: Vector Search Indexの作成
次に、Databricksの公式ドキュメント（日本語版）をデータソースとしたVector Search Indexを作成します。
基本的な手順は、以下の記事を参考に進めますが、一部変更点があります。
参考: 【RAG】Databricksドキュメントを読み込ませてみる
dbdemosからサンプルノートブックのダウンロード
Databricksが提供するデモ用リソース管理ツールdbdemosを使用して、RAGアプリケーションのサンプルノートブックをダウンロードします。このノートブックには、公式ドキュメントを使ってVector Search Indexを作成するコードが含まれています。
%pip install dbdemos
import dbdemos
dbdemos.install('llm-rag-chatbot', catalog='<catalog_name>', schema='<schema_name>')
上記のコードを実行すると、Databricksワークスペース内にllm-rag-chatbotというフォルダが作成されます。
設定ファイルの修正
ダウンロードしたllm-rag-chatbot/configノートブックを開き、設定を一部修正します。
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のカタログ名とスキーマ名に合わせて設定します。
日本語版サイトマップのURLが正しく日本語記事を参照するようになっているため、参考記事にあったようなURL内の言語コード置換処理は不要になりました。
データ準備とインデックス作成ノートブックの実行
次に、llm-rag-chatbot/02-simple-app/01-Data-Preparation-and-Indexノートブックを実行します。
このノートブックでは、以下の処理が行われます。
- 設定ファイルで指定したサイトマップURLに基づき、Databricks公式ドキュメントのコンテンツを取得します。
- 取得したコンテンツをDelta Lakeテーブルに構造化して保存します。
- ドキュメントを適切な長さのチャンクに分割します。
- Step1で作成したEmbeddingモデルのサービングエンドポイントを使用して、各チャンクをベクトル化します。
- ベクトル化されたデータを用いてVector Search Indexを作成します。
-
Embeddingモデルの指定: ノートブック内でEmbeddingモデルを指定する箇所があります。Step1で作成した
text-embedding-3-largeのサービングエンドポイント名に変更してください。 -
クラスター: サーバーレスコンピュート環境で実行した場合、
pandas_udfの処理中にメモリ容量に関するエラーが発生することがありました。そのため、今回は従来のクラスターを使用して実行しました。 - 処理時間: ドキュメントの量にもよりますが、処理完了までには20〜30分程度の時間がかかる場合があります。
Step3: AI PlaygroundでのAgent作成と動作確認
Vector Search Indexの準備ができましたので、AI PlaygroundでAgentを作成します。
- Databricksワークスペースの左メニューから「AI Playground」を開きます。
- 使用するLLM（大規模言語モデル）を選択します。今回は
Claude 3.7 Sonnetを使用しました。 - 「ツール」セクションを開くと、「ベクトル検索」という新しいオプションが追加されています。
- 「ベクトル検索を追加」をクリックし、Step2で作成したVector Search Indexを選択します。
これで設定は完了です。チャット入力欄にDatabricksに関する質問を入力してみましょう。Agentは、設定したVector Search Indexを利用して関連情報を検索し、回答を生成します。
「トレースを表示」をクリックすると、MLflow Tracing（AIアプリケーションの実行フローを追跡・可視化する機能）の結果を確認できます。
トレース結果を見ると、LLMがユーザーの質問に応じてVector Searchツールを呼び出し、関連情報を取得していることがわかります。また、一度の検索だけでなく、必要に応じて複数回の検索を行い、より精度の高い回答を生成しようとしている様子も確認できます。
Step4: Agentのデプロイ
AI Playgroundで作成したAgentは、サービングエンドポイントとしてデプロイし、他のアプリケーションからAPI経由で利用することができます。
AI Playgroundの右上にある「エクスポート」ボタンをクリックすることでデプロイ用のノートブックが生成されます。
自動認証パススルーの設定
生成されたデプロイ用ノートブックを実行する前に、自動認証パススルーの設定を確認・修正する必要があります。
自動認証パススルーは、デプロイされたAgentが、Vector Search IndexやEmbeddingモデルのエンドポイントといった他のDatabricksリソースにアクセスする際に、認証を行うための仕組みです。
参考: 自動認証パススルーのシステム認証
最も一般的な 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モデルのエンドポイント。
embeddingモデルのサービングエンドポイントの追加も必要な点に注意
# 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上で「カスタムエージェント」として選択し、利用することも可能です。
おわりに
これまで、Vector Search Indexを利用するには、データ取得ロジックを実装したUnity Catalog Function（UC関数）を別途定義する必要がありました。
今回のアップデートにより、作成済みのVector Search Indexをツールとして簡単に追加できるようになり、更にシームレスに機能を利用できるようになりました。
また、RAGの中身についても規定回数のRetrievalを行うのではなく、LLMの判断に基づいて複数回のretrievalを実行するフローとして実装されている点や、tracing機能によって検索クエリや、参照されたドキュメントチャンクを簡単に確認できる点も、非常に便利な機能だと感じました。
Databricksでは、Genie Agentや、今回のようにVector Search Indexのように自前のデータを活用する機能をツールとしてシームレスにAgentへ組み込めるようになってきており、Databricks上でAgent開発を行う大きな利点と感じています。