ユーザーの代理認証を使用してエージェントをデプロイする

ユーザーの代理認証を使用すると、Mosaic AI モデルサービスを介してデプロイされたエージェントは、エージェントにクエリを実行した Databricks エンド ユーザーの ID を使用して Databricks リソースにアクセスできます。 これにより、ユーザーごとに機密情報にアクセスでき、Unity カタログでデータ アクセス制御をきめ細かく適用できます。

ユーザーの代理認証では、ダウンスコープによって受信ユーザー トークンがさらに制限されるため、エージェント コードに公開されるトークンが、エージェントの作成者によって定義された特定の API にのみアクセスするように制限されます。 これにより、承認されていないアクションを防ぎ、トークンの誤用のリスクを軽減することで、セキュリティが向上します。

エージェントを作成するときに、既存の SDK を引き続き使用して、ベクター検索インデックスなどの Databricks リソースにアクセスできます。 リソースへのユーザーの代理アクセスを有効にするには:

1. エージェント コードで、エージェント エンド ユーザーの代わりにリソースにアクセスする必要があることを示す SDK 呼び出しを更新します
2. エージェントのログ記録時 (エージェントのデプロイ前) に、エージェントに必要なエンド ユーザー REST API スコープを指定します。 詳細については、ユーザーの代理認証に関するページを参照してください。

ユーザーの代理認証のエンド ツー エンドの例については、 エンド ツー エンドの例を参照してください。

次のリソースは、エージェントを使用したユーザーの代理認証と互換性があります。

ユーザー代理クライアントを使用してツールを初期化する際は、ツールの初期化をtry-exceptブロックで囲むか、エラーをそのままユーザーに通知することができます。 エラーを処理することで、エンド ユーザーが必要なすべてのツールにアクセスできない場合でも、エージェントはベスト エフォート応答を実行できます。 ただし、ツールの初期化中にエラーを処理しないことを選択した場合、ユーザーに必要なリソースがない場合、エージェントはエラーをスローします。

SDK の構成

次のスニペットは、さまざまな SDK を使用して、さまざまな Databricks リソースへのユーザーの代理アクセスを構成する方法を示しています

ベクター検索リトリーバーツール

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials
from databricks_langchain import VectorSearchRetrieverTool

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy = ModelServingUserCredentials())

vector_search_tools = []
# Exclude exception handling if you want the agent to fail
# when users lack access to all required Databricks resources
try:
  tool = VectorSearchRetrieverTool(
    index_name="<index_name>",
    description="...",
    tool_name="...",
    workspace_client=user_client # Specify the user authenticated client
    )
    vector_search_tools.append(tool)
except Exception as e:
    _logger.debug("Skipping adding tool as user does not have permissions)

ベクター検索クライアント

from databricks.vector_search.client import VectorSearchClient
from databricks.vector_search.utils import CredentialStrategy

# Configure a VectorSearch Client to use on behalf of end
# user authentication
user_authenticated_vsc = VectorSearchClient(credential_strategy=CredentialStrategy.MODEL_SERVING_USER_CREDENTIALS)
# Exclude exception handling if you want the agent to fail when
# users lack access to all required Databricks resources
try:
  vs_index = user_authenticated_vsc.get_index(endpoint_name="endpoint_name", index_name="index_name")
  ...
except Exception as e:
  _logger.debug("Skipping Vector Index because user does not have permissions)

モデル サービング エンドポイント

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

# Exclude exception handling if you want the agent to fail
# when users lack access to all required Databricks resources
try:
  user_client.serving_endpoints.query("endpoint_name", input="")
except Exception as e:
  _logger.debug("Skipping Model Serving Endpoint due to no permissions")

UC 接続

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

user_client.serving_endpoints.http_request(
  conn="connection_name",
  method=ExternalFunctionRequestHttpMethod.POST,
  path="/api/v1/resource",
  json={"key": "value"},
  headers={"extra_header_key": "extra_header_value"},
)

ジーニー スペース

from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials
from databricks_langchain.genie import GenieAgent

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy=ModelServingUserCredentials())

genie_agent = GenieAgent(
    genie_space_id="<genie_space_id>",
    genie_agent_name="Genie",
    description="Genie_description",
    client=user_client, # Specify the user client here
  )

エージェントの初期化

ユーザー代理認証は ChatAgent インターフェースと互換性があります。 ユーザーの代理認証を使用する場合、エンド ユーザーの ID は、デプロイされたエージェントのクエリが実行されたとき (つまり、ChatAgent インターフェイスの predict および predict_stream 関数内) にのみ認識されます。 その結果、ChatAgent 実装の __init__ メソッドではなく、これらのメソッド内でリソースへのアクセスをエンドユーザーに代わって行う必要があります（たとえば、エンドユーザーがアクセスできるリストベクター検索インデックスを表示するなど）。 これにより、リソースが呼び出し間で分離されます

from mlflow.pyfunc import ChatAgent


class LangGraphChatAgent(ChatAgent):
  def initialize_agent():
    user_client = WorkspaceClient(
      credentials_strategy=ModelServingUserCredentials()
    )
    system_authorized_client = WorkspaceClient()
    ### Use the clients above to access resources with either system or user authorization

  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    agent = initialize_agent() # Initialize the Agent in Predict
    request = {"messages": self._convert_messages_to_dict(messages)}

    messages = []
    for event in self.agent.stream(request, stream_mode="updates"):
      for node_data in event.values():
        messages.extend(
          ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
        )
    return ChatAgentResponse(messages=messages)

セキュリティに関する考慮事項

エージェントでユーザーの代理認証を有効にする前に、考慮すべきセキュリティ上の影響がいくつかあります。

1. 機密 Databricks リソースへのアクセス: ユーザーの代理認証を有効にすると、エージェントは機密の Databricks リソースにアクセスできます。 開発者がアクセスできるリソースを制限し、トークンの誤用のリスクを軽減するために API スコープを実装しましたが、一部のリスクはまだ残っています。 たとえば、 serving.serving-endpoints API スコープは、ユーザーに代わってサービス エンドポイントを実行するアクセス許可をエージェントに付与します。 ただし、サービス エンドポイント自体は、元のエージェントが使用を許可されていない追加の API スコープにアクセスできる場合があります。
2. エンド ユーザーの同意のサポートなし: 現在のベータ 段階では、エージェント ユーザーはエージェントで必要な Databricks REST API スコープを表示または同意できません。 ユーザーは、サービス エンドポイントに対する "Can Manage" アクセス許可を持つユーザーが Databricks で自分の代わりにアクションを実行することを信頼する責任があります。

エンドツーエンドの例

次のノートブックは、ユーザーの代理認証を使用してベクター検索を使用してエージェントを作成する方法を示しています。

ベクトル検索を使用したユーザー代理認証

ノートブックを入手