Integrar ferramentas do Unity Catalog com frameworks de IA de terceiros

As ferramentas agentes do Unity Catalog podem ser usadas em bibliotecas de IA populares como LangChain, LlamaIndex, OpenAI e Anthropic. Essas integrações combinam a governança da ferramenta Unity Catalog com os recursos das estruturas de criação de agentes de terceiros. Por exemplo:

  • No LangChain, as funções do Unity Catalog podem fazer parte do fluxo de trabalho de um agente para executar tarefas como consultar ou transformar dados.
  • Em integrações OpenAI ou Anthropic, as funções são chamadas diretamente pelo modelo de IA durante a execução.

Selecione o seu framework nos separadores seguintes para criar uma ferramenta Unity Catalog e usá-lo com esse framework. Executa o código num notebook Azure Databricks ou num script Python.

Requirements

  • Instale o Python 3.10 ou superior.

LangChain

Use o Azure Databricks Unity Catalog para integrar funções SQL e Python como ferramentas nos fluxos de trabalho LangChain e LangGraph. Essa integração combina a governança do Unity Catalog com os recursos LangChain para criar aplicativos poderosos baseados em LLM.

Neste exemplo, você cria uma ferramenta Unity Catalog, testa sua funcionalidade e a adiciona a um agente.

Instalar dependências

Instale os pacotes Unity Catalog AI com o Databricks opcional e instale o pacote de integração LangChain.

# Install the Unity Catalog AI integration package with the Databricks extra
%pip install unitycatalog-langchain[databricks]

# Install Databricks Langchain integration package
%pip install databricks-langchain
dbutils.library.restartPython()

Inicializar o cliente da função Databricks

Inicialize o cliente da função Databricks.

from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

Definir a lógica da ferramenta

Crie uma função Unity Catalog contendo a lógica da ferramenta.


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:

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 a consistência entre diferentes bibliotecas 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.

Este exemplo cria um agente simples usando a API do AgentExecutor LangChain para simplificar. Para cargas de trabalho de produção, use o fluxo de trabalho de autoria de agentes visto em Criar um agente e implemente-o nas aplicações Databricks.

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 (replace with your LLM of choice, if desired)
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?"})

LlamaIndex

Use o Azure Databricks Unity Catalog para integrar funções SQL e Python como ferramentas nos fluxos de trabalho do LlamaIndex. Essa integração combina a governança do Unity Catalog com os recursos do LlamaIndex para indexar e consultar grandes conjuntos de dados para LLMs.

  1. Instale o pacote de integração do Databricks Unity Catalog para LlamaIndex.

    %pip install unitycatalog-llamaindex[databricks]
dbutils.library.restartPython()

  2. Crie uma instância do cliente de funções do Unity Catalog.

    from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

  3. Crie uma função Unity Catalog escrita em Python.

    CATALOG = "your_catalog"
SCHEMA = "your_schema"

func_name = f"{CATALOG}.{SCHEMA}.code_function"

def code_function(code: str) -> str:
  """
  Runs Python code.

  Args:
    code (str): The Python code to run.
  Returns:
    str: The result of running the Python code.
  """
  import sys
  from io import StringIO
  stdout = StringIO()
  sys.stdout = stdout
  exec(code)
  return stdout.getvalue()

client.create_python_function(
  func=code_function,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True
)

  4. Crie uma instância da função Unity Catalog como um kit de ferramentas e execute-a para verificar se a ferramenta se comporta corretamente.

    from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit
import mlflow

# Enable traces
mlflow.llama_index.autolog()

# Create a UCFunctionToolkit that includes the UC function
toolkit = UCFunctionToolkit(function_names=[func_name])

# Fetch the tools stored in the toolkit
tools = toolkit.tools
python_exec_tool = tools[0]

