コードで 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_chunk
と predict_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
としてラップします。
LangChain の最終出力を
mlflow.langchain.output_parsers.ChatCompletionOutputParser
でラップして、チャット完了出力の署名を生成します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_key
、text_column
、doc_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_type
や session_id
などの追加のエージェント入力や、将来の対話のためにチャット履歴に含めるべきではない取得ソースリンクなどの出力が必要になる場合があります。
これらのシナリオでは、MLflow ChatModel
ChatParams フィールド custom_input
と custom_output
を使用して OpenAI チャット完了要求と応答をネイティブに拡張できます。
PyFuncおよびLangGraphエージェントのカスタム入力と出力を作成する方法については、次の例を参照してください。
警告
Agent Evaluation レビュー アプリは、現在、追加の入力フィールドを持つエージェントのトレースのレンダリングをサポートしていません。
LangGraph カスタムスキーマ
次のノートブックは、LangGraph を使用したカスタムスキーマの例を示しています。 ノートブックの wrap_output
関数を変更して、メッセージ・ストリームから情報を解析および抽出できます。
AI Playground とエージェント レビュー アプリでcustom_inputs
を提供する
エージェントが custom_inputs
フィールドを使用して追加の入力を受け入れる場合は、 AI Playground と エージェントレビューアプリの両方でこれらの入力を手動で提供できます。
AI Playground または Agent Review App で、歯車アイコン
を選択します。
custom_inputsを有効にします。
エージェントの定義済み入力スキーマに一致するJSONオブジェクトを提供します。
ストリーミング エラーの伝播
Mosaic AI は、ストリーミング中に発生したエラーを伝播し、最後のトークンを databricks_output.error
未満にします。 このエラーを適切に処理して表示するかどうかは、呼び出し元のクライアント次第です。
{
"delta": …,
"databricks_output": {
"trace": {...},
"error": {
"error_code": BAD_REQUEST,
"message": "TimeoutException: Tool XYZ failed to execute"
}
}
}