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

Unity Catalog 関数でカスタム AI エージェントツールを作成

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

必要条件​

  • の create function 文を使用して記述された 関数を作成するためのサーバレス コンピュート接続。Unity CatalogSQLPython 関数にはサーバレス コンピュートは必要ありません。
  • Databricks Runtime 15.0 以降を使用します。

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

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

依存関係のインストール​

[databricks] Extra を使用して 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 つのAPIsのいずれかを使用して作成できます。

  • 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を使用して関数をラップし、エージェントオーサリングライブラリからアクセスできるようにします。このツールキットは、異なるgen 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

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

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

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

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;

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

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

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

サーバレス mode for 本番運用​

サーバレスモードは、本番運用の使用例に推奨されるデフォルトのオプションです。 SQL サーバレスエンドポイントを使用してリモートで実行し、エージェントのプロセスの安全性を確保し、任意のコードをローカルで実行するリスクがないことを保証します。

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

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

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

開発用のローカルモード​

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

エージェントが ローカル ・モードでのツールの実行を要求すると、 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 時間が実時間を超える可能性があります。

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

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

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

環境変数​

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を使用して関数を実行するときに返される最大行数。

次のステップ​

に最終更新