次の方法で共有

AI エージェントの作成とログを行う

  • [アーティクル]

重要

この機能はパブリック プレビュー段階にあります。

この記事では、Mosaic AI エージェント フレームワークを使用して RAG アプリケーションなどの AI エージェントを作成し、ログを行う方法について説明します。

チェーンとエージェントとは

多くの場合、AI システムには多くのコンポーネントがあります。 たとえば、AI システムは、ベクトル インデックスからドキュメントを取得し、それらのドキュメントを使用してプロンプト テキストを補完し、基本モデルを使用して応答を要約することができます。 これらのコンポーネントをリンクするコード (ステップとも呼ばれます) は、"チェーン" と呼ばれます。

"エージェント" は、大規模言語モデルに依存して、入力に基づいて実行する手順を決定する、はるかに高度な AI システムです。 これに対し、チェーンは、特定の結果を達成することを目的とした、ハードコーディングされた一連のステップです。

エージェント フレームワークでは、任意のライブラリまたはパッケージを使用してコードを作成できます。 また、エージェント フレームワークを使用すると、コードを開発してテストするときに、反復処理が簡単になります。 実際のコードを変更しなくても、トレース可能な方法でコード パラメーターを変更できる構成ファイルを設定できます。

要件

Databricks マネージド ベクトル検索インデックスを使用するエージェントの場合、mlflow バージョン 2.13.1 以降では、ベクトル インデックスで自動承認を使用する必要があります。

RAG エージェントの入力スキーマ

チェーンでは次の入力形式がサポートされています。

  • (推奨) OpenAI チャット完了スキーマを使用したクエリ。 messages パラメーターとしてオブジェクトの配列が必要です。 この形式は RAG アプリケーションに最適です。

    question = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-Augmented Generation?",
        },
        {
            "role": "assistant",
            "content": "RAG, or Retrieval Augmented Generation, is a generative AI design pattern that combines a large language model (LLM) with external knowledge retrieval. This approach allows for real-time data connection to generative AI applications, improving their accuracy and quality by providing context from your data to the LLM during inference. Databricks offers integrated tools that support various RAG scenarios, such as unstructured data, structured data, tools & function calling, and agents.",
        },
        {
            "role": "user",
            "content": "How to build RAG for unstructured data",
        },
    ]
}

  • SplitChatMessagesRequest 現在のクエリと履歴を個別に管理する場合は特に、マルチターン チャット アプリケーションにお勧めします。

    {
"query": "What is MLflow",
"history": [
  {
  "role": "user",
  "content": "What is Retrieval-augmented Generation?"
  },
  {
  "role": "assistant",
  "content": "RAG is"
  }
  ]
}

LangChain の場合、Databricks では、LangChain 式言語でチェーンを記述することをお勧めします。 チェーン定義コードでは、itemgetter を使用して、使用している入力形式に応じてメッセージ、あるいは query または history オブジェクトを取得できます。

RAG エージェントの出力スキーマ

コードは、次のサポートされる出力形式のいずれかに準拠している必要があります。

  • (推奨) ChatCompletionResponse。 この形式は、OpenAI 応答形式の相互運用性を持つお客様にお勧めします。
  • StringResponse。 この形式は解釈するのが最も簡単かつシンプルです。

LangChain の場合は、最後のチェーン ステップとして StrOutputParser() を使用します。 出力は 1 つの文字列値を返す必要があります。

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(fake_model)
      | StrOutputParser()
  )

PyFunc を使用している場合、Databricks では、型ヒントを使用して、mlflow.models.rag_signatures で定義されたクラスのサブクラスである入力および出力データ クラスを使用して predict() 関数に注釈を付けることをお勧めします。

predict() 内のデータ クラスから出力オブジェクトを構築して、形式に従うことができます。 返されるオブジェクトは、シリアル化できるようにディクショナリ表現に変換する必要があります。

パラメーターを使用して高品質の反復処理を制御する

