AI エージェントとそのツールを作成する

プレビュー

この機能はパブリックプレビュー段階です。

この記事では、Mosaic AI Agent Framework を使用して AI エージェントとツールを作成する方法について説明します。

AI Playground を使用して、ツール呼び出しエージェントのプロトタイプをすばやく作成し、それらを Mosaic AI Agent Framework にエクスポートする方法を学びます。

要件

「複合 AI システムと AI エージェントとは」で説明されている AI エージェントとツールの概念を理解します。
Databricks では、エージェントを開発するときに最新バージョンの MLflow Python クライアントをインストールすることをお勧めします。 mlflowバージョン要件に関する情報については、依存リソースの認証を参照してください。

AI エージェントツールの作成

AIエージェントは、言語生成以外のアクション(構造化データや非構造化データの取得、コードの実行、リモートサービスとの通信など)を実行するためのツールを使用します(例: Eメール or Slack メッセージを送信してください)。

Mosaic AI Agent Framework を使用してエージェントにツールを提供するには、次の方法を任意に組み合わせて使用できます。

既存の Unity Catalog 関数をツールとして作成または使用します。Unity Catalog の関数は、検出、ガバナンス、共有を容易にし、大規模なデータセットに変換と集計を適用するのに適しています。
ツールをエージェントコード内の Python 関数としてローカルに定義します。このアプローチは、 REST APIsへの呼び出し、任意のコードまたはライブラリの使用、または非常に短いレイテンシでツールを実行する必要がある状況で役立ちます。このアプローチには、Unity Catalog 関数によって提供される組み込みの検出可能性とガバナンスが欠けています。 Databricks では、エージェントを構築する際にこのトレードオフを比較検討し、どのアプローチが最適かを判断することをお勧めします。

どちらのアプローチも、カスタムPythonコードで記述されたエージェント、またはLangGraphなどのエージェントオーサリングライブラリを使用して機能します。

ツールを定義するときは、エージェント LLM がツールをいつ、どのように使用するかを理解できるように、ツール、そのパラメーター、および戻り値が文書化されていることを確認してください。

Unity Catalog 関数を使用したエージェントツールの作成

これらの例では、ノートブック環境または SQL エディターで記述された Unity Catalog 関数を使用して AI エージェントツールを作成します。

ノートブックのセルで次のコードを実行します。 %sqlノートブックマジックを使用して、python_exec という Unity Catalog 関数を作成します。

LLM は、このツールを使用して、ユーザーから提供された Python コードを実行できます。

%sql
CREATE OR REPLACE FUNCTION
main.default.python_exec (
 code STRING COMMENT 'Python code to execute. Remember to print the final result to stdout.'
)
RETURNS STRING
LANGUAGE PYTHON
DETERMINISTIC
COMMENT 'Executes Python code in the sandboxed environment and returns its stdout. The runtime is stateless and you can not read output of the previous tool executions. i.e. No such variables "rows", "observation" defined. Calling another tool inside a Python code is NOT allowed. Use standard python libraries only.'
AS $$
 import sys
 from io import StringIO
 sys_stdout = sys.stdout
 redirected_output = StringIO()
 sys.stdout = redirected_output
 exec(code)
 sys.stdout = sys_stdout
 return redirected_output.getvalue()
$$

SQL エディターで次のコードを実行します。

これは、LLM が架空の customer_data テーブルから構造化データを取得するために使用できる lookup_customer_info という Unity Catalog 関数を作成します。

CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
  customer_name STRING COMMENT 'Name of the customer whose info to look up'
)
RETURNS STRING
COMMENT 'Returns metadata about a particular customer given the customer name, including the customer email and ID. The
customer ID can be used for other queries.'
RETURN SELECT CONCAT(
    'Customer ID: ', customer_id, ', ',
    'Customer Email: ', customer_email
  )
  FROM main.default.customer_data
  WHERE customer_name = customer_name
  LIMIT 1;

AI Playgroundにおけるツール呼び出しエージェントのプロトタイプ

Unity Catalog 関数を作成したら、 AI Playground を使用してそれらを LLM に渡し、エージェントをテストできます。 AI Playground は、ツールコールエージェントのプロトタイプを作成するためのサンドボックスを提供します。

AIエージェントに満足したら、それをエクスポートしてPythonでさらに開発したり、そのままモデルサービングエンドポイントとしてデプロイしたりできます。

注：

