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

プレビュー

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

要件

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

コードベースのログ記録

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

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

このアプローチをAPIs mlflow.models.predict() などのデプロイ前検証と組み合わせて使用することで、エージェントが配信のためにデプロイされたときに確実に実行されるようにすることができます。

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

ログ記録中のモデル署名の推論

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

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

以下の LangChain と PyFunc の例では、モデル署名の推論を使用しています。

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

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 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と呼ばれます）。
ドライバーノートブックで、次のコードを使用して 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 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エージェントは、タスクを完了するために他のリソースに対して認証を必要とすることがよくあります。たとえば、エージェントは、非構造化データをクエリするためにベクトル検索インデックスにアクセスする必要がある場合があります。

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