Skapa AI-agentverktyg med hjälp av Unity Catalog-funktioner

Använd Unity Catalog-funktioner för att skapa AI-agentverktyg som kör anpassad logik och utför specifika uppgifter som utökar funktionerna i LLMs utöver språkgenerering.

Kravspecifikation

För att skapa och använda Unity Catalog-funktioner som AI-agentverktyg behöver du följande:

• Databricks Runtime: Använda Databricks Runtime 15.0 och senare
• Python-version: Installera Python 3.10 eller senare

Så här kör du Unity Catalog-funktioner:

• Serverlös beräkning måste aktiveras på din arbetsyta för att köra Unity Catalog-funktioner som AI-agentverktyg i produktion. Se Krav för serverlös beräkning.
  • Körning av lokalt läge för Python-funktioner kräver inte serverlös allmän beräkning, men lokalt läge är endast avsett för utveckling och testning.

Så här skapar du Unity Catalog-funktioner:

• Serverlös allmän beräkning måste vara aktiverad på arbetsytan för att skapa funktioner med hjälp av Databricks-arbetsyteklienten eller SQL-brödtextinstruktionerna.
  • Python-funktioner kan skapas utan serverlös beräkning.

Skapa ett agentverktyg

I det här exemplet skapar du ett Unity Catalog-verktyg, testar dess funktioner och lägger till det i en agent. Kör följande kod i en Databricks-notebook-fil.

Installera beroenden

Installera UNITY Catalog AI-paket med extra [databricks] och installera Databricks-LangChain integrationspaketet.

I det här exemplet används LangChain, men en liknande metod kan tillämpas på andra bibliotek. Se Integrera Unity Catalog-verktyg med generativa AI-ramverk från tredje part.

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

Initiera Databricks-funktionsklienten

Initiera Databricks-funktionsklienten, som är ett specialiserat gränssnitt för att skapa, hantera och köra Unity Catalog-funktioner i Databricks.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Definiera verktygets logik

Unity Catalog-verktyg är egentligen bara användardefinierade funktioner i Unity Catalog (UDF: er) under huven. När du definierar ett Unity Catalog-verktyg registrerar du en funktion i Unity Catalog. Mer information om UDF:er för Unity-katalogen finns i Användardefinierade funktioner (UDF:er) i Unity Catalog.

Du kan skapa Unity Catalog-funktioner med något av två API:er:

• create_python_function accepterar ett anropbart Python-objekt.
• create_function accepterar en CREATE FUNCTION-sats för SQL-kropp. Se Skapa Python-funktioner.

Använd API:et create_python_function för att skapa funktionen.

För att en Python-anropbar funktion ska kunna identifieras av Unity Catalog-funktionsdatamodellen måste funktionen uppfylla följande krav:

• Typtips: Funktionssignaturen måste definiera giltiga tips av Python-typ. Både de namngivna argumenten och returvärdet måste ha sina typer definierade.
• Använd inte variabelargument: Variabelargument som *args och **kwargs stöds inte. Alla argument måste definieras uttryckligen.
• Typkompatibilitet: Alla Python-typer stöds inte i SQL. Se Spark-datatyper som stöds.
• Beskrivande dokumentsträngar: Verktygslådan för Unity Catalog-funktioner läser, parsar och extraherar viktig information från din dokumentsträng.
  • Docstrings måste formateras enligt Googles docstring-syntax.
  • Skriv tydliga beskrivningar för din funktion och dess argument för att hjälpa LLM att förstå hur och när funktionen ska användas.
• Beroendeimport: Bibliotek måste importeras i funktionens kropp. Importeringar utanför funktionen kommer inte att lösas när verktyget körs.

Följande kodfragment använder create_python_function för att registrera python-anropsbara 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
)

Testa funktionen

Testa funktionen för att kontrollera att den fungerar som förväntat. Ange ett fullständigt kvalificerat funktionsnamn i API:et execute_function för att köra funktionen:

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'

Omslut funktionen med hjälp av UCFunctionToolKit

Omslut funktionen med hjälp av UCFunctionToolkit för att göra den tillgänglig för agentredigeringsbibliotek. Verktygslådan säkerställer konsekvens i olika generativa AI-bibliotek och lägger till användbara funktioner som automatisk spårning för sökfunktioner.

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

Använd verktyget i en agent

Lägg till verktyget till en LangChain-agent med hjälp av egenskapen tools från UCFunctionToolkit.

Det här exemplet skapar en enkel agent med hjälp av LangChain AgentExecutor API för enkelhet. För produktionsarbetsbelastningar använder du det agentredigeringsarbetsflöde som visas i ResponsesAgent exempel.

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

Förbättra användning av verktyg med tydlig dokumentation

