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

[アーティクル]
02/04/2025

重要

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

要件

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

コードベースのログ記録

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

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

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

エージェントまたはエージェントをログするコードは、エージェントコードとは別のノートブックに含める必要があります。このノートブックは、ドライバーノートブックと呼ばれます。ノートブックの例については、「サンプルノートブック」を参照してください。

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

ログ記録中に、エージェントの入力スキーマと出力スキーマを指定する MLflow Model Signatureを定義する必要があります。署名は入力と出力を検証して、エージェントが AI Playground やレビューアプリなどのダウンストリームツールと正しく対話することを確認します。また、エージェントを効果的に使用する方法に関する他のアプリケーションのガイドも示します。

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

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

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

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

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

コードを含むノートブックまたは Python ファイルを作成します。この例では、ノートブックまたはファイルの名前は agent.pyです。ノートブックまたはファイルには、ここで lc_agentと呼ばれる LangChain エージェントが含まれている必要があります。
ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。
ドライバーノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。
ドライバーノートブックで、次のコードを使用して agent.py を実行し、MLflow モデルに結果をログに記録します。
```
mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
resources パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービスエンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージドリソースを宣言します。詳細については、「自動認証パススルーのリソースを指定する」を参照してください。
モデルをデプロイします。「生成 AI アプリケーション用にエージェントをデプロイする」を参照してください。
サービス環境が読み込まれると、agent.py が実行されます。
サービス要求が入ると、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 を使用してエージェントをログに記録する方法を示します。

コードを含むノートブックまたは Python ファイルを作成します。この例では、ノートブックまたはファイルの名前は agent.pyです。ノートブックまたはファイルには、PyFuncClassという名前の PyFunc クラスが含まれている必要があります。
ノートブックまたはファイルに mlflow.models.set_model(PyFuncClass) を含めます。
ドライバーノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。
ドライバーノートブックで、次のコードを使用して agent.py を実行し、MLflow モデルに結果をログに記録します。
```
mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
resources パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービスエンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージドリソースを宣言します。詳細については、「自動認証パススルーのリソースを指定する」を参照してください。
モデルをデプロイします。「生成 AI アプリケーション用にエージェントをデプロイする」を参照してください。
サービス環境が読み込まれると、agent.py が実行されます。
サービス要求が入ると、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-mixtral-8x7b-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)

自動認証パススルーのリソースを指定する

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

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

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

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

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

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-mixtral-8x7b-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"),
    ]
  )

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

Note

mlflow.langchain.log_model(...)を使用して LangChain エージェントをログに記録するときにリソースを指定しない場合、MLflow はリソースのベストエフォート自動推論を実行します。ただし、これによりすべての依存関係がキャプチャされるとはかからず、エージェントにサービスを提供したりクエリを実行したりするときに承認エラーが発生する可能性があります。

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

リソースの種類	リソースをログに記録するために必要な最小 `mlflow` バージョン
ベクター検索インデックス	`mlflow` 2.13.1 以降が必要です
モデルサービングエンドポイント	`mlflow` 2.13.1 以降が必要です
SQL Warehouse	`mlflow` 2.16.1 以降が必要です
Unity Catalog 関数	`mlflow` 2.16.1 以降が必要です
ジーニー空間	`mlflow` 2.17.1 以降が必要です
Unity カタログテーブル	`mlflow` 2.18.0 以降が必要です

エージェントを 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)

次の方法で共有