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

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 関数を作成します。

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

  • create_python_function Python の呼び出し可能オブジェクトを受け入れます。
  • create_function SQL 本体の create function ステートメントを受け入れます。

create_python_function API を使用して関数を作成します。このAPIを正常に使用するには、いくつかの要件があることに注意してください。

  • タイプヒント: 関数シグネチャは、有効なPython型ヒントを定義する必要があります。名前付き引数と戻り値の両方に、その型が定義されている必要があります。
  • 変数引数は使用しないでください。 *args や **kwargs などの変数引数はサポートされていません。すべての引数は明示的に定義する必要があります。
  • 型の互換性: すべての Python 型が SQL でサポートされているわけではありません。
  • 記述的な docstring: Unity Catalog 関数ツールキットは、docstring から重要な情報を読み取って解析し、抽出します。
    • docstring は、Google docstring 構文に従ってフォーマットする必要があります。
    • 関数とその引数に明確な説明を記述して、LLM が関数をいつどのように使用するかを理解できるようにします。

詳細については、 Unity Catalog ドキュメント - Python callable からの関数の作成を参照してください。

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
)

関数をテストする​

関数をテストして、期待どおりに動作することを確認します。

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

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 の機能を管理する​

Databricks Function Client を使用して、Unity Catalog 関数を管理します。Databricks Function Client は、オープンソース の Unity Catalog Function Client に基づいていますが、Databricks独自の改善点がいくつかあります。

このページでは、Databricks Function Client に固有の機能について説明します。Unity Catalog 関数の管理に関する一般的な情報については、 Unity Catalog のドキュメンテーション - 関数クライアントを参照してください。

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

DatabricksFunctionClientを使用して関数は、サーバレスモードとローカルモードの2つのモードで実行できます。

execute_function API で完全修飾関数名を指定して、関数を実行します。Gen AI サービスがリクエストを満たすためにツール呼び出しが必要であると判断すると、統合パッケージ (ツールキットインスタンス) は自動的にこの API を呼び出して関数を実行します。

サーバレスモード​

サーバレスモードは、本番運用の使用例に推奨されるデフォルトのオプションです。 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

サーバレス コンピュート with databricks-connectを使用して関数を実行する場合の最大行数 .

に最終更新