Bra dokumentation hjälper dina agenter att veta när och hur de ska använda varje verktyg. Följ dessa metodtips för att dokumentera dina verktyg:

• För Unity Catalog-funktioner använder du COMMENT -satsen för att beskriva verktygsfunktioner och parametrar.
• Definiera tydligt förväntade indata och utdata.
• Skriv meningsfulla beskrivningar för att göra verktyg enklare för agenter och människor att använda.

Exempel: Dokumentation om effektiva verktyg

I följande exempel visas tydliga COMMENT strängar för ett verktyg som frågar en strukturerad tabell.

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;

Exempel: Dokumentation om ineffektivt verktyg

I följande exempel saknas viktig information, vilket gör det svårare för agenter att använda verktyget effektivt:

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;

Köra funktioner med serverlöst eller lokalt läge

När en generativ AI-tjänst bestämmer att ett verktygsanrop behövs kör integrationspaket (UCFunctionToolkit instanser) DatabricksFunctionClient.execute_function-API:et.

Anropet execute_function kan köra funktioner i två körningslägen: serverlös eller lokal. Det här läget avgör vilken resurs som kör funktionen.

Serverlöst läge för produktion

Serverlöst läge är standardalternativet och rekommenderas för produktionsanvändningsfall när du kör Unity Catalog-funktioner som AI-agentverktyg. Det här läget använder serverlös allmän beräkning (Serverlös Spark Connect) för att fjärrköra funktioner, vilket säkerställer att agentens process förblir säker och fri från riskerna med att köra godtycklig kod lokalt.

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

När din agent begär en verktygskörning i serverlöst läge händer följande:

1. DatabricksFunctionClient Skickar en begäran till Unity Catalog för att hämta funktionsdefinitionen om definitionen inte har cachelagrats lokalt.
2. Extraherar DatabricksFunctionClient funktionsdefinitionen och validerar parameternamnen och typerna.
3. DatabricksFunctionClient skickar körningen som en UDF till serverlös allmän beräkningskapacitet.

Lokalt läge för utveckling

Lokalt läge kör Python-funktioner i en lokal underprocess i stället för att göra begäranden till serverlös allmän beräkning. På så sätt kan du felsöka verktygsanrop mer effektivt genom att tillhandahålla lokala stackspårningar. Den är utformad för att utveckla och felsöka Python Unity Catalog-funktioner.

När din agent begär att ett verktyg ska köras i lokalt läge DatabricksFunctionClient gör du följande:

1. Skickar en begäran till Unity Catalog för att hämta funktionsdefinitionen om definitionen inte har cachelagrats lokalt.
2. Extraherar den anropsbara Python-definitionen, cachelagrar anropsbar lokalt och validerar parameternamnen och typerna.
3. Anropar anropsbar med de angivna parametrarna i en begränsad underprocess med timeout-skydd.

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

När du kör i "local" läget finns följande funktioner:

• Cpu-tidsgräns: Begränsar den totala CPU-körningen för anropsbar körning för att förhindra överdriven beräkningsbelastning.

  Cpu-tidsgränsen baseras på den faktiska CPU-användningen, inte tid på väggklockan. På grund av systemplanering och samtidiga processer kan CPU-tiden överskrida kalendertiden i verkliga scenarier.

• Minnesgräns: Begränsar det virtuella minne som allokeras till processen.

• Timeout-skydd: Tvingar fram en total väggklockstidsgräns för att köra funktioner.

Anpassa dessa gränser med hjälp av miljövariabler (läs vidare).

Begränsningar i lokalt läge

• Endast Python-funktioner: SQL-baserade funktioner stöds inte i lokalt läge.
• Säkerhetsöverväganden för ej betrodd kod: Även om lokalt läge kör funktioner i en underprocess för processisolering finns det en potentiell säkerhetsrisk när godtycklig kod som genereras av AI-system körs. Detta är främst ett problem när funktioner kör dynamiskt genererad Python-kod som inte har granskats.
• Skillnader i biblioteksversioner: Biblioteksversioner kan skilja sig mellan serverlösa och lokala körningsmiljöer, vilket kan leda till olika funktionsbeteenden.

Miljövariabler

Konfigurera hur funktioner körs i DatabricksFunctionClient med hjälp av följande miljövariabler:

Nästa steg

• Lägg till Unity Catalog-verktyg till agenter programmatiskt. Se ResponsesAgent exempel.

• Lägg till Unity Catalog-verktyg till agenter med hjälp av AI Playground-användargränssnittet. Se även Prototypverktygsanropande agenter i AI Playground.

• Hantera Unity Catalog-funktioner med hjälp av funktionsklienten. Se Dokumentation om Unity Catalog – funktionsklient