Unity カタログ関数を使用してカスタム AI エージェントツールを作成する

2025-05-14

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

要求事項

SQL 本文の create function ステートメントを使用して記述された Unity Catalog 関数を作成するためのサーバーレスコンピューティング接続。 Python 関数では、サーバーレスコンピューティングは必要ありません。
Databricks Runtime 15.0 以降を使用します。

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

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

依存関係のインストール

Unity Catalog AI パッケージを [databricks] エクストラ付きでインストールし、Databricks-LangChain 統合パッケージをインストールします。

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

# 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 関数クライアントを初期化します。これは、Databricks で Unity カタログ関数を作成、管理、および実行するための特殊なインターフェイスです。

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

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

Unity カタログツールは、実際には Unity カタログのユーザー定義関数 (UDF) です。 Unity カタログツールを定義すると、Unity カタログに関数が登録されます。 Unity カタログ UDF の詳細については、 Unity カタログのユーザー定義関数 (UDF) に関するページを参照してください。

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

create_python_function は Pythonのコール可能オブジェクトを受け付けます。
create_function は、SQL 本文の作成関数ステートメントを受け入れます。 Python 関数の作成を参照してください。

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

Unity カタログの関数データモデルにおいて Python の呼び出し可能を認識させるには、関数が次の要件を満たしている必要があります。

型ヒント: 関数シグネチャは、有効な Python 型ヒントを定義する必要があります。名前付き引数と戻り値の両方に、型が定義されている必要があります。
変数引数を使用しないでください。*args や **kwargs などの変数引数はサポートされていません。すべての引数を明示的に定義する必要があります。
型の互換性: SQL では、すべての Python 型がサポートされているわけではありません。 Spark でサポートされるデータ型を参照してください。
説明ドキュメント文字列: Unity カタログ関数ツールキットは、docstring から重要な情報を読み取り、解析、抽出します。
- Docstring は、Google docstring 構文に従って書式設定する必要があります。
- LLM が関数を使用する方法とタイミングを理解するのに役立つ、関数とその引数の明確な説明を記述します。
依存関係のインポート: ライブラリは関数の本体内にインポートする必要があります。関数外のインポートは、ツールの実行時に解決されません。

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


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 で完全修飾関数名を指定して、関数を実行します。

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を使用して関数をラップし、エージェント作成ライブラリからアクセスできるようにします。このツールキットは、さまざまな Gen AI ライブラリ間の一貫性を確保し、レトリバーの自動トレースなどの便利な機能を追加します。

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

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

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

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

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 カタログ関数の場合は、 COMMENT 句を使用してツールの機能とパラメーターを記述します。
予想される入力と出力を明確に定義します。
エージェントや人間が使いやすいツールを作成するためのわかりやすい説明を記述します。

例: 有効なツールのドキュメント

次の例は、構造化テーブルに対してクエリを実行するツールの明確な COMMENT 文字列を示しています。

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;

例: 無効なツールのドキュメント

次の例では、重要な詳細が不足しているため、エージェントがツールを効果的に使用することが困難になります。

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;

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

Gen AI サービスがツール呼び出しが必要であると判断すると、統合パッケージ (UCFunctionToolkit インスタンス) によって DatabricksFunctionClient.execute_function API が実行されます。

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

運用環境のサーバーレスモード

サーバーレスモードは、運用環境のユースケースの既定の推奨オプションです。 SQL Server レスエンドポイントを使用してリモートで関数を実行し、エージェントのプロセスを安全に保ち、任意のコードをローカルで実行するリスクを防ぐことができます。

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

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

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

開発用のローカルモード

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

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

定義がローカルにキャッシュされていない場合は、関数定義を取得する要求を Unity カタログに送信します。
Python 呼び出し可能な定義を抽出し、呼び出し可能な定義をローカルにキャッシュし、パラメーターの名前と型を検証します。
指定されたパラメーターを使用し、タイムアウト保護付きの制限されたサブプロセス内で、呼び出し可能な関数を呼び出します。

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

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

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

CPU 時間制限は、実際の CPU 使用率に基づいており、ウォールクロック時間ではありません。システムスケジューリングと同時実行プロセスにより、CPU 時間は実世界のシナリオで壁時計時間を超える可能性があります。
メモリ制限: プロセスに割り当てられた仮想メモリを制限します。
タイムアウト保護: 関数を実行するための合計ウォールクロックタイムアウトを適用します。

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

環境変数

次の環境変数を使用して、 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`を使用して関数を実行するときに返す最大行数。

次のステップ

プログラムを使用して Unity カタログツールをエージェントに追加します。 ChatAgent 例を参照してください。
AI Playground UI を使用してエージェントに Unity カタログツールを追加します。 AI Playground のプロトタイプツール呼び出しエージェント参照してください。
Function Client を使用して Unity カタログ関数を管理します。 Unity カタログのドキュメントを参照する - 関数クライアント

次の方法で共有

Unity カタログ関数を使用してカスタム AI エージェント ツールを作成する