メインコンテンツまでスキップ

Unity Catalog関数を使用してAIエージェント ツールを作成する (レガシー)

注記

Databricks では、エージェント ツールに MCP サーバーの使用を推奨しています。Python Unity Catalog関数の場合、カスタム MCP サーバーに移行するか、エージェント コードでロジックを直接定義して、実行を高速化し、柔軟性を高めます。

決定論的ルックアップ用のSQL Unity Catalog関数は有効なパターンのままです。 「エージェントを構造化データに接続する」を参照してください。

Unity Catalog 関数を使用して、カスタムロジックを実行し、言語生成を超えて LLM の機能を拡張する特定のタスクを実行する AI エージェントツールを作成します。

必要条件

Unity Catalog 関数を AI エージェントツールとして作成して使用するには、次のものが必要です。

  • Databricks Runtime : Databricks Runtime 15.0 以降を使用します
  • Pythonバージョン :Python 3.10以降をインストールします

Unity Catalog 関数を実行するには:

  • 本番運用でAIエージェントツールとしてUnity Catalog機能を実行するには、ワークスペースでサーバレス コンピュート を有効にする必要があります。サーバレス コンピュートの要件を参照してください。
    • Python関数のローカル・モード実行には、サーバレス generic コンピュートは必要ありませんが、ローカル・モードは開発およびテストのみを目的としています。

Unity Catalog関数を作成するには:

  • Databricks Workspace Client または SQL body 文を使用して関数を作成するには、ワークスペースで サーバレス generic コンピュート を有効にする必要があります。
    • Python 関数はサーバレス コンピュートなしで作成できます。

エージェントツールを作成する

この例では、Unity Catalog ツールを作成し、その機能をテストして、エージェントに追加します。Databricks ノートブックで次のコードを実行します。

依存関係のインストール

[databricks] 追記文 を使用して Unity Catalog AI パッケージをインストールし、Databricks と LangChain の統合パッケージをインストールします。

この例では LangChain を使用していますが、同様のアプローチを他のライブラリにも適用できます。Unity Catalogツールとサードパーティの生成AIフレームワークとの統合を参照してください。

Python
# Install Unity Catalog AI integration packages with the Databricks extra
%pip install unitycatalog-ai[databricks]
%pip install unitycatalog-langchain[databricks]

# Install the Databricks LangChain integration package
%pip install databricks-langchain

dbutils.library.restartPython()

Databricks 関数クライアントを初期化する

Databricks で Unity Catalog 関数を作成、管理、実行するための特殊なインターフェイスである Databricks Function Client を初期化します。

Python
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

ツールのロジックを定義する

Unity Catalog ツールは、実際には Unity Catalog ユーザー定義関数 (UDF) にすぎません。Unity Catalogツールを定義すると、Unity Catalogに関数が登録されます。Unity Catalog UDF の詳細については、Unity Catalog のユーザー定義関数 (UDF)を参照してください。

警告

エージェント ツールで任意のコードを実行すると、エージェントがアクセスできる機密情報や個人情報が公開される可能性があります。顧客は、信頼できるコードのみを実行し、データへの意図しないアクセスを防ぐためのガードレールと適切な権限を構成する責任を負います。

Unity Catalog関数は、次の 2 つのAPIのいずれかを使用して作成できます。

  • create_python_function Python の呼び出し可能オブジェクトを受け入れます。
  • create_function SQL 本体の create function ステートメントを受け入れます。Python 関数の作成を参照してください。

create_python_function API を使用して関数を作成します。

Python の呼び出し可能オブジェクトを Unity Catalog 関数データ モデルで認識できるようにするには、関数が次の要件を満たしている必要があります。

  • 型ヒント : 関数シグネチャは、有効な Python 型ヒントを定義する必要があります。名前付き引数と戻り値の両方に、その型が定義されている必要があります。

  • 変数引数を使用しない : *args や **kwargs などの変数引数はサポートされていません。すべての引数は明示的に定義する必要があります。

  • 型の互換性 : すべての Python 型が SQL でサポートされているわけではありません。Spark でサポートされるデータ型を参照してください。

  • 記述的な docstrings : Unity Catalog 関数ツールキットは、docstring から重要な情報を読み取り、解析し、抽出します。

    • docstring は、Google docstring 構文に従ってフォーマットする必要があります。
    • 関数とその引数に明確な説明を記述して、LLM が関数をいつどのように使用するかを理解できるようにします。
  • 依存関係のインポート : ライブラリは、関数の本体内でインポートする必要があります。関数の外部からのインポートは、ツールの実行時に解決されません。

