Databricks AI Playgroundの新機能！Vector Search IndexをAgentツールとして使ってみた #AIエージェント

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

自動認証パススルーは、デプロイされた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モデルのエンドポイント。

# 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=''), # 追記
    DatabricksVectorSearchIndex(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上で「カスタムエージェント」として選択し、利用することも可能です。

おわりに