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

この記事では、Mosaic AI Agent Framework と、LangGraph、PyFunc、OpenAI などの一般的なエージェントオーサリングライブラリを使用して、Python で AI エージェントを作成する方法について説明します。

必要条件

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

この記事のアプローチを使用してエージェントを作成およびデプロイするには、次の最小パッケージバージョンが必要です。

databricks-agents バージョン0.16.0以降
mlflow バージョン 2.20.2 以降
Python 3.10 以降。この要件を満たすには、サーバレスコンピュートまたは Databricks Runtime 13.3 LTS 以降を使用できます。

%pip install -U -qqqq databricks-agents>=0.16.0 mlflow>=2.20.2

Databricks では、エージェントの作成時に Databricks AI Bridge 統合パッケージをインストールすることもお勧めします。これらの統合パッケージ (databricks-langchain 、databricks-openai など)は、エージェント作成フレームワークやSDKにおけるDatabricks AI/BI Genie、Vector Searchなどの機能と対話するためのAPIの共有レイヤーを提供します。

LangChain/LangGraph
OpenAI
Pure Python agents

%pip install -U -qqqq databricks-langchain

%pip install -U -qqqq databricks-openai

%pip install -U -qqqq databricks-ai-bridge

`ChatAgent`を使用してエージェントを作成する

Databricksでは本番運用グレードのエージェントを作成するには、MLflowの ChatAgent インターフェースを使用することをお勧めします。このチャットスキーマ仕様は、エージェントのシナリオ向けに設計されており、OpenAI ChatCompletion スキーマと似ていますが、厳密には互換性がありません。また、ChatAgentではマルチターンのツールコールエージェントの機能も追加されています。

ChatAgent を使用してエージェントを作成すると、次の利点があります。

高度なエージェント機能
- ストリーミング 出力: 小さなチャンクで出力をストリーミングすることで、対話型のユーザーエクスペリエンスを実現します。
- 包括的なツール呼び出しメッセージ履歴 : 中間ツール呼び出しメッセージを含む複数のメッセージを返し、品質と会話管理を向上させます。
- ツールコール確認支援
- マルチエージェントシステムのサポート
開発、デプロイ、モニタリングの効率化
- フレームワークに依存しない Databricks 機能の統合 :任意のフレームワークでエージェントを作成し、 AI Playground、エージェント評価、エージェントモニタリングとすぐに互換性を持たせることができます。
- 型付きオーサリングインターフェイス :型付きPythonクラスを使用してエージェントコードを記述し、IDEとノートブックのオートコンプリートの恩恵を受けます。
- シグネチャの自動推論 : MLflow は、エージェントのログ記録時に ChatAgent シグネチャを自動的に推論するため、登録とデプロイが簡素化されます。ログ記録中のモデルのシグネチャの推論を参照してください。
- AI Gateway 拡張推論テーブル : AI Gateway 推論テーブルは、デプロイされたエージェントに対して自動的に有効になり、詳細なリクエストログメタデータにアクセスできます。

ChatAgentの作成方法については、次のセクションの例と MLflow のドキュメント - ChatAgent インターフェイスとはを参照してください。

`ChatAgent` の例

次のノートブックでは、一般的なライブラリである OpenAI と LangGraph を使用して、ストリーミングと非ストリーミングの ChatAgents を作成する方法を示します。

ツールを追加してこれらのエージェントの機能を拡張する方法については、 AI エージェントツールを参照してください。

マルチエージェントシステム `ChatAgent` 例

Genie を使用してマルチエージェント・システムを作成する方法については、マルチエージェント・システムでの Genie の使用を参照してください。

ストリーミング出力エージェントの作成

ストリーミングエージェントは、より小さな増分チャンクの連続ストリームで応答を配信します。ストリーミング出力により、エンドユーザーはエージェント出力が生成されたときに読み取ることができるため、知覚されるレイテンシーが短縮され、会話型エージェントの全体的なユーザーエクスペリエンスが向上します。

ストリーミングChatAgentを作成するには、ChatAgentChunkオブジェクトを生成するジェネレータを返す predict_stream メソッドを定義します。各オブジェクトにはレスポンスの一部が含まれます。理想的な ChatAgent ストリーミング動作の詳細については、 MLflow のドキュメントを参照してください。

次のコードは predict_stream 関数の例を示しています。ストリーミングエージェントの完全な例については、「 ChatAgent の例」を参照してください。

Python
def predict_stream(
  self,
  messages: list[ChatAgentMessage],
  context: Optional[ChatContext] = None,
  custom_inputs: Optional[dict[str, Any]] = None,
) -> Generator[ChatAgentChunk, None, None]:
  # Convert messages to a format suitable for your agent
  request = {"messages": self._convert_messages_to_dict(messages)}

  # Stream the response from your agent
  for event in self.agent.stream(request, stream_mode="updates"):
    for node_data in event.values():
      # Yield each chunk of the response
      yield from (
        ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
      )

Databricks モデルサービングにデプロイメントできる`ChatAgent`の開発

Databricks は、 Databricks モデルサービングの分散環境にChatAgentをデプロイするため、マルチターンの会話中に、同じサービングレプリカがすべてのリクエストを処理できない場合があります。エージェントの状態の管理には、次の影響に注意してください。