# Run the tool directly
result = python_exec_tool.call(code="print(1 + 1)")
print(result)  # Outputs: {"format": "SCALAR", "value": "2\n"}

  5. Use a ferramenta em um LlamaIndex ReActAgent definindo a função Unity Catalog como parte de uma coleção de ferramentas LlamaIndex. Em seguida, verifique se o agente se comporta corretamente chamando a coleção de ferramentas LlamaIndex.

    from llama_index.llms.openai import OpenAI
from llama_index.core.agent import ReActAgent

llm = OpenAI()

agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)

agent.chat("Please run the following python code: `print(1 + 1)`")

OpenAI

Use o Azure Databricks Unity Catalog para integrar funções SQL e Python como ferramentas nos fluxos de trabalho da OpenAI. Esta integração combina a governação do Unity Catalog com a OpenAI para criar aplicações de IA poderosas.

  1. Instale o pacote de integração do Databricks Unity Catalog para OpenAI.

    %pip install unitycatalog-openai[databricks]
%pip install mlflow -U
dbutils.library.restartPython()

  2. Crie uma instância do cliente de funções do Unity Catalog.

    from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

  3. Crie uma função Unity Catalog escrita em Python.

    CATALOG = "your_catalog"
SCHEMA = "your_schema"

func_name = f"{CATALOG}.{SCHEMA}.code_function"

def code_function(code: str) -> str:
  """
  Runs Python code.

  Args:
    code (str): The python code to run.
  Returns:
    str: The result of running the Python code.
  """
  import sys
  from io import StringIO
  stdout = StringIO()
  sys.stdout = stdout
  exec(code)
  return stdout.getvalue()

client.create_python_function(
  func=code_function,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True
)

  4. Crie uma instância da função Unity Catalog como um kit de ferramentas e verifique se a ferramenta se comporta corretamente executando a função.

    from unitycatalog.ai.openai.toolkit import UCFunctionToolkit
import mlflow

# Enable tracing
mlflow.openai.autolog()

# Create a UCFunctionToolkit that includes the UC function
toolkit = UCFunctionToolkit(function_names=[func_name])

# Fetch the tools stored in the toolkit
tools = toolkit.tools
client.execute_function = tools[0]

  5. Envie a solicitação para o modelo OpenAI junto com as ferramentas.

    import openai

messages = [
  {
    "role": "system",
    "content": "You are a helpful customer support assistant. Use the supplied tools to assist the user.",
  },
  {"role": "user", "content": "What is the result of 2**10?"},
]
response = openai.chat.completions.create(
  model="gpt-4o-mini",
  messages=messages,
  tools=tools,
)
# check the model response
print(response)

  6. Depois que o OpenAI retornar uma resposta, invoque a chamada da função Unity Catalog para gerar a resposta de volta para o OpenAI.

    import json

# OpenAI sends only a single request per tool call
tool_call = response.choices[0].message.tool_calls[0]
# Extract arguments that the Unity Catalog function needs to run
arguments = json.loads(tool_call.function.arguments)

# Run the function based on the arguments
result = client.execute_function(func_name, arguments)
print(result.value)

  7. Uma vez que a resposta tenha sido retornada, você pode construir a carga útil de resposta para chamadas subsequentes para OpenAI.

    # Create a message containing the result of the function call
function_call_result_message = {
  "role": "tool",
  "content": json.dumps({"content": result.value}),
  "tool_call_id": tool_call.id,
}
assistant_message = response.choices[0].message.to_dict()
completion_payload = {
  "model": "gpt-4o-mini",
  "messages": [*messages, assistant_message, function_call_result_message],
}

# Generate final response
openai.chat.completions.create(
  model=completion_payload["model"], messages=completion_payload["messages"]
)

Utilities

Para simplificar o processo de elaboração da resposta da ferramenta, o ucai-openai pacote tem um utilitário, generate_tool_call_messages, que converte mensagens de resposta OpenAI ChatCompletion para que elas possam ser usadas para geração de resposta.

