Pular para o conteúdo principal

Criar ferramentas personalizadas do agente AI com funções Unity Catalog

Use as funções Unity Catalog para criar ferramentas de agente AI que executam lógica personalizada e realizam tarefas específicas que ampliam os recursos dos LLMs além da geração de linguagem.

Requisitos​

  • Uma conexão serverless compute para criar funções Unity Catalog escritas usando instruções de função de criação de corpo SQL. Python As funções não requerem serverless compute.
  • Use o Databricks Runtime 15.0 e o acima.

Crie uma ferramenta de agente​

Neste exemplo, o senhor cria uma ferramenta do Unity Catalog, testa sua funcionalidade e a adiciona a um agente. Execute o seguinte código em um Notebook Databricks.

Instalar dependências​

Instale o pacote Unity Catalog AI com o extra [databricks] e instale o pacote de integração Databricks-LangChain.

Este exemplo usa o site LangChain, mas uma abordagem semelhante pode ser aplicada a outras bibliotecas. Consulte Integrar ferramentas Unity Catalog com estruturas geradoras de AI de terceiros.

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

Inicializar o cliente da função Databricks​

Inicialize o Databricks Function Client, que é uma interface especializada para criar, gerenciar e executar funções do Unity Catalog no Databricks.

Python
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Defina a lógica da ferramenta​

Crie uma função do Unity Catalog que contenha a lógica da ferramenta.

O senhor pode criar funções Python usando uma das duas APIs:

  • create_python_function aceita um chamável Python.
  • create_function aceita uma instrução de função de criação de corpo SQL.

Use a API create_python_function para criar a função. Lembre-se de que há alguns requisitos para o uso bem-sucedido dessa API:

  • Dicas de tipo: A assinatura da função deve definir dicas de tipo Python válidas. Tanto os argumentos nomeados quanto o valor de retorno devem ter seus tipos definidos.
  • Não use argumentos variáveis: argumentos variáveis como *args e **kwargs não são suportados. Todos os argumentos devem ser definidos explicitamente.
  • Compatibilidade de tipos: Nem todos os tipos Python são compatíveis com o SQL.
  • Docstrings descritivos: O kit de ferramentas das funções Unity Catalog lê, analisa e extrai informações importantes de sua docstring.
    • As docstrings devem ser formatadas de acordo com a sintaxe docstring do Google.
    • Escreva descrições claras para sua função e seus argumentos para ajudar o LLM a entender como e quando usar a função.

Para obter mais informações, consulte a documentaçãoUnity Catalog - Creating functions from Python callables.

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
)

Teste a função​

Teste sua função para garantir que ela funcione conforme o esperado:

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'

Envolva a função usando o UCFunctionToolkit​

Envolva a função usando o endereço UCFunctionToolkit para torná-la acessível à biblioteca de autoria do agente. O kit de ferramentas garante a consistência entre diferentes bibliotecas e acrescenta recursos úteis, como o rastreamento automático para recuperadores.

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

Use a ferramenta em um agente​

Adicione a ferramenta a um agente LangChain usando a propriedade tools de UCFunctionToolkit.

Este exemplo cria um agente simples usando a API AgentExecutor da LangChain para simplificar. Para cargas de trabalho de produção, use o fluxo de trabalho de criação de agentes visto nos exemplos do siteChatAgent.

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

gerenciar Unity Catalog funções​

Use o Databricks Function Client para gerenciar as funções do Unity Catalog. O Databricks Function Client é baseado no Unity Catalog Function Client de código aberto, mas oferece vários aprimoramentos exclusivos da Databricks.

Esta página documenta a funcionalidade específica do Databricks Function Client. Para obter informações gerais sobre o gerenciamento das funções do Unity Catalog, consulte a documentação doUnity Catalog - Function client.

Executar funções usando serverless ou o modo local​

O senhor pode executar funções com o DatabricksFunctionClient em dois modos: modo serverless e modo local.

Especifique um nome de função totalmente qualificado no endereço execute_function API para executar a função. Quando um serviço gen AI determina que uma chamada de ferramenta é necessária para atender a uma solicitação, o pacote de integração (instâncias do kit de ferramentas) chama automaticamente esse API para executar a função.

modo sem servidor​

O modo sem servidor é a opção default e recomendada para casos de uso de produção. Ele executa funções remotamente usando um SQL serverless endpoint, garantindo que o processo do seu agente permaneça seguro e livre dos riscos de execução de código arbitrário localmente.

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

Quando seu agente solicita a execução de uma ferramenta no modo serverless acontece o seguinte:

  1. O site DatabricksFunctionClient envia uma solicitação ao Unity Catalog para recuperar a definição da função, caso a definição não tenha sido armazenada em cache localmente.
  2. O DatabricksFunctionClient extrai a definição da função e valida os nomes e tipos dos parâmetros.
  3. O DatabricksFunctionClient envia a execução como um UDF para uma instância serverless.

Modo local​

O modo local foi projetado para desenvolvimento e depuração. Ele executa funções em um subprocesso local em vez de fazer solicitações a um servidor sem servidor SQL endpoint. Isso permite que você solucione problemas de chamadas de ferramentas com mais eficiência, fornecendo rastreamentos de pilha locais.

Quando seu agente solicita a execução de uma ferramenta no modo local , o DatabricksFunctionClient faz o seguinte:

  1. Envia uma solicitação ao Unity Catalog para recuperar a definição da função se a definição não tiver sido armazenada em cache localmente.
  2. Extrai a definição do chamável Python e armazena em cache o chamável localmente, além de validar os nomes e tipos de parâmetros.
  3. Invoca o callable com os parâmetros especificados em um subprocesso restrito com proteção de tempo limite.
Python
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="local")

A execução no modo"local" oferece o seguinte recurso:

  • Limite de tempo de CPU: restringe o tempo de execução total da CPU para execução que pode ser chamada para evitar cargas computacionais excessivas.

    O limite de tempo da CPU é baseado no uso real da CPU, não no tempo do relógio de parede. Devido aos processos de programar e concorrente do sistema, o tempo da CPU pode exceder o tempo do relógio de parede em cenários reais.

  • Limite de memória: restringe a memória virtual alocada ao processo.

  • Proteção de tempo limite: impõe um tempo limite total do relógio de parede para executar funções.

Personalize esses limites usando a variável de ambiente (leia mais).

variável de ambiente​

Configure a execução das funções no site DatabricksFunctionClient usando as seguintes variáveis de ambiente:

Variável de ambiente

Valor padrão

Descrição

EXECUTOR_MAX_CPU_TIME_LIMIT

10 Segundos

Tempo máximo permitido de execução da CPU (somente no modo local).

EXECUTOR_MAX_MEMORY_LIMIT

100 MB

Alocação máxima permitida de memória virtual para o processo (somente no modo local).

EXECUTOR_TIMEOUT

20 Segundos

Tempo máximo total do relógio de parede (somente no modo local).

UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS

5

Número máximo de tentativas para atualizar novamente o cliente de sessão em caso de expiração de tokens.

UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT

100

Número máximo de linhas ao executar funções usando serverless compute com databricks-connect.

Última atualização em