Projektowanie niestandardowych narzędzi agenta AI za pomocą funkcji katalogu Unity

2025-05-14

Funkcje Unity Catalog umożliwiają tworzenie narzędzi agenta sztucznej inteligencji, które realizują niestandardową logikę i wykonują określone zadania, rozszerzając możliwości LLMs poza generowaniem języka.

Wymagania

Bezserwerowa obsługa obliczeń do tworzenia funkcji katalogu Unity, utworzonych za pomocą instrukcji tworzenia funkcji SQL. Funkcje języka Python nie wymagają przetwarzania bezserwerowego.
Użyj środowiska Databricks Runtime 15.0 lub nowszego.

Stwórz narzędzie agenta

W tym przykładzie utworzysz narzędzie Katalogu Unity, przetestujesz jego funkcjonalność i dodasz je do agenta. Uruchom następujący kod w notatniku usługi Databricks.

Instalowanie zależności

Zainstaluj pakiety AI Unity Catalog z dodatkiem [databricks] i zainstaluj pakiet integracyjny Databricks-LangChain.

W tym przykładzie użyto języka LangChain, ale podobne podejście można zastosować do innych bibliotek. Zobacz narzędzia Unity Catalog z ramami generatywnej sztucznej inteligencji innych firm.

# 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()

Zainicjuj klienta funkcji Databricks

Zainicjuj klienta funkcji Databricks, który jest wyspecjalizowanym interfejsem służącym do tworzenia, zarządzania i uruchamiania funkcji w Unity Catalog w usłudze Databricks.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Definiowanie logiki narzędzia

Narzędzia katalogu Unity to naprawdę tylko funkcje zdefiniowane przez użytkownika katalogu Unity (UDF, User-Defined Functions) w tle. Podczas definiowania narzędzia Unity Catalog rejestrujesz funkcję w Unity Catalog. Aby dowiedzieć się więcej o funkcjach zdefiniowanych przez użytkownika (UDF) w katalogu Unity, zobacz Funkcje zdefiniowane przez użytkownika (UDF) w katalogu Unity.

Funkcje katalogu Unity można utworzyć przy użyciu jednego z dwóch interfejsów API.

create_python_function akceptuje język Python, który można wywołać.
create_function akceptuje instrukcję SQL create function. Zobacz Tworzenie funkcji języka Python.

Użyj interfejsu create_python_function API, aby utworzyć funkcję.

Aby wywołanie w Pythonie było rozpoznawalne w modelu danych funkcji Katalogu Unity, funkcja musi spełniać następujące wymagania:

Wskazówki dotyczące typów: podpis funkcji musi definiować prawidłowe wskazówki dotyczące typu języka Python. Zarówno nazwane argumenty, jak i wartość zwracana muszą mieć zdefiniowane typy.
Nie używaj argumentów zmiennych: argumenty zmiennych, takie jak *args i **kwargs, nie są obsługiwane. Wszystkie argumenty muszą być jawnie zdefiniowane.
Zgodność typów: nie wszystkie typy języka Python są obsługiwane w języku SQL. Zobacz Typy danych obsługiwanych przez platformę Spark.
Opisowe docstrings: zestaw narzędzi funkcji Unity Catalog odczytuje, analizuje i wyodrębnia ważne informacje z docstringu.
- Docstringi muszą być sformatowane zgodnie ze składnią docstringów Google.
- Napisz jasne opisy swojej funkcji i jej argumentów, aby pomóc LLM zrozumieć, jak i kiedy używać funkcji.
Import zależności: biblioteki muszą być importowane w treści funkcji. Importy poza funkcją nie zostaną rozwiązane podczas uruchamiania narzędzia.

Następujące fragmenty kodu używają create_python_function do zarejestrowania funkcji do wywołania języka 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
)

Testowanie funkcji

Przetestuj funkcję, aby sprawdzić, czy działa zgodnie z oczekiwaniami. Określ w pełni kwalifikowaną nazwę funkcji w interfejsie API execute_function, aby uruchomić funkcję.

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'

Owiń funkcję za pomocą zestawu UCFunctionToolKit

Opakuj funkcję przy użyciu elementu UCFunctionToolkit, aby była dostępna dla bibliotek tworzenia agentów. Zestaw narzędzi zapewnia spójność w różnych bibliotekach generatywnej sztucznej inteligencji i dodaje przydatne funkcje, takie jak automatyczne śledzenie dla narzędzi wyszukiwania.

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

Korzystanie z narzędzia w agencie

Dodaj narzędzie do agenta LangChain przy użyciu właściwości tools z UCFunctionToolkit.

W tym przykładzie tworzony jest prosty agent korzystający z interfejsu API LangChain AgentExecutor dla uproszczenia. W przypadku obciążeń produkcyjnych użyj przepływu pracy tworzenia agenta, jak pokazano w ChatAgent przykładach.

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?"})

Ulepszanie wywoływania narzędzi za pomocą przejrzystej dokumentacji

Dobra dokumentacja pomaga agentom wiedzieć, kiedy i jak używać każdego narzędzia. Postępuj zgodnie z poniższymi najlepszymi rozwiązaniami dotyczącymi dokumentowania narzędzi:

W przypadku funkcji Unity Catalog użyj klauzuli COMMENT, aby opisać parametry i funkcjonalność narzędzia.
Jasno zdefiniuj oczekiwane dane wejściowe i wyjściowe.
Napisz znaczące opisy, aby ułatwić agentom i ludziom korzystanie z narzędzi.

