エージェントログに記録して登録する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 ファイルに記録します。
- デプロイ中に、 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 のリソースファイルの例を参照してください。
モデルをデプロイします。「アプリケーションを生成するためのエージェントのデプロイ」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の場合、リソースは自動的に取得され、モデルとともにログに記録されます。

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

注：

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

次に、 resources パラメーターで指定して、サービングエンドポイントとベクトル検索インデックスの依存関係を追加する方法を示します。

 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="rag.studio_bugbash.databricks_docs_index")
        ]
   )

また、 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

要件