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によるコードベースのロギング

あなたのコードでノートブックまたは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エージェントのリソースの指定を参照してください。
- デプロイ中に、 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 エージェントのリソースの指定」を参照してください。
モデルをデプロイします。「生成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 トークンが作成され、デプロイされたエージェントに提供されます。

注：

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

次のコードでは、 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)