AI エージェントをログに記録して登録する

Mosaic AI Agent Framework を使用して AI エージェントをログに記録します。 エージェントのログは、開発プロセスの基礎となります。 ログ記録では、エージェントのコードと構成の "特定の時点" がキャプチャされるため、構成の品質を評価できます。

要件

ログに記録する前に AI エージェントを作成します。

Databricks では、最新バージョンの databricks-sdkをインストールすることをお勧めします。

% pip install databricks-sdk

コード ベースのログ記録

Databricks では、エージェントのログ記録時に MLflow の コードからのモデル機能 を使用することをお勧めします。

この方法では、エージェントのコードが Python ファイルとしてキャプチャされ、Python 環境がパッケージの一覧としてキャプチャされます。 エージェントがデプロイされると、Python 環境が復元され、エージェントのコードが実行され、エージェントがメモリに読み込まれるため、エンドポイントが呼び出されたときに呼び出すことができます。

この方法と、mlflow.models.predict() などのデプロイ前検証 API を使用して、サービスを提供するためにデプロイされたときにエージェントが確実に実行されるようにすることができます。

コードベースのログ記録の例については、ChatAgent の作成例 ノートブックを参照してください。

ログ記録中にモデル署名を推論する

ChatAgent インターフェイスを使用しない場合は、次のいずれかの方法を使用して、ログ記録時にエージェントの MLflow Model Signature を指定する必要があります。

1. 署名を手動で定義する
2. MLflow の Model Signature 推論機能を使用して、指定した入力例に基づいてエージェントの署名を自動的に生成します。 この方法は、署名を手動で定義するよりも便利です。

MLflow モデルシグネチャは、入力と出力を検証して、エージェントが AI Playground やレビュー アプリなどのダウンストリーム ツールと正しく対話することを確認します。 また、エージェントを効果的に使用する方法に関する他のアプリケーションのガイドも示します。

以下の LangChain と PyFunc の例では、Model Signature 推論を使用します。

ログ記録時にモデル署名を自分で明示的に定義する場合は、「MLflow のドキュメント - 署名を使用してモデルをログに記録する方法」を参照してください。

LangChain を使用したコードベースのログ

次の手順とコード サンプルでは、LangChain を使用してエージェントをログに記録する方法を示します。

1. コードを含むノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は agent.pyです。 ノートブックまたはファイルには、ここで lc_agentと呼ばれる LangChain エージェントが含まれている必要があります。

2. ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。

3. ドライバー ノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。

4. ドライバー ノートブックで、次のコードを使用して agent.py を実行し、MLflow モデルに結果をログに記録します。

   mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
   

   resources パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービス エンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージド リソースを宣言します。 詳細については、「 Databricks リソースの認証」を参照してください。

5. モデルをデプロイします。 生成 AI アプリケーションのエージェントのデプロイを参照してください。

6. サービス環境が読み込まれると、agent.py が実行されます。

7. サービス要求が入ると、lc_agent.invoke(...) が呼び出されます。

import mlflow

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

# example using langchain
with mlflow.start_run():
  logged_agent_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
    artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
    input_example=input_example, # Must be a valid input to the agent
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)

PyFunc を使用したコードベースのログ記録

次の手順とコード サンプルでは、PyFunc を使用してエージェントをログに記録する方法を示します。

1. コードを含むノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は agent.pyです。 ノートブックまたはファイルには、PyFuncClassという名前の PyFunc クラスが含まれている必要があります。

2. ノートブックまたはファイルに mlflow.models.set_model(PyFuncClass) を含めます。

3. ドライバー ノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。

4. ドライバー ノートブックで、次のコードを使用して agent.py を実行し、 log_model() を使用して結果を MLflow モデルにログに記録します。

   mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
   

   resources パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービス エンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージド リソースを宣言します。 詳細については、「 Databricks リソースの認証」を参照してください。

5. モデルをデプロイします。 生成 AI アプリケーションのエージェントのデプロイを参照してください。

6. サービス環境が読み込まれると、agent.py が実行されます。

7. サービス要求が入ると、PyFuncClass.predict(...) が呼び出されます。

import mlflow
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex,
)

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
    ]
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)

Databricks リソースの認証

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

「依存リソースの 認証」で説明されているように、Model Serving では、エージェントのデプロイ時に Databricks で管理されるリソースと外部リソースの両方に対する認証がサポートされます。

Model Serving では、Databricks マネージド リソースに対して次の 2 種類の認証がサポートされています。