次のコード スニペットでは、 create_python_function を使用して、 Python 呼び出し可能 add_numbersを登録します。

Python

CATALOG = "my_catalog"
SCHEMA = "my_schema"

def add_numbers(number_1: float, number_2: float) -> float:
"""
A function that accepts two floating point numbers adds them,
and returns the resulting sum as a float.

Args:
number_1 (float): The first of the two numbers to add.
number_2 (float): The second of the two numbers to add.

Returns:
float: The sum of the two input numbers.
"""
return number_1 + number_2

function_info = client.create_python_function(
func=add_numbers,
catalog=CATALOG,
schema=SCHEMA,
replace=True
)

関数をテストする

関数をテストして、期待どおりに動作することを確認します。execute_function API で完全修飾関数名を指定して、関数を実行します。

Python
result = client.execute_function(
function_name=f"{CATALOG}.{SCHEMA}.add_numbers",
parameters={"number_1": 36939.0, "number_2": 8922.4}
)

result.value # OUTPUT: '45861.4'

UCFunctionToolKit を使用して関数をラップします

UCFunctionToolkitを使用して関数をラップし、エージェントオーサリングライブラリからアクセスできるようにします。このツールキットは、異なる生成AI ライブラリ間での一貫性を確保し、レトリーバーの自動トレースなどの便利な機能を追加します。

Python
from databricks_langchain import UCFunctionToolkit

# Create a toolkit with the Unity Catalog function
func_name = f"{CATALOG}.{SCHEMA}.add_numbers"
toolkit = UCFunctionToolkit(function_names=[func_name])

tools = toolkit.tools

エージェントでツールを使用する

UCFunctionToolkittools プロパティを使用して、ツールを LangChain エージェントに追加します。

注記

この例では、LangChainを使用します。ただし、Unity Catalog ツールを LlamaIndex、OpenAI、Anthropic などの他のフレームワークと統合することはできます。Unity Catalogツールをサードパーティ製AIフレームワークと統合するを参照してください。

この例では、わかりやすくするために LangChain AgentExecutor API を使用して単純なエージェントを作成します。本番運用ワークロードの場合は、 ResponsesAgentの例に示されているエージェント作成ワークフローを使用します。

Python
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.prompts import ChatPromptTemplate
from databricks_langchain import (
ChatDatabricks,
UCFunctionToolkit,
)
import mlflow

# Initialize the LLM (optional: replace with your LLM of choice)
LLM_ENDPOINT_NAME = "databricks-meta-llama-3-3-70b-instruct"
llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME, temperature=0.1)

# Define the prompt
prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"You are a helpful assistant. Make sure to use tools for additional functionality.",
),
("placeholder", "{chat_history}"),
("human", "{input}"),
("placeholder", "{agent_scratchpad}"),
]
)

# Enable automatic tracing
mlflow.langchain.autolog()

# Define the agent, specifying the tools from the toolkit above
agent = create_tool_calling_agent(llm, tools, prompt)

# Create the agent executor
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
agent_executor.invoke({"input": "What is 36939.0 + 8922.4?"})

明確なドキュメントでツールコールを改善

優れたドキュメントは、エージェントが各ツールをいつ、どのように使用するかを知るのに役立ちます。ツールを文書化するための次のベストプラクティスに従ってください。

  • Unity Catalog 関数の場合は、 COMMENT 句を使用して、ツールの機能とパラメーターを記述します。
  • 期待される入力と出力を明確に定義します。
  • エージェントや人間がツールを使いやすくするために、意味のある説明を記述します。

例: 効果的なツールの文書化

次の例は、構造化テーブルをクエリするツールのクリア COMMENT 文字列を示しています。

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 specific customer including their email and ID.'
RETURN SELECT CONCAT(
'Customer ID: ', customer_id, ', ',
'Customer Email: ', customer_email
)
FROM main.default.customer_data
WHERE customer_name = customer_name
LIMIT 1;

例: 効果のないツールの文書化

次の例では、重要な詳細が欠けているため、エージェントがツールを効果的に使用するのが難しくなっています。

SQL
CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
customer_name STRING COMMENT 'Name of the customer.'
)
RETURNS STRING
COMMENT 'Returns info about a customer.'
RETURN SELECT CONCAT(
'Customer ID: ', customer_id, ', ',
'Customer Email: ', customer_email
)
FROM main.default.customer_data
WHERE customer_name = customer_name
LIMIT 1;

サーバレスまたはローカルモードを使用した関数の実行

