Databricks AIエージェントにおけるリソース認証：ユーザ代理認証を試す

Posted at 2025-04-27

今回試すDatabricksの機能はまだベータ版です。
将来的に動作が大きく変更になる可能性がありますので、ご注意ください。

導入

DatabricksでAIエージェントを作成する際、利用するリソースの認証をどうするかは悩みの種です。
例えばDatabricksのリソース(Unity Catalog上のテーブルやGenie APIなど)や外部APIを利用する際、固定の権限で実行するのではなく、AIエージェントを利用するエンドユーザの権限で行いたい時があります。

Databricksでそれを簡単にしようと思ったらどうするのがいいんだろう、と思ってドキュメントを読んでいたら「ユーザ代理認証」の機能が追加されていました。

以下のドキュメントより、抜粋。

Databricks リソースの認証

AIエージェントは、タスクを完了するために他のリソースに対して認証する必要があることがよくあります。たとえば、エージェントは、非構造化データをクエリするためにベクトル検索インデックスにアクセスする必要がある場合があります。

依存リソースの認証で説明されているように、モデルサービングは、エージェントをデプロイするときに、Databricksマネージドリソースと外部リソースの両方に対する認証をサポートします。

モデルサービングは、 Databricks管理リソースに対して 2 種類の認証をサポートしています。

システム認証: エージェントサービスプリンシパルが、エージェントのログ時に指定された任意の依存リソースにアクセスできるようにします。これは、共有リソースや機密性の低いリソース(公開ドキュメントを含むベクトル検索インデックスなど)にアクセスする場合に便利です

[ベータ版] ユーザー代理認証: エージェントがエンドユーザーの資格情報を使用して Databricks リソースにアクセスできるようにします。これは、エージェントが機密データにアクセスしたり、リモート APIにクエリを実行してユーザーごとにアクションを実行したりする必要があるシナリオに役立ちます

2025年4月現在、まだベータ版のようですが、エンドユーザの資格情報を使ってDatabricksリソースにアクセスできるようです。まさに求めている機能。

実際に挙動を確かめたく、以下の公式サンプルをもとにAIエージェントを作って試してみます。

ただ、そのまま動かしても（個人的に）面白くないので、公式サンプルはVector Searchでの対応をしていますが、今回はGenie Spaceに対応してみます。

検証はDatabricks on AWSで行いました。
ノートブックはサーバレスクラスタで動かしています。

Step1. 準備

本機能はベータ版につき標準だと有効化されていないようなので、ワークスペースのプレビュー機能より有効化してください。

次にノートブックを作成し、AIエージェント作成に必要なパッケージをインストールします。

%pip install -U -qqqq mlflow langgraph==0.3.34 langgraph-supervisor databricks-langchain databricks-agents uv

%restart_python

次にサンプルとしてGenie Spaceと連携するAIエージェント作成します。
（複数Genieスペースに拡張できるようlanggraph-supervisorを利用しているため、少し冗長になっています)

今回はこちらの記事でサンプル用に作成した売上データ探索Genieスペースを再利用しています。

%%writefile obo_agent.py
from typing import Any, Generator, Optional
import uuid

import mlflow
from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials
from databricks_langchain import ChatDatabricks
from databricks_langchain.genie import GenieAgent
from langgraph_supervisor import create_supervisor
from langchain_core.messages import BaseMessage, convert_to_openai_messages

from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
    ChatAgentChunk,
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)

mlflow.langchain.autolog()

###################################################
## Genie Spaceのエージェント作成用のパラメータ
###################################################

GENIE_SPACE_ID = "01f000e7c3141060b962a5b99f980efd" # 売上探索GenieスペースのID
genie_agent_description = "このエージェントは売上情報に関する質問に答えることができます。"

############################################
# LLMエンドポイントとシステムプロンプトを定義
############################################
LLM_ENDPOINT_NAME = "databricks-claude-3-7-sonnet"
llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME)

supervisor_prompt = "あなたは売上などの経営データに関する専門家です。基本的にGenieエージェントを使ってください。"


def create_genie_agents():
    """Genieエージェントを作成する"""

    user_authenticated_client = WorkspaceClient(
        credentials_strategy=ModelServingUserCredentials()
    )
    genie_agents = []
    try:
        genie_agent = GenieAgent(
            genie_space_id=GENIE_SPACE_ID,
            genie_agent_name="Genie",
            description=genie_agent_description,
            client=user_authenticated_client,
        )
        genie_agent.name = "sales_genie_agent"
        genie_agents.append(genie_agent)
    except Exception as e:
        print(f"Genieエージェントツールをスキップ: {e}")

    return genie_agents