1. システム認証: エージェント サービス プリンシパルが、エージェントのログ記録時に指定された依存リソースにアクセスできるようにします。 これは、共有リソースや機密性の高いリソース (パブリック ドキュメントを含むベクター検索インデックスなど) にアクセスする場合に便利です。
2. [ベータ] ユーザーの代理認証: エージェントがエンド ユーザーの資格情報を使用して Databricks リソースにアクセスできるようにします。 これは、エージェントが機密データにアクセスしたり、リモート API にクエリを実行してユーザーごとにアクションを実行したりする必要があるシナリオに役立ちます。

自動認証パススルー (システム認証) のリソースを指定する

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

自動認証パススルーを有効にするには、次のコードに示すように、resources API の log_model() パラメーターを使用して依存リソースを指定します。

import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection
)

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    example_no_conversion=True,
    # Specify resources for automatic authentication passthrough
    resources=[
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
      DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
      DatabricksTable(table_name="your_table_name"),
      DatabricksUCConnection(connection_name="your_connection_name"),
    ]
  )

Databricks では、すべてのエージェント フレーバーに resources を手動で指定することをお勧めします。

次の表に、自動認証パススルーをサポートする Databricks リソースと、リソースのログ記録に必要な最小 mlflow バージョンを示します。

ユーザーの代行認証

ユーザーの代理認証を使用するエージェントをログに記録する場合は、エージェント コード内でエンド ユーザーとしてアクションを実行するために必要な Databricks API スコープの最小セットを指定します。 これにより、エージェントは、デプロイ時にエンド ユーザーに代わってアクションを実行するための最小限の特権アクセスが確保され、承認されていないアクションを防ぎ、トークンの誤用のリスクを最小限に抑えることでセキュリティが強化されます。

いくつかの一般的な種類の Databricks リソースにアクセスするために必要なスコープの一覧を次に示します。

ユーザーの代理認証を有効にするには、次の例に示すように、MLflow AuthPolicy を log_model()に渡します。 MLflow AuthPolicy には、次の 2 つのコンポーネントがあります。

• system_auth_policy: システム認証のリソースを指定します。 通常、エージェントは、機密性の高いリソースまたは API にアクセスするためのユーザー代理認証と組み合わせて、共有リソースのシステム認証を使用します (例: モデル サービス エンドポイントにクエリを実行する)。
• user_auth_policy: エージェントがユーザーに代わって認証するために必要な API スコープを指定します

from mlflow.models.resources import DatabricksServingEndpoint
from mlflow.models.auth_policy import SystemAuthPolicy, UserAuthPolicy, AuthPolicy

resources = [
    DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct")
]
# Specify resources here for system authentication
system_auth_policy = SystemAuthPolicy(resources=resources)

# Specify the minimal set of API scopes needed for on-behalf-of-user authentication
	# When deployed, the agent can access Databricks resources and APIs
	# on behalf of the end user, but only via REST APIs that are covered by the list of
	# scopes below

user_auth_policy = UserAuthPolicy(
    api_scopes=[
        "serving.serving-endpoints",
        "vectorsearch.vector-search-endpoints",
        "vectorsearch.vector-search-indexes",
    ]
)

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        ...
        # Instead of passing `resources` (which only supports system authentication),
	        # pass an auth_policy to log_model to enable both system authentication and
	        # on-behalf-of-user authentication
        auth_policy=AuthPolicy(
            system_auth_policy=system_auth_policy,
            user_auth_policy=user_auth_policy
        )
    )

OpenAI クライアントの自動認証

エージェントが OpenAI クライアントを使用している場合は、Databricks SDK を使用してデプロイ時に自動的に認証します。 Databricks SDK には、 get_open_ai_client()自動的に構成された承認を使用して OpenAI クライアントを構築するためのラッパーが用意されています。 ノートブックで次を実行します。

% pip install databricks-sdk[openai]

from databricks.sdk import WorkspaceClient
def openai_client(self):
  w = WorkspaceClient()
  return w.serving_endpoints.get_open_ai_client()

次に、デプロイ時に自動的に認証するために、resources の一部として Model Serving エンドポイントを指定します。

エージェントを Unity カタログに登録する

エージェントをデプロイする前に、エージェントを Unity カタログに登録する必要があります。 エージェントパッケージをモデルとして Unity カタログに登録する。 その結果、エージェント内のリソースの承認に Unity カタログのアクセス許可を使用できます。

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)

mlflow.register_model()を参照してください。

次のステップ

• AI エージェントにトレースを追加します。
• AI エージェントを展開します。