エージェント フレームワークでは、パラメーターを使用してエージェントの実行方法を制御できます。 これにより、コードを変更することなく、エージェントのさまざまな特性を使用して迅速に反復処理できます。 パラメーターは、Python ディクショナリまたは .yaml ファイルで定義するキーと値のペアです。

コードを構成するには、キー値パラメーターのセットである ModelConfig を作成します。 ModelConfig は Python ディクショナリまたは .yaml ファイルです。 たとえば、開発中にディクショナリを使用し、運用環境のデプロイと CI/CD 用の .yaml ファイルに変換できます。 ModelConfig の詳細については、MLflow のドキュメントを参照してください。

ModelConfig の例を次に示します。

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-dbrx-instruct
vector_search_index: ml.docs.databricks_docs_index
prompt_template: 'You are a hello world bot. Respond with a reply to the user''s
  question that indicates your prompt template came from a YAML file. Your response
  must use the word "YAML" somewhere. User''s question: {question}'
prompt_template_input_vars:
- question

コードから構成を呼び出すには、次のいずれかを使用します。

# Example for loading from a .yml file
config_file = "configs/hello_world_config.yml"
model_config = mlflow.models.ModelConfig(development_config=config_file)

# Example of using a dictionary
config_dict = {
    "prompt_template": "You are a hello world bot. Respond with a reply to the user's question that is fun and interesting to the user. User's question: {question}",
    "prompt_template_input_vars": ["question"],
    "model_serving_endpoint": "databricks-dbrx-instruct",
    "llm_parameters": {"temperature": 0.01, "max_tokens": 500},
}

model_config = mlflow.models.ModelConfig(development_config=config_dict)

# Use model_config.get() to retrieve a parameter value
value = model_config.get('sample_param')

エージェントのログを記録する

エージェントのログは、開発プロセスの基礎となります。 ログにより、"特定の時点" でのエージェントのコードと構成がキャプチャされるため、構成の品質を評価できます。 エージェントを開発する場合、Databricks では、シリアル化ベースのログではなく、コードベースのログを使用することをお勧めします。 各種類のログの長所と短所の詳細については、「コードベースとシリアル化ベースのログ」を参照してください。

このセクションでは、コードベースのログを使用する方法について説明します。 シリアル化ベースのログの使用方法の詳細については、「シリアル化ベースのログ ワークフロー」を参照してください。

コードベースのログ ワークフロー

コードベースのログの場合、エージェントまたはチェーンをログに記録するコードは、チェーン コードとは別のノートブックに含まれている必要があります。 このノートブックは、ドライバー ノートブックと呼ばれます。 ノートブックの例については、「サンプル ノートブック」を参照してください。

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

  1. コードを含むノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は chain.py です。 ノートブックまたはファイルには、ここでは lc_chain と呼ばれる LangChain チェーンが含まれている必要があります。
  2. ノートブックまたはファイルに mlflow.models.set_model(lc_chain) を含めます。
  3. ドライバー ノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。
  4. ドライバー ノートブックに、mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”) の呼び出しを含めます。 この呼び出しは chain.py を実行し、結果のログを MLflow モデルに記録します。
  5. モデルをデプロイします。
  6. サービス環境が読み込まれると、chain.py が実行されます。
  7. サービス要求が入ると、lc_chain.invoke(...) が呼び出されます。

PyFunc を使用したコードベースのログ ワークフロー

  1. コードを含むノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は chain.py です。 ノートブックまたはファイルには、ここでは PyFuncClass と呼ばれる PyFunc クラスが含まれている必要があります。
  2. ノートブックまたはファイルに mlflow.models.set_model(PyFuncClass) を含めます。
  3. ドライバー ノートブックとして機能する新しいノートブックを作成します (この例では driver.py と呼ばれます)。
  4. ドライバー ノートブックに、mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”) の呼び出しを含めます。 この呼び出しは chain.py を実行し、結果のログを MLflow モデルに記録します。
  5. モデルをデプロイします。
  6. サービス環境が読み込まれると、chain.py が実行されます。
  7. サービス要求が入ると、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 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
  )

# or use a PyFunc model
# 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,
#   )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

