コードで 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

%pip install -U -qqqq databricks-langchain

OpenAI

%pip install -U -qqqq databricks-openai

Pure Python agents

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

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

Databricks では、本番運用グレードのエージェントを作成するための MLflow ChatAgent インターフェイスを推奨しています。このチャットスキーマ仕様は、OpenAI ChatCompletion スキーマと似ていますが、厳密には互換性がありません。

ChatAgent 次の利点があります。

• 高度なエージェント機能

  • マルチエージェントのサポート
  • ストリーミング 出力: 小さなチャンクで出力をストリーミングすることで、対話型のユーザー エクスペリエンスを実現します。
  • 包括的なツール呼び出しメッセージ履歴 : 中間ツール呼び出しメッセージを含む複数のメッセージを返し、品質と会話管理を向上させます。
  • ツールコール確認支援

• 開発、デプロイ、モニタリングの効率化

  • 任意のフレームワークを使用してエージェントを作成する : ChatAgent インターフェイスを使用して既存のエージェントをラップし、 AI Playground、エージェント評価、およびエージェント モニタリングとすぐに使用できる互換性を実現します。
  • 型付きオーサリングインターフェイス :型付きPythonクラスを使用してエージェントコードを記述し、IDEとノートブックのオートコンプリートの恩恵を受けます。
  • シグネチャの自動推論 : MLflow は、エージェントのログ記録時に ChatAgent シグネチャを自動的に推論するため、登録とデプロイが簡素化されます。 ログ記録中のモデルのシグネチャの推論を参照してください。
  • AI Gateway 拡張推論テーブル : AI Gateway 推論テーブルは、デプロイされたエージェントに対して自動的に有効になり、詳細なリクエストログメタデータにアクセスできます。

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

ChatAgent の例

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

LangGraph

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

Open notebook in new tab

OpenAI

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

Open notebook in new tab

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

Open notebook in new tab

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

Open notebook in new tab

AutoGen

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

Open notebook in new tab

DSPy

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

Open notebook in new tab

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

マルチエージェントの例

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

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

ストリーミング エージェントは、より小さな増分チャンクの連続ストリームで応答を配信します。ストリーミングにより、知覚される遅延が短縮され、会話型エージェントのユーザーエクスペリエンスが向上します。

ストリーミングChatAgentを作成するには、ChatAgentChunkオブジェクトを生成するジェネレータを返す predict_stream メソッドを定義します。各ChatAgentChunkにはレスポンスの一部が含まれます。理想的な 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カスタムスキーマエージェントノートブック

Open notebook in new tab

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

Open notebook in new tab

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

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

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

2. custom_inputs を有効にします。

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

   

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

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-meta-llama-3-3-70b-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-meta-llama-3-3-70b-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."
    }
  }
}

次のステップ

• 独自のエージェントツールを作成します。
• AI エージェントをログに記録します。
• AIエージェントにトレースを追加します。
• AI エージェントをデプロイします。