使用 Unity 目錄函式來建立 AI 代理程式工具,以執行自定義邏輯,並執行特定工作,以擴充超越語言世代的 LLM 功能。
何時使用 Unity 目錄函式與 MCP 伺服器
Databricks 建議在查詢事先已知且代理提供參數時,使用 Unity Catalog 函式作為代理工具,專門用於結構化資料檢索工具。 請參見 「連接代理與結構化資料」。
在大多數其他使用情境中,Databricks 建議使用 MCP 伺服器或直接在代理程式中定義邏輯,以加快執行速度、支援每用戶驗證及增加彈性。
需求
若要建立和使用 Unity 目錄函式作為 AI 代理程式工具,您需要下列專案:
- Databricks Runtime:使用 Databricks Runtime 15.0 和更新版本
- Python 版本:安裝 Python 3.10 或更新版本
若要執行 Unity 目錄函式:
- 必須在您的工作區中啟用無伺服器計算,才能在生產環境中執行 Unity 目錄函式作為 AI 代理程式工具。 請參閱 無伺服器運算需求。
- Python 函式的本機模式執行不需要無伺服器泛型計算才能執行,不過本機模式只適用於開發和測試用途。
若要建立 Unity 目錄函式:
- 您的工作區中必須啟用無伺服器泛型計算,才能使用 Databricks Workspace Client 或 SQL 主體語句來建立函式。
- 不需要無伺服器計算即可建立 Python 函式。
建立代理程式工具
在此範例中,您會建立 Unity 目錄工具、測試其功能,並將其新增至代理程式。 在 Databricks 筆記本中執行下列程式代碼。
安裝依賴項
安裝具有 [databricks] 額外功能的 Unity 目錄 AI 套件,並安裝 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 Catalog 工具其實本質上只是 Unity Catalog 使用者定義函數 (UDF)。 當您定義 Unity 目錄工具時,您會在 Unity 目錄中註冊函式。 若要深入瞭解 Unity 目錄 UDF,請參閱 Unity 目錄中的使用者定義函式 (UDF)。
警告
在代理程式工具中執行任意程式碼可能會暴露代理程式有權存取的敏感或私人資訊。 客戶只需執行受信任的程式碼,並設定防護欄及適當權限,以防止資料被誤導存取。
您可以使用兩個 API 的其中一個來建立 Unity 目錄函式:
-
create_python_function接受 Python 可調用對象。 -
create_function接受 SQL 主體 create 函式語句。 請參閱 建立 Python 函式。
使用 create_python_function API 來建立函式。
若要讓 Python 函式能被 Unity Catalog 函數數據模型辨識,您的函式必須符合下列條件:
- 類型提示:函式簽章必須定義有效的 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
)
測試函式
測試您的函式,以檢查它是否如預期般運作。 在 API execute_function 中指定完整的函式名稱以運行函式:
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 函式庫的一致性,並新增實用功能,例如檢索器的自動追蹤。
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。 不過,您可以將 Unity 目錄工具與其他架構整合,例如 LlamaIndex、OpenAI、Anthropic 等。 請參閱 整合 Unity 目錄工具與第三方產生 AI 架構。
本範例使用 LangChain AgentExecutor API 撰寫簡單的代理程式,以求簡單。 對於生產工作負載,請使用 Author a AI Agent 中看到的代理編寫工作流程 ,並部署到 Databricks 應用程式上。
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 可以在兩種執行模式中執行函式:無伺服器或本機。 此模式會決定執行函式的資源。
生產環境的無伺服器模式
無伺服器模式是執行 Unity 目錄函式做為 AI 代理程式工具時,生產使用案例的預設和建議選項。 此模式使用無伺服器的通用計算平台(Spark Connect 無伺服器模式)來遠端執行函式,而 Lakeguard 確保代理程式的安全,避免本地執行任意程式碼的風險。
備註
以 AI 代理程式工具身分執行的 Unity 目錄函式需要無伺服器泛型計算(Spark Connect 無伺服器),而不是無伺服器 SQL 倉儲。 嘗試在沒有無伺服器泛型計算的情況下執行工具會產生類似的錯誤 PERMISSION_DENIED: Cannot access Spark Connect。
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="serverless")
當您的代理程式要求 在無伺服器 模式中執行工具時,會發生下列情況:
-
DatabricksFunctionClient將要求傳送至 Unity Catalog,以擷取函式定義,如果尚未在本機快取該定義。 -
DatabricksFunctionClient會擷取函式定義,並驗證參數名稱和類型。 - 會將
DatabricksFunctionClient執行以 UDF 的形式提交至無伺服器泛型計算。
用於開發的本機模式
本機模式會在本機子程式中執行 Python 函式,而不是對無伺服器泛型計算提出要求。 這可讓您藉由提供本地堆棧追蹤,更有效率地對工具調用進行疑難解答。 其設計目的是開發及偵錯 Python Unity 目錄函式。
當您的代理程式要求以 本機 模式執行工具時,會 DatabricksFunctionClient 執行下列動作:
- 將請求傳送至 Unity Catalog,以擷取函式定義,若此定義尚未在本機快取。
- 擷取 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執行函式時要傳回的數據列數目上限。 |
後續步驟
以程序設計方式將 Unity 目錄工具新增至代理程式。 請參閱 建立 AI 代理並將其部署到 Databricks 應用程式中。
使用 AI 遊樂場 UI 將 Unity Catalog 工具新增至代理程式。 請參見 :不需編碼即可查詢大型語言模型 (LLM) 並開發原型 AI 代理程式。
使用函式用戶端管理 Unity 目錄函式。 請參閱 Unity Catalog 文件 - 函數客戶端