コードで AI エージェントを作成する

この記事では、MLflow ChatModelを使用して、コードで AI エージェントを作成する方法について説明します。 Databricks は MLflow ChatModel を活用して、評価、トレース、デプロイなどの Databricks AI エージェント機能との互換性を確保します。

ChatModelとは

ChatModel は、会話型 AI エージェントの作成を簡略化するように設計された MLflow クラスです。 OpenAIのChatCompletion APIと互換性のあるモデルを構築するための標準化されたインターフェースを提供します。

ChatModel OpenAI の ChatCompletion スキーマを拡張します。 このアプローチにより、ChatCompletion 標準をサポートするプラットフォームとの広範な互換性を維持しながら、独自のカスタム機能を追加できます。

ChatModelを使用することで、開発者は、本番運用可能なモデルをデプロイするために不可欠なエージェントの追跡、評価、ライフサイクル管理のための Databricks ツールや MLflow ツールと互換性のあるエージェントを作成できます。

「MLflow: ChatModel の概要」を参照してください。

要件

Databricks では、エージェントを開発するときに最新バージョンの MLflow Python クライアントをインストールすることをお勧めします。

この記事のアプローチを使用してエージェントを作成およびデプロイするには、次の要件を満たす必要があります。

  • databricks-agentsバージョン0.15.0以降をインストールします

  • mlflowバージョン2.20.0以降をインストールします

%pip install -U -qqqq databricks-agents>=0.15.0 mlflow>=2.20.0

ChatModelエージェントの作成

エージェントは 、mlflow.pyfunc.ChatModel のサブクラスとして作成できます。 この方法には、次の利点があります。

  • 型指定された Python クラスを使用して、ChatCompletion スキーマと互換性のあるエージェントコードを記述できます。

  • MLflow は、エージェントのログを記録するときに、 input_exampleがなくても、チャット完了と互換性のある署名を自動的に推論します。 これにより、エージェントの登録とデプロイのプロセスが簡素化されます。 ログ記録中のモデルの署名の推論を参照してください。

次のコードは、Databricks ノートブックで実行するのが最適です。 ノートブックは、エージェントの開発、テスト、および反復処理に便利な環境を提供します。

MyAgent クラスは mlflow.pyfunc.ChatModelを拡張し、必要な predict メソッドを実装します。これにより、Mosaic AI Agent Framework との互換性が確保されます。

このクラスには、ストリーミング出力を処理するためのオプションのメソッド ( _create_chat_completion_chunkpredict_stream ) も含まれています。

from dataclasses import dataclass
from typing import Optional, Dict, List, Generator
from mlflow.pyfunc import ChatModel
from mlflow.types.llm import (
    # Non-streaming helper classes
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionChunk,
    ChatMessage,
    ChatChoice,
    ChatParams,
    # Helper classes for streaming agent output
    ChatChoiceDelta,
    ChatChunkChoice,
)

