エージェント ログに記録して登録するAI
プレビュー
この機能はパブリックプレビュー段階です。
Mosaic AI Agent Framework を使用して AI エージェント をログに記録します。 エージェントのログ記録は、開発プロセスの基本です。 ログ記録は、エージェントのコードと構成の「時点」をキャプチャするため、構成の品質を評価できます。
要件
ログに記録する前に AI エージェントを作成します。
コードベースのログ記録とシリアライズベースのログ記録
コードベースの MLflow ログ記録またはシリアル化ベースの MLflow ログ記録を使用できます。 Databricks では、コードベースのログ記録を使用することをお勧めします。
コードベースの MLflow ログ記録: チェーンのコードは Python ファイルとしてキャプチャされます。 Python 環境は、パッケージのリストとしてキャプチャされます。 チェーンがデプロイされると、Python 環境が復元され、チェーンのコードが実行されてチェーンがメモリにロードされ、エンドポイントが呼び出されたときに呼び出すことができます。
シリアル化ベースの MLflow ログ記録: Python 環境でのチェーンのコードと現在の状態は、多くの場合、pickle
や joblib
などのライブラリを使用してディスクにシリアル化されます。チェーンがデプロイされると、Python 環境が復元され、シリアル化されたオブジェクトがメモリにロードされるため、エンドポイントが呼び出されたときに呼び出すことができます。
表は、各方法の長所と短所を示しています。
メソッド |
利点 |
デメリット |
---|---|---|
コードベースのMLflowログ記録 |
|
|
シリアル化ベースのMLflowログ記録 |
|
|
コードベースのログ記録では、エージェントまたはチェーンを記録するコードは、チェーンコードとは別のノートブックに保存する必要があります。このノートブックは、ドライバーノートブックと呼ばれます。ノートブックの例については、「ノートブックの例」を参照してください。
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)