AI エージェントをログに記録して登録する (モデル サービス)

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

Requirements

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

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

% pip install databricks-sdk

コード ベースのログ記録

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

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

この方法と、mlflow.models.predict() などのデプロイ前検証 API を使用して、サービスを提供するためにデプロイされたときにエージェントが確実に実行されるようにすることができます。

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

ログ記録中にモデル署名を推論する

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

1. 署名を手動で定義する
2. MLflow の Model Signature 推論機能を使用して、指定した入力例に基づいてエージェントの署名を自動的に生成します。 この方法は、署名を手動で定義するよりも便利です。

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

以下の LangChain と PyFunc の例では、Model Signature 推論を使用します。

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

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

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

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

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

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

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

   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(...) が呼び出されます。

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 モデルにログに記録します。

   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(...) が呼び出されます。

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 エージェントでは、タスクを完了するために他のリソースに対して認証を行う必要があることがよくあります。 たとえば、デプロイされたエージェントでは、ベクター検索インデックスにアクセスして非構造化データを照会したり、Prompt Registry にアクセスして動的プロンプトを読み込んだりする必要がある場合があります。

自動認証パススルーと代理認証は、エージェントのログイン時に構成が必要です。

エージェントを Unity カタログに登録する

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

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()を参照してください。

次のステップ

• AI エージェントにトレースを追加します。
• AI エージェントを展開します。