def convert_lc_message_to_chat_message(lc_message: BaseMessage) -> ChatAgentMessage:
    """LangChainメッセージをChatAgentMessageに変換する"""
    msg = convert_to_openai_messages(lc_message)
    if not "id" in msg:
        msg.update({"id": str(uuid.uuid4())})

    return ChatAgentMessage(**msg)


class LangGraphChatAgent(ChatAgent):
    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        # ここでpredict呼び出し時にエージェントを初期化
        workflow = create_supervisor(
            create_genie_agents(),
            model=llm,
            prompt=supervisor_prompt,
        ).compile()
        request = {"messages": self._convert_messages_to_dict(messages)}

        messages = []
        for event in workflow.stream(request, stream_mode="updates"):
            for node_data in event.values():
                print(node_data)
                messages.extend(
                    convert_lc_message_to_chat_message(msg)
                    for msg in node_data.get("messages", [])
                )
        return ChatAgentResponse(messages=messages)

    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> Generator[ChatAgentChunk, None, None]:
        workflow = create_supervisor(
            create_genie_agents(),
            model=llm,
            prompt=supervisor_prompt,
        ).compile()

        request = {"messages": self._convert_messages_to_dict(messages)}
        for event in workflow.stream(request, stream_mode="updates"):
            for node_data in event.values():
                yield from (
                    ChatAgentChunk(**{"delta": convert_lc_message_to_chat_message(msg)})
                    for msg in node_data["messages"]
                )


# エージェントオブジェクトを作成し、mlflow.models.set_model()を介して推論のためにエージェントオブジェクトとして指定
AGENT = LangGraphChatAgent()
mlflow.models.set_model(AGENT)

動かすと以下のような結果がMLflow Tracingで得られます。

from obo_agent import AGENT

AGENT.predict({"messages": [{"role": "user", "content": "先月の売上合計を教えて"}]})

Step2. エージェントのロギング

Step1で定義したエージェントをMLflowにロギングします。
ここが今回のポイント。（解説？はコード後）

import mlflow
from obo_agent import LLM_ENDPOINT_NAME
from mlflow.models.resources import DatabricksServingEndpoint
from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
from pkg_resources import get_distribution

resources = [DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)]
systemAuthPolicy = SystemAuthPolicy(resources=resources)

userAuthPolicy = UserAuthPolicy(
    api_scopes=[
        "serving.serving-endpoints",
        "dashboards.genie",
    ]
)
with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        artifact_path="obo_agent",
        python_model="obo_agent.py",
        extra_pip_requirements=[
            f"databricks-connect=={get_distribution('databricks-connect').version}"
        ],
        # auth_policyでSystem Auth PolicyとUser Auth Policyを指定。ここが今回の重要ポイント。
        auth_policy=AuthPolicy(
            system_auth_policy=systemAuthPolicy, user_auth_policy=userAuthPolicy
        ),
    )

簡単な解説ですが、log_modelでモデルを保管する際に、auth_policyという珍しい？パラメータを指定しています。
ここでsystem_auth_policy(従来のresourcesに相当するパラメータ)とuser_auth_policyを指定します。
エンドユーザの資格情報を使う場合、user_auth_policyに適切なapi_scopeを指定する必要があります。
逆に言えば、ここのapi_scopeを設定すれば、エンドユーザの資格情報を使って各APIに対して実行できるようになる、ということです。

以下のドキュメントにも記載があるのですが、利用には若干のリスクや考慮点があります。
注意して使いましょう。

ひとまずここまででエージェントをロギングすることができます。

Step3. デプロイする

Mosaic AI Agent Frameworkを使ってエンドポイントにエージェントをデプロイします。
(本来は事前検証などするべきですが、説明が長くなるので割愛します）

まずは先ほどロギングしたエージェントをUnity Catalogに登録。

mlflow.set_registry_uri("databricks-uc")

catalog = "training"
schema = "llm"
model_name = "obo_sample_agent"
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"

# register the model to UC
uc_registered_model_info = mlflow.register_model(
    model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME
)

次にAgent Frameworkを使ってUnity Catalogに登録したモデルをデプロイ。

from databricks import agents

agents.deploy(
    UC_MODEL_NAME,
    uc_registered_model_info.version,
    tags={"endpointSource": "docs"},
    scale_to_zero=True,
)