Unity Catalog、サーバレスコンピュート、 Mosaic AI Agent Framework、およびトークン単位の従量課金基盤モデルまたは外部モデルが現在のワークスペースで使用可能であることAI Playgroundでエージェントをプロトタイプ化する必要があります。

ツール呼び出しエンドポイントのプロトタイプを作成する。

Playground から、 Function calling ラベルを持つモデルを選択します。
[ツール] を選択し、ドロップダウンで Unity Catalog 関数名を指定します。
チャットして、LLM、ツール、およびシステムプロンプトの現在の組み合わせをテストし、バリエーションを試します。

AI Playground エージェントのエクスポートとデプロイ

ツールを追加し、エージェントをテストした後、Playground エージェントを Python ノートブックにエクスポートします。

「エージェント・コードのエクスポート」をクリックしてPythonAIエージェントの開発およびデプロイに役立つノートブック・ノートブックを生成します。

エージェントコードをエクスポートすると、ワークスペースに3つのファイルが保存されます。
- agent ノートブック: LangChainを使用してエージェントを定義するPythonコードが含まれています。
- driver ノートブック: エージェントフレームワークを使用してPython エージェントをログに記録、トレース、登録するデプロイするためのAI Mosaic AIコードが含まれています。
- config.yml: ツール定義を含む、エージェントに関する構成情報が含まれます。
agentノートブックを開いて、エージェントを定義する LangChain コードを確認し、このノートブックを使用して、他のツールの定義やエージェントのパラメーターの調整など、エージェントをプログラムでテストおよび反復処理します。

注：

エクスポートされたコードは、AI プレイグラウンドセッションとは異なる動作をする場合があります。 Databricks では、エクスポートされたノートブックを実行してさらに反復処理とデバッグを行い、エージェントの品質を評価してから、エージェントをデプロイして他のユーザーと共有することをお勧めします。
エージェントの出力に問題がなければ、 driver ノートブックを実行して、エージェントをログに記録し、モデルサービングエンドポイントにデプロイできます。

コードでエージェントを定義する

AI Playground からエージェントコードを生成するだけでなく、LangChain や Python コードなどのフレームワークを使用して、自分でコードでエージェントを定義することもできます。 Agent Framework を使用してエージェントをデプロイするには、その入力がサポートされている入力形式と出力形式のいずれかに準拠している必要があります。

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

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

サポートされている入力形式

以下は、エージェントでサポートされている入力形式です。

（推奨）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特に現在のクエリ―と履歴を個別に管理したい場合など、マルチターンチャットアプリケーションに推奨されます。

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

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

サポートされている出力形式

エージェントは、次のいずれかのサポートされている出力形式を持っている必要があります。

（推奨）ChatCompletionResponse。この形式は、OpenAI応答形式の相互運用性を持つ顧客に推奨されます。
StringObjectResponse です。この形式は、最も簡単で解釈が簡単です。

LangChain の場合は、最後のチェーンステップとして MLflow の StringResponseOutputParser() または ChatCompletionsOutputParser() を使用します。これにより、LangChain AIメッセージがエージェント互換の形式にフォーマットされます。

  from mlflow.langchain.output_parsers import StringResponseOutputParser, ChatCompletionsOutputParser

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(fake_model)
      | StringResponseOutputParser() # use this for StringObjectResponse
      # ChatCompletionsOutputParser() # or use this for ChatCompletionResponse
  )

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

形式が確実に守られるように、predict()内のデータクラスから出力オブジェクトを構築できます。返されたオブジェクトは、シリアル化できるように辞書表現に変換する必要があります。

  from mlflow.models.rag_signatures import ChatCompletionRequest, ChatCompletionResponse, ChainCompletionChoice, Message

  class RAGModel(PythonModel):
    ...
      def predict(self, context, model_input: ChatCompletionRequest) -> ChatCompletionResponse:
        ...
        return asdict(ChatCompletionResponse(
            choices=[ChainCompletionChoice(message=Message(content=text))]
        ))

ノートブックの例

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

AI エージェントとそのツールを作成する

要件

AI エージェント ツールの作成

Unity Catalog 関数を使用したエージェントツールの作成

AI Playgroundにおけるツール呼び出しエージェントのプロトタイプ

AI Playground エージェントのエクスポートとデプロイ

コードでエージェントを定義する

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

サポートされている入力形式

サポートされている出力形式

ノートブックの例

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

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

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

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

次のステップ

AI エージェントツールの作成