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

プレビュー

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

この記事では、Mosaic AI Agent Framework を使用して AI エージェント用のツールを作成する方法について説明します。詳細については、「複合 AI システムと AI エージェントとは」を参照してください。

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

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

Unity Catalog 関数でツールを作成する: Unity Catalog 関数を使用すると、検出、ガバナンス、共有が容易になります。 Unity Catalog 関数は、大規模なデータセットに変換と集計を適用する場合に特に適しています。
エージェントコードを使用してツールを作成する: このアプローチは、任意のコードやライブラリを使用して REST APIsを呼び出す場合や、非常に短いレイテンシでツールを実行する場合に便利です。ただし、このアプローチには、Unity Catalog 関数によって提供される組み込みの検出可能性とガバナンスが欠けています。エージェントを構築する際には、このトレードオフを比較検討して、どのアプローチが最適かを判断してください。

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

要件

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

ドキュメントによるツールコールの改善

AIエージェントは、ドキュメントを使用して、提供するツールをいつ、どのように使用するかを理解します。各ツールのパラメーターと戻り値に関する明確で詳細なドキュメントを提供します。

Unity Catalog 関数の場合は、 COMMENT を使用してツールを記述します。

Unity Catalog 関数を使用したツールの作成

次の例は、 Unity Catalog 関数を使用してエージェントツールを作成する方法を示しています。

Python エグゼキューターツール

この例では、Python コードを実行するツールを作成します。

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

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

%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 the output of the previous tool executions. No such variables, "rows", or "observations" were 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()
$$

テーブルクエリツール

この例では、Unity Catalog テーブルから顧客データをクエリするツールを作成します。

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

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

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;

ベクトル検索 retriever

この例では、 Mosaic AI Vector Search インデックスからデータをクエリする取得ツールを作成します。

Unity Catalog 関数は、Databricksドキュメントを含む架空のベクトル検索インデックスをクエリdatabricks_docs_vector_search。標準の vector_search() 関数をラップし、エイリアス page_content と metadata を使用して、その出力を MLflow 取得者スキーマと一致させます。

注：

MLflow 取得者スキーマに準拠するには、最上位の出力キーとしてではなく、SQL マップ関数を使用して、追加のメタデータ列を metadata 列に追加する必要があります。

ノートブックまたは SQL エディターで次のコードを実行します。

CREATE OR REPLACE FUNCTION main.default.databricks_docs_vector_search (
  -- The agent uses this comment to determine how to generate the query string parameter.
  query STRING
  COMMENT 'The query string for searching Databricks documentation.'
) RETURNS TABLE
-- The agent uses this comment to determine when to call this tool. It describes the types of documents and information contained within the index.
COMMENT 'Executes a search on Databricks documentation to retrieve text documents most relevant to the input query.' RETURN
SELECT
  chunked_text as page_content,
  map('doc_uri', url, 'chunk_id', chunk_id) as metadata
FROM
  vector_search(
    -- Specify your vector search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

このレトリーバーツールには、次の注意事項があります。

MLflow は、この Unity Catalog 関数を RETRIEVER スパンタイプではなく TOOL スパンタイプとしてトレースします。その結果、エージェントレビューアプリやAI PlaygroundなどのダウンストリームのAgent Frameworkアプリケーションには、チャンクへのリンクなど、レトリーバー固有の詳細が表示されません。スパンタイプの詳細については、「 MLflow トレーシングスキーマ」を参照してください。
SQL クライアントでは、返される行数またはバイト数の制限がある場合があります。データの切り捨てを防ぐには、UDF によって返される列値を切り捨てる必要があります。たとえば、 substring(chunked_text, 0, 8192) を使用して、大きなコンテンツ列のサイズを縮小し、実行中の行の切り捨てを回避できます。
このツールは vector_search() 関数のラッパーであるため、 vector_search() 関数と同じ制限が適用されます。制限事項を参照してください。

この例がユースケースに適していない場合は、代わりにカスタムエージェントコードを使用してベクトル検索レトリーバを作成してください。「ベクトル検索レトリーバーを作成する」を参照してください。

Unity Catalog ツールをエージェントに割り当てる

Unity Catalog ツールを作成したら、エージェントに割り当てて使用します。

LangChain エージェントは、UCFunctionToolkit を使用して Unity Catalog ツールを割り当てることができます。

次に、LangChainツール呼び出しエージェントをAI Playgroundからコードにエクスポートする方法について説明します。

AI Playground からのツールコールエージェントのエクスポート

AI Playground を使用して、Unity Catalog ツールを LLM に割り当て、エージェントをテストしてから、ツール呼び出しエージェントを定義するコードをエクスポートします。

AI Playground を使用してエージェントをエクスポートするには、ワークスペースが次の要件を満たしている必要があります。

Playground から、[ ツールの有効化 ] ラベルが付いたモデルを選択します。
[ツール] を選択し、[ツールの追加] をクリックします。
ドロップダウンメニューで、 Unity Catalog 機能を選択します。
プレイグラウンドを使用して、LLM、ツール、およびシステムプロンプトの現在の組み合わせをチャットしてテストします。バリエーションを試して、現在のセットアップがどのように機能するかを感じてください。

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

[ エクスポート ] をクリックして、エージェントを定義してデプロイする Python ノートブックを生成します。

エージェントコードをエクスポートすると、ワークスペースに保存された3つのファイルが表示されます。
- agent ノートブック: LangChainを使用してエージェントを定義するPythonコードが含まれています。
- driver ノートブック: エージェントフレームワークを使用してPython エージェントをログに記録、トレース、登録するデプロイするためのAI Mosaic AIコードが含まれています。
- config.yml: ツール定義など、エージェントに関する構成情報が含まれます。
agentノートブックを開いて、エージェントを定義するLangChainコードを確認します。このノートブックを使用して、エージェントをプログラムでテストおよび反復処理します (他のツールの定義など)。
エージェントの出力に問題がなければ、 driver ノートブックを実行して、エージェントをログに記録し、モデルサービングエンドポイントにデプロイできます。

エージェントコードを使用したツールの作成

また、Unity Catalog 関数ではなく、エージェントのコードでツールを作成することもできます。

このアプローチは、任意のコードまたはライブラリを使用して REST APIsを呼び出す場合、または低遅延ツールを実行する場合に便利です。ただし、エージェントのコードでツールを定義するには、Unity Catalog 関数によって提供される検出可能性とガバナンスが不足しています。

エージェントコードで定義されたベクトル検索取得ツールの例については、「ベクトル検索取得子を作成する」を参照してください。

次のステップ

エージェントを作成します。

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

要件