メインコンテンツまでスキップ

AIエージェントの登録

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

必要条件​

ロギングの前に AI エージェントを作成します。

Databricks では、最新バージョンの databricks-sdkをインストールすることをお勧めします。

Python
% pip install databricks-sdk

コードベースの記録​

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

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

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

コードベースのログの例については、「 ResponsesAgent オーサリングのノートブックの例」を参照してください。

ロギングにおけるモデルシグネチャの推論​

注記

Databricks では、 ResponsesAgent インターフェイスを使用してエージェントを作成することをお勧めします。ResponsesAgent を使用している場合は、このセクションをスキップできます。MLflow は、エージェントの有効な署名を自動的に推測します。

ResponsesAgent インターフェイスを使用しない場合は、次のいずれかの方法を使用して、ログ時にエージェントの MLflow モデルシグネチャを指定する必要があります。

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

MLflow モデル シグネチャは、入力と出力を検証して、エージェントが AI Playground やレビュー アプリなどのダウンストリーム ツール。 また、エージェントの使用方法に関する他のアプリケーションのガイドも示します 効果的。

以下の LangChain と PyFunc の例では、モデルシグネチャの推論を使用しています。

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

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

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

  1. コードを使用してノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は agent.pyです。 ノートブックまたはファイルには、ここでは LangChain エージェント ( lc_agentと呼びます) が含まれている必要があります。

  2. ノートブックまたはファイルに mlflow.models.set_model(lc_agent) を含めます。

  3. ドライバーノートブックとして機能する新しいノートブックを作成します（この例ではdriver.pyと呼ばれます）。

  4. ドライバー ノートブックで、次のコードを使用して ML agent.py を実行し、結果を MLflow モデルに記録します。

    Python
    mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)

    resources パラメーターは、ベクトル検索インデックスや基盤モデルを提供するサービング エンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージド リソースを宣言します。詳細については、 情報 「自動認証パススルーの実装」を参照してください。

  5. モデルをデプロイします。生成AIアプリケーションのためのエージェントのデプロイを参照してください。

  6. サービング環境が読み込まれると、 agent.py が実行されます。

  7. サービングリクエストが来るとlc_agent.invoke(...)が呼び出されます。

Python

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 を使用してエージェントをログに記録する方法を示しています。

  1. コードを使用してノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は agent.pyです。 ノートブックまたはファイルには、 PyFuncClassという名前の PyFunc クラスが含まれている必要があります。

  2. ノートブックまたはファイルに mlflow.models.set_model(PyFuncClass) を含めます。

  3. ドライバーノートブックとして機能する新しいノートブックを作成します（この例ではdriver.pyと呼ばれます）。

  4. ドライバー ノートブックで、次のコードを使用して agent.py を実行し、log_model() を使用して結果を MLflow モデルに記録します。

    Python
    mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)

    resources パラメーターは、ベクトル検索インデックスや基盤モデルを提供するサービング エンドポイントなど、エージェントにサービスを提供するために必要な Databricks マネージド リソースを宣言します。詳細については、 情報 「自動認証パススルーの実装」を参照してください。

  5. モデルをデプロイします。生成AIアプリケーションのためのエージェントのデプロイを参照してください。

  6. サービング環境が読み込まれると、 agent.py が実行されます。

  7. サービングリクエストが来るとPyFuncClass.predict(...)が呼び出されます。

Python
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-meta-llama-3-3-70b-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)

Databricks リソースの認証​

AI エージェントは、多くの場合、タスクを完了するために他のリソースに対して認証する必要があります。たとえば、デプロイされたエージェントは、非構造化データをクエリするためにベクトル検索インデックスにアクセスしたり、動的プロンプトをロードするためにプロンプトレジストリにアクセスしたりする必要がある場合があります。

自動認証パススルー と 代理認証では、 エージェントのログ記録中に構成する必要があります。

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

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

Python
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)

詳しくは、mlflow.register_model()を参照してください。

次のステップ​

に最終更新