Bagikan melalui

Membuat alat agen AI kustom dengan fungsi Unity Catalog

Gunakan fungsi Unity Catalog untuk membuat alat agen AI yang menjalankan logika kustom dan melakukan tugas tertentu yang memperluas kemampuan LLM di luar pembuatan bahasa.

Persyaratan

  • Koneksi komputasi tanpa server untuk membuat fungsi Katalog Unity yang ditulis menggunakan pernyataan fungsi buat isi SQL. Fungsi Python tidak memerlukan komputasi tanpa server.
  • Gunakan Databricks Runtime 15.0 ke atas.

Membuat perangkat agen

Dalam contoh ini, Anda membuat alat Unity Catalog, menguji fungsionalitasnya, dan menambahkannya ke agen. Jalankan kode berikut di buku catatan Databricks.

Pasang dependensi

Pasang paket AI Unity Catalog dengan tambahan [databricks] serta pasang paket integrasi Databricks-LangChain.

Contoh ini menggunakan LangChain, tetapi pendekatan serupa dapat diterapkan ke pustaka lain. Lihat Mengintegrasikan alat Unity Catalog dengan kerangka kerja AI generatif pihak ketiga.

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

Menginisialisasi Klien Fungsi Databricks

Inisialisasi Klien Fungsi Databricks, yang merupakan antarmuka khusus untuk membuat, mengelola, dan menjalankan fungsi Unity Catalog di Databricks.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Menentukan logika alat

Alat Katalog Unity sebenarnya hanyalah fungsi yang ditentukan pengguna (UDF) di balik layar. Saat Anda menentukan alat Katalog Unity, Anda mendaftarkan fungsi di Unity Catalog. Untuk mempelajari selengkapnya tentang Unity Catalog UDF, lihat Fungsi yang ditentukan pengguna (UDF) di Unity Catalog.

Anda dapat membuat fungsi Katalog Unity menggunakan salah satu dari dua API:

  • create_python_function menerima fungsi Python yang dapat dipanggil.
  • create_function menerima pernyataan pembuatan fungsi struktur SQL. Lihat Membuat fungsi Python.

create_python_function Gunakan API untuk membuat fungsi.

Agar fungsi Python yang dapat dipanggil diakui oleh model data fungsi Unity Catalog, fungsi tersebut harus memenuhi persyaratan berikut:

  • Petunjuk jenis: Tanda tangan fungsi harus menentukan petunjuk jenis Python yang valid. Argumen bernama dan nilai pengembalian harus memiliki jenis yang ditentukan.
  • Jangan gunakan argumen variabel: Argumen variabel seperti *args dan **kwargs tidak didukung. Semua argumen harus didefinisikan secara eksplisit.
  • Kompatibilitas jenis: Tidak semua jenis Python didukung di SQL. Lihat Jenis Data yang Didukung Spark.
  • Docstring yang deskriptif: Perangkat fungsi Unity Catalog membaca, mengurai, dan mengekstrak informasi penting dari docstring Anda.
    • Docstring harus diformat sesuai dengan sintaks Google docstring.
    • Tulis deskripsi yang jelas untuk fungsi Anda dan argumennya untuk membantu LLM memahami bagaimana dan kapan menggunakan fungsi.
  • Impor dependensi: Pustaka harus diimpor dalam isi fungsi. Impor di luar fungsi tidak akan diatasi saat menjalankan alat.

Cuplikan kode berikut menggunakan create_python_function untuk mendaftarkan callable 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
)

Uji fungsi

Uji fungsi Anda untuk memastikan berfungsi seperti yang diharapkan. Tentukan nama fungsi yang sepenuhnya memenuhi syarat di execute_function API untuk menjalankan fungsi:

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'

Membungkus fungsi menggunakan UCFunctionToolKit

Bungkus fungsi menggunakan UCFunctionToolkit untuk membuatnya dapat diakses oleh perpustakaan untuk penulisan agen. Toolkit memastikan konsistensi di berbagai pustaka AI generatif dan menambahkan fitur bermanfaat seperti pelacakan otomatis untuk pencari informasi.

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

Menggunakan alat di dalam agen

Tambahkan alat ke agen LangChain dengan menggunakan properti tools dari UCFunctionToolkit.

Contoh ini menulis agen sederhana menggunakan LangChain AgentExecutor API untuk kesederhanaan. Gunakan alur kerja perancangan agen yang terdapat dalam ChatAgent contoh untuk pekerjaan produksi.

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

Meningkatkan penggunaan alat dengan dokumentasi yang jelas

