次の方法で共有


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

重要

この機能はパブリック プレビュー段階にあります。

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

要件

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

コードベースのログ記録

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

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

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

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

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

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

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

以下の LangChainPyFunc の例では、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 マネージド リソースを宣言します。 詳細については、「自動認証パススルーのリソースを指定する」を参照してください。

  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 を実行し、MLflow モデルに結果をログに記録します。

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

    resources パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービス エンドポイントなど、エージェントにサービスを提供するために必要な 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-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)

次のステップ