AIエージェントの記録と登録

プレビュー

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

要件

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

コードベースのログ記録とシリアライズベースのログ記録

コードベースの MLflow ログ記録またはシリアル化ベースの MLflow ログ記録を使用できます。 Databricks では、コードベースのログ記録を使用することをお勧めします。

コードベースの MLflow ログ記録: エージェントのコードは Python ファイルとしてキャプチャされます。 Python 環境は、パッケージのリストとしてキャプチャされます。エージェントがデプロイされると、Python環境が復元され、エージェントのコードが実行されてエージェントがメモリにロードされ、エンドポイントが呼び出されたときにエージェントを呼び出せるようになります。

シリアル化ベースの MLflow ログ記録: Python 環境でのエージェントのコードと現在の状態は、多くの場合、pickle や joblibなどのライブラリを使用してディスクにシリアル化されます。エージェントがデプロイされると、Python環境が復元され、シリアル化されたオブジェクトがメモリにロードされるため、エンドポイントが呼び出されたときに呼び出すことができます。

表は、各方法の長所と短所を示しています。

メソッド	利点	デメリット
コードベースのMLflowログ記録	多くの一般的な生成AIライブラリでサポートされていないシリアル化の固有の制限を克服します。後で参照できるように、元のコードのコピーを保存します。コードをシリアル化できる単一のオブジェクトに再構築する必要はありません。	`log_model(...)` エージェントのコード (ドライバーノートブックと呼ばれます) とは異なるノートブックから呼び出す必要があります。
シリアル化ベースのMLflowログ記録	`log_model(...)` モデルが定義されている同じノートブックから呼び出すことができます。	元のコードは使用できません。エージェントで使用されるすべてのライブラリとオブジェクトは、シリアル化をサポートしている必要があります。

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

LangChainによるコードベースのロギング

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

コードを使用してノートブックまたは Python ファイルを作成します。この例では、ノートブックまたはファイルの名前は agent.pyです。ノートブックまたはファイルには、ここでは LangChain エージェント ( lc_agentと呼びます) が含まれている必要があります。
ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。
ドライバーノートブックとして機能する新しいノートブックを作成します（この例ではdriver.pyと呼ばれます）。
ドライバーノートブックで、次のコードを使用して ML agent.py を実行し、結果を MLflow モデルに記録します。
```
mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
resources パラメーターは、エージェントにサービスを提供する Databricksために必要な -managed リソース (ベクトル検索インデックスや、基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、「自動認証パススルーのリソースを指定する」を参照してください。
モデルをデプロイします。「生成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 = {
    "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 is the configuration that is used for training the model. 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 your 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と呼ばれます）。
ドライバーノートブックで、次のコードを使用して ML agent.py を実行し、結果を MLflow モデルに記録します。
```
mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
```
resources パラメーターは、エージェントにサービスを提供する Databricksために必要な -managed リソース (ベクトル検索インデックスや、基盤モデルを提供する提供エンドポイントなど) を宣言します。詳細については、「自動認証パススルーのリソースを指定する」を参照してください。
モデルをデプロイします。「生成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 = {
    "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エージェントは、タスクを完了するために他のリソースに対して認証を必要とすることがよくあります。たとえば、エージェントは、非構造化データをクエリするためにベクトル検索インデックスにアクセスする必要がある場合があります。

「依存リソースの認証」で説明されているように、モデルサービングは、エージェントをデプロイするときに、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 を手動で指定することをお勧めします。

注：

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

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

リソースタイプ	リソースをログに記録するために必要な最小 `mlflow` バージョン
ベクトル検索インデックス	`mlflow` 2.13.1 以降が必要です
モデルサービングendpoint	`mlflow` 2.13.1 以降が必要です
SQLウェアハウス	`mlflow` 2.16.1以降が必要です
Unity Catalog 機能	`mlflow` 2.16.1以降が必要です
Genieスペース	`mlflow` 2.17.1 以降が必要です
Unity Catalog テーブル	`mlflow` 2.18.0 以降が必要です

エージェントを Unity Catalog に登録する

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

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)