生成AI サービスがツール呼び出しが必要であると判断すると、統合パッケージ (UCFunctionToolkit インスタンス) は DatabricksFunctionClient.execute_function API を実行します。

execute_function 呼び出しは、サーバレスとローカルの 2 つの実行モードで関数を実行できます。このモードは、関数を実行するリソースを決定します。

本番運用におけるサーバレスモード

サーバーレス モードは、 Unity Catalog機能をAIエージェント ツールとして実行する場合の本番運用ユースケースにとって、安心かつ推奨されるオプションです。 このモードでは、サーバレスの汎用コンピュート ( Spark Connect サーバレス) を使用して機能をリモートで実行します。Lakeguard、エージェントのプロセスが安全に保たれ、ローカルで任意のコードを実行するリスクから解放されることを保証します。

注記

AIエージェントツールとして実行されるUnity Catalog機能には、サーバレス SQLウェアハウスではなく、サーバレス generic コンピュート (Spark Connect サーバレス) が必要です。サーバレス generic コンピュートなしでツールを実行しようとすると、 PERMISSION_DENIED: Cannot access Spark Connectのようなエラーが発生します。

Python
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="serverless")

エージェントが サーバレス モードでツールの実行を要求すると、次の処理が行われます。

  1. DatabricksFunctionClient は、定義がローカルにキャッシュされていない場合に関数定義を取得するように Unity Catalog に要求を送信します。
  2. DatabricksFunctionClientは、関数定義を抽出し、パラメーターの名前と型を検証します。
  3. DatabricksFunctionClientは、実行をUDFとしてサーバレス generic コンピュートに送信します。

開発用のローカルモード

ローカルモードは、サーバレス generic コンピュートにリクエストを行う代わりに、ローカルサブプロセスで Python 関数を実行します。 これにより、ローカルスタックトレースを提供することで、ツール呼び出しのトラブルシューティングをより効果的に行うことができます。これは、 Python Unity Catalog 関数の開発とデバッグ用に設計されています。

エージェントが ローカル モードでのツールの実行を要求すると、 DatabricksFunctionClient は次の処理を行います。

  1. 関数定義がローカルにキャッシュされていない場合は、関数定義を取得するための要求を Unity Catalog に送信します。
  2. Python 呼び出し可能オブジェクト定義を抽出し、呼び出し可能オブジェクトをローカルにキャッシュし、パラメーターの名前と型を検証します。
  3. 指定されたパラメーターを使用して、タイムアウト保護付きの制限付きサブプロセスで呼び出し可能オブジェクトを呼び出します。
Python
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="local")

"local" モードで実行すると、次の機能が提供されます。

  • CPU 時間制限: 呼び出し可能実行の合計 CPU ランタイムを制限して、過剰な計算負荷を防ぎます。

    CPU 時間制限は、実時間ではなく、実際の CPU 使用率に基づいています。システムのスケジューリングと並列プロセスにより、実際のシナリオでは CPU 時間が実時間を超える可能性があります。

  • メモリ制限: プロセスに割り当てられる仮想メモリを制限します。

  • タイムアウト保護: 実行中の関数に対して合計ウォールクロックタイムアウトを強制します。

環境変数を使用してこれらの制限をカスタマイズします(詳細を参照)。

ローカルモードの制限

  • Python 関数のみ : SQL ベースの関数は、ローカル モードではサポートされていません。
  • 信頼できないコードのセキュリティに関する考慮事項 : ローカル モードでは、プロセスを分離するためにサブプロセスで関数が実行されますが、AI システムによって生成された任意のコードを実行すると、潜在的なセキュリティ リスクが発生します。これは主に、関数がレビューされていない動的に生成された Python コードを実行する場合に問題になります。
  • ライブラリのバージョンの違い : サーバレスとローカル実行環境ではライブラリのバージョンが異なる場合があり、関数の動作が異なる可能性があります。

環境変数

DatabricksFunctionClientでの関数の実行方法を設定するには、次の環境変数を使用します。

環境変数

デフォルト値

説明

EXECUTOR_MAX_CPU_TIME_LIMIT

10

最大許容 CPU 実行時間 (ローカル モードのみ)。

EXECUTOR_MAX_MEMORY_LIMIT

100 MB

プロセスに許容される仮想メモリの最大割り当て (ローカル モードのみ)。

EXECUTOR_TIMEOUT

20

最大合計ウォールクロック時間(ローカルモードのみ)。

UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS

5

トークンの有効期限が切れた場合にセッションクライアントの更新を再試行する最大試行回数。

UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT

100

サーバレス コンピュート と databricks-connectを使用して関数を実行するときに返される最大行数。

次のステップ