from unitycatalog.ai.openai.utils import generate_tool_call_messages

messages = generate_tool_call_messages(response=response, client=client)
print(messages)

Observação

Se a resposta contiver entradas de múltipla escolha, você poderá passar o argumento choice_index ao chamar generate_tool_call_messages para escolher qual entrada de escolha utilizar. Atualmente, não há suporte para o processamento de entradas de múltipla escolha.

Anthropic

Use o Azure Databricks Unity Catalog para integrar funções SQL e Python como ferramentas nas chamadas LLM do Anthropic SDK. Esta integração combina a governação do Unity Catalog com modelos Anthropic para criar aplicações de IA poderosas.

Observação

A integração com o Anthropic requer o Databricks Runtime 15.0 ou superior.

  1. Instale o pacote de integração do Databricks Unity Catalog para Anthropic.

    %pip install unitycatalog-anthropic[databricks]
dbutils.library.restartPython()

  2. Crie uma instância do cliente de funções do Unity Catalog.

    from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

  3. Crie uma função Unity Catalog escrita em Python.

    CATALOG = "your_catalog"
SCHEMA = "your_schema"

func_name = f"{CATALOG}.{SCHEMA}.weather_function"

def weather_function(location: str) -> str:
  """
  Fetches the current weather from a given location in degrees Celsius.

  Args:
    location (str): The location to fetch the current weather from.
  Returns:
    str: The current temperature for the location provided in Celsius.
  """
  return f"The current temperature for {location} is 24.5 celsius"

client.create_python_function(
  func=weather_function,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True
)

  4. Crie uma instância da função Unity Catalog como um kit de ferramentas.

    from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit

# Create an instance of the toolkit
toolkit = UCFunctionToolkit(function_names=[func_name], client=client)

  5. Use uma chamada de ferramenta em Anthropic.

    import anthropic

# Initialize the Anthropic client with your API key
anthropic_client = anthropic.Anthropic(api_key="YOUR_ANTHROPIC_API_KEY")

# User's question
question = [{"role": "user", "content": "What's the weather in New York City?"}]

# Make the initial call to Anthropic
response = anthropic_client.messages.create(
  model="claude-3-5-sonnet-20240620",  # Specify the model
  max_tokens=1024,  # Use 'max_tokens' instead of 'max_tokens_to_sample'
  tools=toolkit.tools,
  messages=question  # Provide the conversation history
)

# Print the response content
print(response)

  6. Construa uma resposta de ferramenta. A resposta do modelo Claude contém um bloco de metadados de solicitação de ferramenta se uma ferramenta precisar ser chamada.

    from unitycatalog.ai.anthropic.utils import generate_tool_call_messages

# Call the UC function and construct the required formatted response
tool_messages = generate_tool_call_messages(
  response=response,
  client=client,
  conversation_history=question
)

# Continue the conversation with Anthropic
tool_response = anthropic_client.messages.create(
  model="claude-3-5-sonnet-20240620",
  max_tokens=1024,
  tools=toolkit.tools,
  messages=tool_messages,
)

print(tool_response)

O unitycatalog.ai-anthropic pacote inclui um utilitário manipulador de mensagens para simplificar a análise e o tratamento de uma chamada para a função Unity Catalog. O utilitário faz o seguinte:

  1. Deteta os requisitos de chamada da ferramenta.
  2. Extrai informações sobre a chamada da ferramenta a partir da consulta.
  3. Executa a chamada para a função Unity Catalog.
  4. Analisa a resposta da função Unity Catalog.
  5. Crie o próximo formato de mensagem para continuar a conversa com Claude.

Observação

Todo o histórico da conversa deve ser fornecido no conversation_history argumento para a generate_tool_call_messages API. Os modelos Claude exigem a inicialização da conversa (a pergunta de entrada original do usuário) e todas as respostas subsequentes geradas pelo LLM e os resultados de chamadas de ferramentas consecutivas.

  • Last updated on