class MyAgent(ChatModel):
    """
    Defines a custom agent that processes ChatCompletionRequests
    and returns ChatCompletionResponses.
    """
    def predict(self, context, messages: list[ChatMessage], params: ChatParams) -> ChatCompletionResponse:
        last_user_question_text = messages[-1].content
        response_message = ChatMessage(
            role="assistant",
            content=(
                f"I will always echo back your last question. Your last question was: {last_user_question_text}. "
            )
        )
        return ChatCompletionResponse(
            choices=[ChatChoice(message=response_message)]
        )

    def _create_chat_completion_chunk(self, content) -> ChatCompletionChunk:
        """Helper for constructing a ChatCompletionChunk instance for wrapping streaming agent output"""
        return ChatCompletionChunk(
                choices=[ChatChunkChoice(
                    delta=ChatChoiceDelta(
                        role="assistant",
                        content=content
                    )
                )]
            )

    def predict_stream(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> Generator[ChatCompletionChunk, None, None]:
        last_user_question_text = messages[-1].content
        yield self._create_chat_completion_chunk(f"Echoing back your last question, word by word.")
        for word in last_user_question_text.split(" "):
            yield self._create_chat_completion_chunk(word)

agent = MyAgent()
model_input = ChatCompletionRequest(
    messages=[ChatMessage(role="user", content="What is Databricks?")]
)
response = agent.predict(context=None, model_input=model_input)
print(response)

エージェント クラス MyAgent は 1 つのノートブックで定義されていますが、別のドライバー ノートブックを作成する必要があります。 driver ノートブックは、エージェントを Model Registry にログに記録し、モデルサービングを使用してエージェントをデプロイします。

この分離は、MLflow の Models from Code 手法を使用してモデルをログ記録するための Databricks の推奨ワークフローに従います。

例: ChatModel で LangChain をラップする

既存の LangChain モデルがあり、それを他の Mosaic AI エージェント機能と統合する場合は、互換性を確保するために MLflow ChatModel にラップできます。

このコード サンプルでは、次の手順を実行して、LangChain の実行可能オブジェクトを ChatModelとしてラップします。

  1. LangChain の最終出力を mlflow.langchain.output_parsers.ChatCompletionOutputParser でラップして、チャット完了出力の署名を生成します

  2. LangchainAgent クラスは mlflow.pyfunc.ChatModel を拡張し、次の 2 つの主要なメソッドを実装します。

    • predict: チェーンを呼び出し、書式設定された応答を返すことで、同期予測を処理します。

    • predict_stream: チェーンを呼び出し、応答のチャンクを生成することで、ストリーミング予測を処理します。

from mlflow.langchain.output_parsers import ChatCompletionOutputParser
from mlflow.pyfunc import ChatModel
from typing import Optional, Dict, List, Generator
from mlflow.types.llm import (
    ChatCompletionResponse,
    ChatCompletionChunk
)

chain = (
    <your chain here>
    | ChatCompletionOutputParser()
)

class LangchainAgent(ChatModel):
    def _prepare_messages(self, messages: List[ChatMessage]):
        return {"messages": [m.to_dict() for m in messages]}

    def predict(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> ChatCompletionResponse:
        question = self._prepare_messages(messages)
        response_message = self.chain.invoke(question)
        return ChatCompletionResponse.from_dict(response_message)

    def predict_stream(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> Generator[ChatCompletionChunk, None, None]:
        question = self._prepare_messages(messages)
        for chunk in chain.stream(question):
          yield ChatCompletionChunk.from_dict(chunk)

パラメーターを使用してエージェントを構成する

エージェントフレームワークでは、パラメータを使用してエージェントの実行方法を制御できます。これにより、コードを変更せずにエージェントの特性を変えて素早く反復処理を行うことができます。パラメータは、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')

レトリーバーのスキーマ設定

AIエージェントは、ベクトル検索インデックスを使用して関連ドキュメントを検索して返すエージェントツールの一種であるレトリーバーをよく使用します。 取得子の詳細については、「 非構造化取得 AI エージェント ツール」を参照してください。

レトリーバーが正しくトレースされていることを確認するには、 mlflow.models.set_retriever_schema エージェントをコードで定義する場合。 set_retriever_schema を使用して、返されたテーブルの列名を MLflow の想定フィールド (primary_keytext_columndoc_uriなど) にマップします。

# Define the retriever's schema by providing your column names
# These strings should be read from a config dictionary
mlflow.models.set_retriever_schema(
    name="vector_search",
    primary_key="chunk_id",
    text_column="text_column",
    doc_uri="doc_uri"
    # other_columns=["column1", "column2"],
)

注:

doc_uri列は、レトリーバーのパフォーマンスを評価するときに特に重要です。doc_uri は、レトリーバーによって返されるドキュメントの主な識別子であり、それらをグラウンド トゥルース評価セットと比較できます。 評価セットを参照してください。

また、レトリーバーのスキーマで追加の列を指定するには、 other_columns フィールドに列名のリストを指定します。

複数のレトリーバーがある場合は、各レトリーバースキーマに一意の名前を使用して、複数のスキーマを定義できます。

カスタム入力と出力

一部のシナリオでは、 client_typesession_idなどの追加のエージェント入力や、将来の対話のためにチャット履歴に含めるべきではない取得ソースリンクなどの出力が必要になる場合があります。

これらのシナリオでは、MLflow ChatModel ChatParams フィールド custom_inputcustom_output を使用して OpenAI チャット完了要求と応答をネイティブに拡張できます。

PyFuncおよびLangGraphエージェントのカスタム入力と出力を作成する方法については、次の例を参照してください。

警告

Agent Evaluation レビュー アプリは、現在、追加の入力フィールドを持つエージェントのトレースのレンダリングをサポートしていません。

PyFuncカスタムスキーマ

次のノートブックは、PyFunc を使用したカスタムスキーマの例を示しています。

PyFunc カスタムスキーマエージェントノートブック

ノートブックを新しいタブで開く

PyFuncカスタムスキーマドライバーノートブック

ノートブックを新しいタブで開く

LangGraph カスタムスキーマ

次のノートブックは、LangGraph を使用したカスタムスキーマの例を示しています。 ノートブックの wrap_output 関数を変更して、メッセージ・ストリームから情報を解析および抽出できます。

LangGraph カスタム スキーマ エージェント ノートブック

ノートブックを新しいタブで開く

LangGraph カスタム スキーマ ドライバー ノートブック

ノートブックを新しいタブで開く

AI Playground とエージェント レビュー アプリでcustom_inputsを提供する

エージェントが custom_inputs フィールドを使用して追加の入力を受け入れる場合は、 AI Playgroundエージェントレビューアプリの両方でこれらの入力を手動で提供できます。

  1. AI Playground または Agent Review App で、歯車アイコン 歯車アイコン を選択します。

  2. custom_inputsを有効にします。

  3. エージェントの定義済み入力スキーマに一致するJSONオブジェクトを提供します。

    AIプレイグラウンドでcustom_inputsを提供します。

ストリーミング エラーの伝播

Mosaic AI は、ストリーミング中に発生したエラーを伝播し、最後のトークンを databricks_output.error未満にします。 このエラーを適切に処理して表示するかどうかは、呼び出し元のクライアント次第です。

{
  "delta": …,
  "databricks_output": {
    "trace": {...},
    "error": {
      "error_code": BAD_REQUEST,
      "message": "TimeoutException: Tool XYZ failed to execute"
    }
  }
}

ノートブックの例

これらのノートブックは、Databricks でのエージェントの作成を示すための単純な "Hello, world" チェーンを作成します。 最初の例では単純なチェーンを作成し、2 番目のサンプル ノートブックでは、開発中のコード変更を最小限に抑えるためにパラメーターを使用する方法を示します。

シンプルなチェーンノートブック

ノートブックを新しいタブで開く

シンプルチェーンドライバーノート

ノートブックを新しいタブで開く

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

ノートブックを新しいタブで開く

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

ノートブックを新しいタブで開く