Partilhar via

Crie ferramentas de agente de IA usando as funções do Unity Catalog

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

Advertência

A execução de código arbitrário em uma ferramenta de agente pode expor informações confidenciais ou privadas às quais o agente tem acesso. Os clientes são responsáveis por executar apenas código confiável e implementar guarda-corpos e permissões adequadas para evitar o acesso não intencional aos dados.

Requerimentos

Para criar e usar as funções do Unity Catalog como ferramentas de agente de IA, você precisa do seguinte:

  • Databricks Runtime: Use o Databricks Runtime 15.0 e superior
  • Versão do Python: Instale o Python 3.10 ou superior

Para executar as funções do Catálogo Unity:

  • A computação sem servidor deve ser habilitada em seu espaço de trabalho para executar funções do Unity Catalog como ferramentas de agente de IA em produção. Consulte Requisitos de computação sem servidor.
    • A execução em modo local para funções Python não requer computação genérica sem servidor para ser executada, no entanto, o modo local destina-se apenas a fins de desenvolvimento e teste.

Para criar funções do Catálogo Unity:

  • A computação genérica sem servidor deve ser habilitada em seu espaço de trabalho para criar funções usando as instruções Databricks Workspace Client ou SQL body.
    • As funções Python podem ser criadas sem computação sem servidor.

Criar uma ferramenta de agente

Neste exemplo, você cria uma ferramenta Unity Catalog, testa sua funcionalidade e a adiciona a um agente. Execute o código a seguir em um bloco de anotações Databricks.

Instalar dependências

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

Este exemplo usa LangChain, mas uma abordagem semelhante pode ser aplicada a outras bibliotecas. Consulte Integrar ferramentas do Unity Catalog com estruturas de IA generativas de terceiros.

# 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.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Definir a lógica da ferramenta

As ferramentas do Unity Catalog são, na verdade, apenas funções definidas pelo utilizador do Unity Catalog, na sua essência. Ao definir uma ferramenta Unity Catalog, você está registrando uma função no Unity Catalog. Para saber mais sobre UDFs do Unity Catalog, consulte Funções definidas pelo usuário (UDFs) no Unity Catalog.

Você pode criar funções do Unity Catalog usando uma das duas APIs:

  • create_python_function aceita um objeto chamável em Python.
  • create_function aceita uma instrução SQL para criar uma função. Consulte Criar funções Python.

Use a create_python_function API para criar a função.

Para tornar uma função Python chamada reconhecível para o modelo de dados de funções do Unity Catalog, a sua função deve atender aos seguintes requisitos:

  • Dicas de tipo: A assinatura de 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 explicitamente definidos.
  • Compatibilidade de tipos: Nem todos os tipos Python são suportados em SQL. Consulte Tipos de dados suportados pelo Spark.
  • Docstrings descritivos: O kit de ferramentas de funções do Unity Catalog lê, processa e extrai informações importantes dos seus docstrings.
    • Docstrings devem ser formatados 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.
  • Importações de dependência: As bibliotecas devem ser importadas dentro do corpo da função. As importações fora da função não serão resolvidas ao executar a ferramenta.

Os trechos de código a seguir usam o create_python_function para registrar o Python chamável 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
)

Testar a função

Teste sua função para verificar se ela funciona conforme o esperado. Especifique um nome de função totalmente qualificado execute_function na API para executar a função:

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 UCFunctionToolkit para torná-la acessível às bibliotecas de desenvolvimento de agentes. O kit de ferramentas garante consistência entre diferentes bibliotecas de IA de geração e adiciona recursos úteis, como rastreamento automático para retrievers.

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

Usar a ferramenta em um agente

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

Observação

Este exemplo usa LangChain. No entanto, você pode integrar as ferramentas do Unity Catalog com outras estruturas, como LlamaIndex, OpenAI, Anthropic e muito mais. Consulte Integrar ferramentas do Unity Catalog com estruturas de IA generativas de terceiros.

Este exemplo cria um agente simples usando a API LangChain AgentExecutor para simplificar. Para operações de produção, utilize o fluxo de construção do agente visto nos ResponsesAgent exemplos.

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

Melhore a ativação de ferramentas com documentação clara

Uma boa documentação ajuda seus agentes a saber quando e como usar cada ferramenta. Siga estas práticas recomendadas para documentar suas ferramentas:

  • Para as funções do Unity Catalog, use a cláusula COMMENT para descrever a funcionalidade da ferramenta e os parâmetros.
  • Definir claramente as entradas e saídas esperadas.
  • Escreva descrições significativas para tornar as ferramentas mais fáceis de usar por agentes e humanos.