Dokumentasi yang baik membantu agen Anda mengetahui kapan dan bagaimana menggunakan setiap alat. Ikuti praktik terbaik ini untuk mendokumen alat Anda:

  • Untuk fungsi Katalog Unity, gunakan COMMENT klausul untuk menjelaskan fungsionalitas dan parameter alat.
  • Tentukan input dan output yang diharapkan dengan jelas.
  • Tulis deskripsi yang bermakna untuk mempermudah alat bagi agen, dan manusia, untuk digunakan.

Contoh: Dokumentasi alat yang efektif

Contoh berikut menunjukkan string yang jelas COMMENT untuk alat yang meminta tabel terstruktur.

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;

Contoh: Dokumentasi alat yang tidak efektif

Contoh berikut tidak memiliki detail penting, sehingga lebih sulit bagi agen untuk menggunakan alat ini secara efektif:

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;

Menjalankan fungsi menggunakan mode tanpa server atau lokal

Ketika layanan AI generatif menentukan bahwa pemanggilan fungsi diperlukan, paket integrasi (instans UCFunctionToolkit) menjalankan API DatabricksFunctionClient.execute_function.

Panggilan execute_function dapat menjalankan fungsi dalam dua mode eksekusi: tanpa server atau lokal. Mode ini menentukan sumber daya mana yang menjalankan fungsi.

Mode tanpa server untuk produksi

Mode tanpa server adalah opsi default dan direkomendasikan untuk kasus penggunaan produksi. Ini menjalankan fungsi dari jarak jauh menggunakan titik akhir tanpa server SQL, memastikan bahwa proses agen Anda tetap aman dan bebas dari risiko menjalankan kode acak secara lokal.

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

Saat agen Anda meminta eksekusi alat dalam mode tanpa server , hal berikut ini terjadi:

  1. DatabricksFunctionClient mengirim permintaan ke Katalog Unity untuk mengambil definisi fungsi jika definisi belum di-cache secara lokal.
  2. DatabricksFunctionClient mengekstrak definisi fungsi dan memvalidasi nama dan jenis parameter.
  3. DatabricksFunctionClient mengirimkan eksekusi sebagai UDF ke instans tanpa server.

Mode lokal untuk pengembangan

Mode lokal dirancang untuk pengembangan dan debugging. Ini menjalankan fungsi dalam subproses lokal alih-alih membuat permintaan ke titik akhir SQL Serverless. Ini memungkinkan Anda memecahkan masalah panggilan alat secara lebih efektif dengan menyediakan jejak tumpukan lokal.

Saat agen Anda meminta menjalankan alat dalam mode lokal , DatabricksFunctionClient lakukan hal berikut:

  1. Mengirim permintaan ke Katalog Unity untuk mengambil definisi fungsi jika definisi belum di-cache secara lokal.
  2. Mengekstrak definisi yang dapat dipanggil Python, menyimpan cache yang dapat dipanggil secara lokal, dan memvalidasi nama dan jenis parameter.
  3. Memanggil fungsi yang dapat dipanggil dengan parameter yang ditentukan dalam subproses terbatas dengan pelindung batas waktu.
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="local")

Berjalan dalam "local" mode menyediakan fitur berikut:

  • Batas waktu CPU: Membatasi total runtime CPU untuk eksekusi yang dapat dipanggil untuk mencegah beban komputasi yang berlebihan.

    Batas waktu CPU didasarkan pada penggunaan CPU aktual, bukan waktu jam dinding. Karena penjadwalan sistem dan proses bersamaan, waktu CPU dapat melebihi waktu jam dinding dalam skenario dunia nyata.

  • Batas memori: Membatasi memori virtual yang dialokasikan untuk proses.

  • Perlindungan batas waktu: Memberlakukan batas waktu keseluruhan untuk menjalankan fungsi.

Sesuaikan batas ini menggunakan variabel lingkungan (baca lebih lanjut).

Variabel lingkungan

Konfigurasikan bagaimana fungsi berjalan dalam DatabricksFunctionClient menggunakan variabel lingkungan berikut:

Variabel lingkungan Nilai bawaan Deskripsi
EXECUTOR_MAX_CPU_TIME_LIMIT 10 detik Waktu eksekusi CPU maksimum yang diizinkan (hanya mode lokal).
EXECUTOR_MAX_MEMORY_LIMIT 100 MB Alokasi memori virtual maksimum yang diizinkan untuk proses (hanya mode lokal).
EXECUTOR_TIMEOUT 20 detik Total waktu jam dinding maksimum (hanya mode lokal).
UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS 5 Jumlah maksimum upaya untuk mencoba menyegarkan kembali klien sesi jika token kedaluwarsa.
UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT 100 Jumlah maksimum baris yang akan dikembalikan saat menjalankan fungsi menggunakan komputasi tanpa server dan databricks-connect.

Langkah selanjutnya