モデルが正しくログに記録されたことを確認するには、チェーンを読み込み、invoke を呼び出します。

# Using LangChain
model = mlflow.langchain.load_model(logged_chain_info.model_uri)
model.invoke(example)

# Using PyFunc
model = mlflow.pyfunc.load_model(logged_chain_info.model_uri)
model.invoke(example)

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)

ノートブックの例

これらのノートブックは、Databricks でチェーン アプリケーションを作成する方法を示す単純な "Hello World" チェーンを作成します。 最初の例では、単純なチェーンを作成します。 2 番目のノートブック例では、開発中にパラメーターを使用してコードの変更を最小限に抑える方法を示します。

単純なチェーン ノートブック

ノートブックを入手

単純なチェーン ドライバー ノートブック

ノートブックを入手

パラメーター化されたチェーン ノートブック

ノートブックを入手

パラメーター化されたチェーン ドライバー ノートブック

ノートブックを入手

コードベースとシリアル化ベースのログ

チェーンを作成してログに記録するには、コードベースの MLflow ログまたはシリアル化ベースの MLflow ログを使用できます。 Databricks では、コードベースのログを使用することをお勧めします。

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

シリアル化ベースの MLflow ログでは、Python 環境のチェーンのコードと現在の状態は、多くの場合、pickle や joblib などのライブラリを使用して、ディスクにシリアル化されます。 チェーンがデプロイされると、Python 環境が復元され、シリアル化されたオブジェクトがメモリに読み込まれるため、エンドポイントが呼び出されたときに起動できます。

各方法の長所と短所を表に示します。

メソッド 長所 短所
コード ベースの MLflow ログ * 多くの一般的な GenAI ライブラリではサポートされていない、シリアル化に固有の制限を克服します。
* 後で参照できるように元のコードのコピーを保存します。
* シリアル化できる 1 つのオブジェクトにコードを再構築する必要はありません。		 log_model(...) は、チェーンのコード (ドライバー ノートブックと呼ばれます) とは "異なる" ノートブックから呼び出す必要があります。
シリアル化ベースの MLflow ログ log_model(...) は、モデルが定義されているのと同じノートブックから呼び出すことができます。 * 元のコードは使用できません。
* チェーンで使用されるすべてのライブラリとオブジェクトは、シリアル化をサポートする必要があります。

シリアル化ベースのログ ワークフロー

Databricks では、シリアル化ベースのログではなく、コードベースのログを使用することをお勧めします。 コードベースのログの使用方法の詳細については、「コードベースのログ ワークフロー」を参照してください。

このセクションでは、既存のシリアル化ベースのログの使用方法について説明します。

LangChain を使用したシリアル化ベースのログ ワークフロー

  1. コードを含むノートブックまたは Python ファイルを作成します。 ノートブックまたはファイルには、ここでは lc_chain と呼ばれる LangChain チェーンが含まれている必要があります。
  2. ノートブックまたはファイルに mlflow.lang_chain.log_model(lc_model=lc_chain) を含めます。
  3. PyFuncClass() のシリアル化されたコピーが MLflow モデルに記録されます。
  4. モデルをデプロイします。
  5. サービス環境が読み込まれると、PyFuncClass が逆シリアル化されます。
  6. サービス要求が入ると、lc_chain.invoke(...) が呼び出されます。

PyFunc を使用したシリアル化ベースのログ ワークフロー

  1. コードを含むノートブックまたは Python ファイルを作成します。 この例では、ノートブックまたはファイルの名前は notebook.py です。 ノートブックまたはファイルには、ここでは PyFuncClass と呼ばれる PyFunc クラスが含まれている必要があります。
  2. mlflow.pyfunc.log_model(python_model=PyFuncClass()) を notebook.py に含めます。
  3. PyFuncClass() のシリアル化されたコピーが MLflow モデルに記録されます。
  4. モデルをデプロイします。
  5. サービス環境が読み込まれると、PyFuncClass が逆シリアル化されます。
  6. サービス要求が入ると、PyFuncClass.predict(...) が呼び出されます。