Exemplo: Documentação eficaz da ferramenta

O exemplo a seguir mostra cadeias de caracteres claras COMMENT para uma ferramenta que consulta uma tabela estruturada.

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;

Exemplo: documentação ineficaz da ferramenta

O exemplo a seguir carece de detalhes importantes, dificultando o uso eficaz da ferramenta pelos agentes:

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;

Executando funções usando o modo sem servidor ou local

Quando um serviço de IA de geração determina que uma chamada de ferramenta é necessária, os pacotes de integração (UCFunctionToolkit instâncias) executam a DatabricksFunctionClient.execute_function API.

A execute_function chamada pode executar funções em dois modos de execução: sem servidor ou local. Este modo determina qual recurso executa a função.

Modo sem servidor para produção

O modo sem servidor é a opção padrão e recomendada para casos de uso de produção ao executar funções do Unity Catalog como ferramentas de agente de IA. Esse modo usa computação genérica sem servidor (Spark Connect serverless) para executar funções remotamente, garantindo que o processo do seu agente permaneça seguro e livre dos riscos de executar código arbitrário localmente.

Observação

As funções do Unity Catalog executadas como ferramentas de agente de IA exigem computação genérica sem servidor (Spark Connect sem servidor), não armazéns SQL sem servidor. Tentativas de executar ferramentas sem computação genérica sem servidor produzirão erros como PERMISSION_DENIED: Cannot access Spark Connect.

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

Quando o agente solicita a execução de uma ferramenta no modo sem servidor , acontece o seguinte:

  1. O DatabricksFunctionClient envia uma solicitação ao Unity Catalog para recuperar a definição de função se a definição não tiver sido armazenada localmente em cache.
  2. O DatabricksFunctionClient extrai a definição da função e valida os nomes e tipos de parâmetros.
  3. O DatabricksFunctionClient envia a execução como um UDF para computação genérica sem servidor.

Modo local de desenvolvimento

O modo local executa funções Python em um subprocesso local em vez de fazer solicitações para computação genérica sem servidor. Isso permite que você solucione problemas de chamadas de ferramentas de forma mais eficaz, fornecendo rastreamentos de pilha locais. Ele é projetado para desenvolver e depurar funções Python Unity Catalog.

Quando o 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 de função se a definição não tiver sido armazenada em cache localmente.
  2. Extrai a definição chamável do Python, armazena em cache o chamável localmente e valida os nomes e tipos de parâmetros.
  3. Invoca o chamável com os parâmetros especificados em um subprocesso restrito com proteção de tempo limite.
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="local")

A execução no "local" modo fornece os seguintes recursos:

  • Limite de tempo da CPU: Restringe o tempo de execução total da CPU para execução chamável 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 ao agendamento do sistema e processos simultâneos, o tempo da CPU pode exceder o tempo de relógio de parede em cenários do mundo real.

  • Limite de memória: Restringe a memória virtual alocada para o processo.

  • Proteção contra tempo limite: Impõe um tempo limite total de relógio de parede para funções em execução.

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

Limitações do modo local

  • Somente funções Python: funções baseadas em SQL não são suportadas no modo local.
  • Considerações de segurança para código não confiável: enquanto o modo local executa funções em um subprocesso para isolamento de processos, há um risco potencial de segurança ao executar código arbitrário gerado por sistemas de IA. Isso é principalmente uma preocupação quando as funções executam código Python gerado dinamicamente que não foi revisado.
  • Diferenças de versão da biblioteca: As versões da biblioteca podem diferir entre ambientes de execução sem servidor e locais, o que pode levar a um comportamento de função diferente.

Variáveis de ambiente

Configure como as funções são executadas usando as seguintes variáveis de DatabricksFunctionClient ambiente:

Variável de ambiente Valor predefinido Descrição
EXECUTOR_MAX_CPU_TIME_LIMIT 10 segundos Tempo máximo permitido de execução da CPU (apenas no modo local).
EXECUTOR_MAX_MEMORY_LIMIT 100 MB Alocação máxima de memória virtual permitida para o processo (somente no modo local).
EXECUTOR_TIMEOUT 20 segundos Tempo total máximo do relógio de parede (apenas no modo local).
UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS 5 O número máximo de tentativas de atualizar novamente o cliente de sessão em caso de expiração do token.
UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT 100 O número máximo de linhas a serem retornadas ao executar funções usando computação sem servidor e databricks-connect.

Próximos passos

  • Last updated on