ローカルキャッシュの回避 : ChatAgentをデプロイするときは、同じレプリカが複数ターンの会話のすべての要求を処理すると想定しないでください。各ターンのディクショナリ ChatAgentRequest スキーマを使用して内部状態を再構築します。
スレッドセーフな状態 : エージェントの状態をスレッドセーフに設計し、マルチスレッド環境での競合を防ぎます。
predict関数で状態を初期化 する: predict関数が呼び出されるたびに、ChatAgent初期化中ではなく、状態を初期化します。ChatAgent レベルで状態を格納すると、1 つの ChatAgent レプリカが複数のメッセージ交換からの要求を処理できるため、会話間で情報が漏洩し、競合が発生する可能性があります。

カスタム入力と出力

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

これらのシナリオでは、MLflow ChatAgent は custom_inputs フィールドと custom_outputsフィールドをネイティブにサポートします。

警告

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

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

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

Open notebook in new tab

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

Open notebook in new tab

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

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

AI Playground または Agent Review App で、歯車アイコンを選択します。
custom_inputs を有効にします。
エージェントの定義済み入力スキーマに一致するJSONオブジェクトを提供します。

カスタム取得者スキーマの指定

AIエージェントは一般的に、レトリーバーを使用して、ベクトル検索インデックスから非構造化データを検索し、クエリを実行します。たとえば、取得ツールについては、非構造化取得 AI エージェントツールを参照してください。

MLflow RETRIEVER スパンを使用してエージェント内でこれらのレトリーバーをトレースし、次のような Databricks 製品の機能を有効にします。

取得したソースドキュメントへのリンクを AI Playground UI に自動的に表示する
Agent Evaluation で自動的に実行される検索の根拠と関連性のジャッジ

注記

Databricks では、 databricks_langchain.VectorSearchRetrieverTool や databricks_openai.VectorSearchRetrieverTool などの Databricks AI Bridge パッケージによって提供されるレトリーバーツールの使用をお勧めします。これは、既に MLflow レトリーバースキーマに準拠しているためです。 AI Bridge を使用したベクトル検索取得ツールをローカルで開発するを参照してください。

エージェントがカスタムスキーマを持つレトリーバースパンを含んでいる場合は、 mlflow.models.set_retriever_schemaエージェントをコードで定義する場合。これにより、レトリーバーの出力列が MLflow の想定フィールド (primary_key、 text_column、 doc_uri) にマップされます。

Python
import mlflow
# Define the retriever's schema by providing your column names
# For example, the following call specifies the schema of a retriever that returns a list of objects like
# [
#     {
#         'document_id': '9a8292da3a9d4005a988bf0bfdd0024c',
#         'chunk_text': 'MLflow is an open-source platform, purpose-built to assist machine learning practitioners...',
#         'doc_uri': 'https://mlflow.org/docs/latest/index.html',
#         'title': 'MLflow: A Tool for Managing the Machine Learning Lifecycle'
#     },
#     {
#         'document_id': '7537fe93c97f4fdb9867412e9c1f9e5b',
#         'chunk_text': 'A great way to get started with MLflow is to use the autologging feature. Autologging automatically logs your model...',
#         'doc_uri': 'https://mlflow.org/docs/latest/getting-started/',
#         'title': 'Getting Started with MLflow'
#     },
# ...
# ]
mlflow.models.set_retriever_schema(
    # Specify the name of your retriever span
    name="mlflow_docs_vector_search",
    # Specify the output column name to treat as the primary key (ID) of each retrieved document
    primary_key="document_id",
    # Specify the output column name to treat as the text content (page content) of each retrieved document
    text_column="chunk_text",
    # Specify the output column name to treat as the document URI of each retrieved document
    doc_uri="doc_uri",
    # Specify any other columns returned by the retriever
    other_columns=["title"],
)

注記

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

環境間でのデプロイのためのエージェントコードのパラメータ化

エージェント・コードをパラメーター化して、異なる環境で同じエージェント・コードを再利用できます。

パラメータは、Pythonディクショナリーまたは.yamlファイルで定義するキーと値のペアです。

コードを構成するには、Python ディクショナリまたは .yamlファイルを使用してModelConfigを作成します。ModelConfig は、柔軟な構成管理を可能にするキーと値のパラメーターのセットです。たとえば、開発中にディクショナリを使用し、それを本番運用デプロイメントおよび CI/CD用の.yamlファイルに変換できます。

ModelConfigの詳細については、MLflow のドキュメントを参照してください。

ModelConfigの例を以下に示します。

YAML
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

エージェントコードでは、 .yaml ファイルまたはディクショナリからデフォルト(開発)設定を参照できます。

Python
import mlflow
# 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
# You can also use model_config.to_dict() to convert the loaded config object
# into a dictionary
value = model_config.get('sample_param')

次に、エージェントをログに記録するときに、log_modelするパラメーターmodel_configを指定しますログに記録されたエージェントをロードするときに使用するパラメーターのカスタムセットを指定します。MLflow のドキュメント - ModelConfigをご覧ください。

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

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

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

必要条件​

ChatAgentを使用してエージェントを作成する​

ChatAgent の例​

LangGraphツールコールエージェント

OpenAIツールコールエージェント

OpenAI Responses API ツール呼び出しエージェント

OpenAIチャット専用エージェント

マルチエージェントシステム ChatAgent 例​

ストリーミング出力エージェントの作成​

Databricks モデルサービングにデプロイメントできるChatAgentの開発​

カスタム入力と出力​

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

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

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

カスタム取得者スキーマの指定​

環境間でのデプロイのためのエージェントコードのパラメータ化​

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

次のステップ​