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

[アーティクル]
10/30/2024

重要

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

要件

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

コードベースとシリアル化ベースのログ

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

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

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

各方法の長所と短所を表に示します。

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

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

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

コードを含むノートブックまたは Python ファイルを作成します。この例では、ノートブックまたはファイルの名前は chain.py です。ノートブックまたはファイルには、ここでは lc_chain と呼ばれる LangChain チェーンが含まれている必要があります。
ノートブックまたはファイルに mlflow.models.set_model(lc_chain) を含めます。
ドライバーノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。
ドライバーノートブックで、mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”) を使用して chain.py を実行し、結果を MLflow モデルにログします。
モデルをデプロイします。「生成 AI アプリケーション用にエージェントをデプロイする」を参照してください。エージェントのデプロイは、ベクトル検索インデックスやモデル提供エンドポイントなどの他の Databricks リソースに依存する場合があります。 LangChain エージェントの場合:
- MLflow log_model は、チェーンに必要な依存関係を推論し、ログされたモデル成果物の MLmodel ファイルにログします。 Mlflow バージョン 2.17.0 以降では、これらの推論された依存関係をオーバーライドできます。 PyFunc エージェントまたは LangChain エージェント Specify リソースを参照してください。
- デプロイ時に、databricks.agents.deploy はこれらの推測リソースの依存関係にアクセスして通信するために必要な M2M OAuth トークンを自動的に作成します。
サービス環境が読み込まれると、chain.py が実行されます。
サービス要求が入ると、lc_chain.invoke(...) が呼び出されます。


import mlflow

code_path = "/Workspace/Users/first.last/chain.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_chain_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="chain", # 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 chain
    example_no_conversion=True, # Required
  )

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

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

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

コードを含むノートブックまたは Python ファイルを作成します。この例では、ノートブックまたはファイルの名前は chain.py です。ノートブックまたはファイルには、ここでは PyFuncClass と呼ばれる PyFunc クラスが含まれている必要があります。
ノートブックまたはファイルに mlflow.models.set_model(PyFuncClass) を含めます。
ドライバーノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。
ドライバーノートブックで、mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”) を使用して chain.py を実行し、結果を MLflow モデルにログします。 resources パラメーターは、ベクター検索インデックスや基盤モデルを提供するサービスエンドポイントなど、モデルの提供に必要なリソースを宣言します。 PyFunc のリソースファイルの例については、PyFunc エージェントまたは LangChain エージェント Specify リソースを参照してください。
モデルをデプロイします。「生成 AI アプリケーション用にエージェントをデプロイする」を参照してください。
サービス環境が読み込まれると、chain.py が実行されます。
サービス要求が入ると、PyFuncClass.predict(...) が呼び出されます。

import mlflow

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

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using PyFunc model

resources_path = "/Workspace/Users/first.last/resources.yml"

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
  )

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

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

PyFunc または LangChain エージェントのリソースを指定する

モデルを提供するために必要なリソース (ベクトル検索インデックスや提供エンドポイントなど) を指定できます。

LangChain の場合、リソースは自動的に検出され、ベストエフォートアプローチを使用してモデルでログに記録されます。 MLflow バージョン 2.17.0 以降では、次に示すようなコードを使用して、自動的に推論されたリソースをオーバーライドできます。これは、エージェントが必要な依存関係と共にログに記録されるようにするため、運用環境のユースケースに推奨されます。

pyfunc でフレーバーされたエージェントをデプロイするときは、デプロイされたエージェントのリソース依存関係を手動で追加する必要があります。 resources パラメーターで指定されたすべてのリソースにアクセスできる M2M OAuth トークンが作成され、デプロイされたエージェントに提供されます。

Note

チェーンのログ時にリソースを手動で指定することで、エンドポイントがアクセス許可を持つリソースをオーバーライドできます。

次のコードでは、 resources パラメーターを使用して依存関係を指定します。

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

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
    ]
  )

resources.yaml ファイルでリソースを指定しても、追加できます。そのファイルパスは、resources パラメーターで参照します。 resources.yaml 内のすべての指定されたリソースにアクセスできる M2M OAuth トークンが作成され、デプロイされたエージェントに提供されます。

モデル提供エンドポイントとベクトル検索インデックスを定義する resources.yaml ファイルの例を次に示します。


api_version: "1"
databricks:
  vector_search_index:
    - name: "catalog.schema.my_vs_index"
  serving_endpoint:
    - name: databricks-dbrx-instruct
    - name: databricks-bge-large-en

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 = "chain_name"

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

次の方法で共有

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

要件

コードベースとシリアル化ベースのログ

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

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

PyFunc または LangChain エージェントのリソースを指定する

Unity Catalog にチェーンを登録する

次のステップ

フィードバック

その他のリソース