Integrar ferramentas do catálogo Unity com estruturas de IA generativas de terceiros.

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

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

Selecione sua estrutura nas guias a seguir para criar uma ferramenta do Catálogo do Unity e usá-la com essa estrutura. Execute o código em um notebook Azure Databricks ou Python script.

Requirements

  • Instale o Python 3.10 ou superior.

LangChain

Use o Unity Catalog do Azure Databricks para integrar funções SQL e Python como ferramentas em fluxos de trabalho do LangChain e do LangGraph. Essa integração combina a governança do Catálogo do Unity com as funcionalidades do LangChain para criar aplicativos avançados baseados em LLM.

Neste exemplo, você cria uma ferramenta do Catálogo do Unity, testa sua funcionalidade e a adiciona a um agente.

Instalar dependências

Instale pacotes de IA do Catálogo do Unity com o Databricks opcional e instale o pacote de integração do 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 de funções do Databricks

Inicie o cliente de funções do 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 de Catálogo do Unity que contenha 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'

Encapsular a função usando o UCFunctionToolKit

Encapsule a função usando o UCFunctionToolkit para torná-la acessível às bibliotecas de criação de agentes. O kit de ferramentas garante a consistência em bibliotecas diferentes e adiciona recursos úteis, como o rastreamento automático para recuperadores.

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 de 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 criação de agente visto em Criar um agente de IA e implantá-lo nos Aplicativos do 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 Catálogo do Unity do Azure Databricks para integrar funções SQL e Python como ferramentas em fluxos de trabalho do LlamaIndex. Essa integração combina a governança do Catálogo do Unity 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 do Catálogo do Unity 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 catálogo do Unity 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 catálogo do Unity 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 aos fluxos de trabalho do OpenAI. Essa integração combina a governança do Catálogo do Unity com o OpenAI para criar aplicativos de IA de geração avançada.

  1. Instale o pacote de integração do Catálogo Unity do Databricks 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 do Catálogo do Unity 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 catálogo do Unity 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 a OpenAI retornar uma resposta, invoque a chamada de função do Catálogo Unity para gerar a resposta de volta à 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. Depois que a resposta for retornada, você poderá construir o conteúdo da resposta para chamadas subsequentes ao 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 criação da resposta da ferramenta, o pacote ucai-openai possui um utilitário, generate_tool_call_messages, que converte mensagens de ChatCompletion do OpenAI para que possam ser usadas para geração de respostas.

from unitycatalog.ai.openai.utils import generate_tool_call_messages

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

Note

Se a resposta contiver várias entradas de escolha, você poderá passar o argumento choice_index ao chamar generate_tool_call_messages para escolher qual entrada de opção utilizar. No momento, não há suporte para o processamento de entradas de várias opções.

Anthropic

Use o Azure Databricks Unity Catalog para integrar funções SQL e Python como ferramentas em chamadas do SDK Anthropic LLM. Essa integração combina a governança do Catálogo do Unity com modelos antropáticos para criar aplicativos de IA de geração avançada.

Note

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 do Catálogo do Unity 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 catálogo do Unity 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. Usar uma chamada de ferramenta na 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. Construir uma resposta da 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 de manipulador de mensagens para simplificar a análise e o tratamento de uma chamada para a função catálogo do Unity. A ferramenta faz o seguinte:

  1. Detecta os requisitos de chamada da ferramenta.
  2. Extrai informações de chamada da ferramenta da consulta.
  3. Executa a chamada para a função do Catálogo do Unity.
  4. Analisa a resposta da função catálogo do Unity.
  5. Crie o próximo formato de mensagem para continuar a conversa com Claude.

Note

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

  • Last updated on