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

AIエージェントの登録

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

必要条件​

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

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

Python
% pip install databricks-sdk

コードベースの記録​

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

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

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

コードベースのログ記録の例については、ChatAgent 作成ノートブックの例を参照してください。

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

注記

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

ChatAgent インターフェイスを使用しない場合は、次のいずれかの方法を使用して、ログ時にエージェントの 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. ドライバー ノートブックで、次のコードを使用して ML agent.py を実行し、結果を 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-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 パラメーターを使用して依存リソースを指定します。

Python
import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection
)

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"),
      DatabricksUCConnection(connection_name="your_connection_name"),
    ]
  )

Databricks では、すべてのエージェント フレーバーの resources を手動で指定することをお勧めします。

注記

mlflow.langchain.log_model(...)を使用して LangChain エージェントをログに記録するときにリソースを指定しない場合、MLflow はリソースのベストエフォート型自動推論を実行します。ただし、これによりすべての依存関係がキャプチャされるわけではないため、エージェントの提供時またはクエリ時に認証エラーが発生する可能性があります。

次の表に、自動認証パススルーをサポートする Databricks リソースと、リソースをログに記録するために必要な最小mlflowバージョンを示します。

リソースタイプ

リソースをログに記録するために必要な最小 mlflow バージョン

ベクトル検索インデックス

mlflow 2.13.1 以降が必要です

モデルサービングエンドポイント

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 接続

mlflow 2.17.1 以降が必要です

OpenAIクライアントの自動認証​

エージェントが OpenAI クライアントを使用している場合は、Databricks SDK を使用してデプロイ中に自動的に認証します。 Databricks SDK には、承認が自動的に構成された OpenAI クライアントを構築するためのラッパーが用意されています。 ノートブックで次のコマンドを実行します。

Python
% pip install databricks-sdk[openai]
Python
from databricks.sdk import WorkspaceClient
def openai_client(self):
  w = WorkspaceClient()
  return w.serving_endpoints.get_open_ai_client()

次に、モデルサービングエンドポイントを resources の一部として指定して、デプロイ時に自動的に認証します。

エージェントを 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)

次のステップ​

に最終更新