Przykład: Dokumentacja skutecznego narzędzia

W poniższym przykładzie pokazano jasne COMMENT ciągi dla narzędzia, które wysyła zapytanie do tabeli ustrukturyzowanej.

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;

Przykład: dokumentacja narzędzia nieskutecznego

W poniższym przykładzie brakuje ważnych szczegółów, co utrudnia agentom efektywne korzystanie z narzędzia:

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;

Uruchamianie funkcji przy użyciu trybu bezserwerowego lub lokalnego

Gdy usługa generatywnej sztucznej inteligencji określa, że potrzebne jest wywołanie narzędzia, pakiety integracji (UCFunctionToolkit wystąpienia) uruchamiają interfejs API DatabricksFunctionClient.execute_function.

Wywołanie execute_function może uruchamiać funkcje w dwóch trybach wykonywania: bezserwerowych lub lokalnych. Ten tryb określa, który zasób uruchamia funkcję.

Tryb bezserwerowy dla środowiska produkcyjnego

Tryb bezserwerowy jest domyślną i zalecaną opcją dla przypadków użycia w środowisku produkcyjnym. Uruchamia funkcje zdalnie przy użyciu bezserwerowego punktu końcowego SQL, zapewniając, że proces agenta pozostaje bezpieczny i wolny od ryzyka wykonywania dowolnego kodu w lokalnym środowisku.

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

Gdy agent żąda wykonania narzędzia w trybie bezserwerowym , wykonywane są następujące czynności:

Obiekt DatabricksFunctionClient wysyła żądanie do Unity Catalog w celu pobrania definicji funkcji, jeśli definicja nie została zapisana w lokalnej pamięci podręcznej.
Funkcja DatabricksFunctionClient wyodrębnia definicję funkcji i weryfikuje nazwy parametrów i typy.
Funkcja DatabricksFunctionClient przesyła wykonanie jako funkcję zdefiniowaną przez użytkownika do wystąpienia bezserwerowego.

Tryb lokalny na potrzeby rozwoju

Tryb lokalny jest przeznaczony do programowania i debugowania. Wykonuje funkcje w lokalnym podprocesie zamiast wysyłać żądania do punktu końcowego bezserwerowego programu SQL Server. Dzięki temu można efektywniej rozwiązywać problemy z wywołaniami narzędzi, zapewniając ślady stosu lokalnego.

Gdy agent żąda uruchomienia narzędzia w trybie lokalnym , DatabricksFunctionClient wykonuje następujące czynności:

Wysyła żądanie do katalogu Unity, aby pobrać definicję funkcji, jeśli definicja nie została lokalnie zapisana w pamięci podręcznej.
Wyodrębnia definicję wywoływalnej funkcji w języku Python, buforuje tę funkcję lokalnie i sprawdza nazwy oraz typy parametrów.
Wywołuje obiekt wywoływalny z określonymi parametrami w ograniczonym podprocesie z zabezpieczeniem czasowym.

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

Uruchamianie w "local" trybie zapewnia następujące funkcje:

Limit czasu procesora: Ogranicza całkowity czas działania procesora dla wywoływalnego wykonania, aby zapobiec nadmiernemu obciążeniu obliczeniowemu.

Limit czasu procesora CPU zależy od rzeczywistego użycia procesora CPU, a nie czasu zegara. Ze względu na planowanie systemu i współbieżne procesy czas procesora CPU może przekraczać czas zegara w rzeczywistych scenariuszach.
Limit pamięci: Ogranicza pamięć wirtualną przydzieloną do procesu.
Ochrona limitu czasu: Wymusza całkowity limit czasu zegara ściany dla uruchomionych funkcji.

Dostosuj te limity przy użyciu zmiennych środowiskowych (przeczytaj więcej).

Zmienne środowiskowe

Skonfiguruj sposób działania funkcji w programie DatabricksFunctionClient przy użyciu następujących zmiennych środowiskowych:

Zmienna środowiskowa	Wartość domyślna	Opis
`EXECUTOR_MAX_CPU_TIME_LIMIT`	`10` Sekund	Maksymalny dozwolony czas wykonywania procesora CPU (tylko tryb lokalny).
`EXECUTOR_MAX_MEMORY_LIMIT`	`100` MB	Maksymalna dozwolona alokacja pamięci wirtualnej dla procesu (tylko tryb lokalny).
`EXECUTOR_TIMEOUT`	`20` Sekund	Maksymalny całkowity czas zegara ściany (tylko tryb lokalny).
`UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS`	`5`	Maksymalna liczba prób ponawiania próby odświeżenia klienta sesji w przypadku wygaśnięcia tokenu.
`UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT`	`100`	Maksymalna liczba wierszy, które zostaną zwrócone podczas uruchamiania funkcji z użyciem obliczeń bezserwerowych i `databricks-connect`.

Dalsze kroki

Programistycznie dodawanie narzędzi Katalogu Unity do agentów. Zobacz ChatAgent przykłady.
Dodaj narzędzia Unity Catalog do agentów, używając interfejsu użytkownika AI Playground. Zobacz prototypowe narzędzia wywołujące agentów w środowisku AI Playground.
Zarządzanie funkcjami Unity Catalog przy użyciu klienta Function. Zobacz dokumentację Unity Catalog - klient funkcji

Udostępnij za pośrednictwem