Integrar ferramentas do Unity Catalog com estruturas de AI de terceiros
As ferramentas de agente do Unity Catalog podem ser usadas em bibliotecas de AI populares como LangChain, LlamaIndex, OpenAI e Anthropic. Essas integrações combinam a governança de ferramentas do Unity Catalog com os recursos de estruturas de autoria de agentes de terceiros. Por exemplo:
- Na 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 seu framework nas abas a seguir para criar uma ferramenta do Unity Catalog e usá-la com esse framework. Execute o código em um Notebook Databricks ou script Python.
Requisitos
- Instale o Python 3.10 ou acima.
- LangChain
- LlamaIndex
- OpenAI
- Anthropic
Utilize o Databricks Unity Catalog para integrar funções SQL e Python como ferramentas em fluxos de trabalho LangChain e LangGraph. Esta integração combina a governança do Unity Catalog com recursos de LangChain para construir aplicativos poderosos baseados em LLM.
Neste exemplo, o senhor cria uma ferramenta do Unity Catalog, testa sua funcionalidade e a adiciona a um agente.
Instalar dependências
Instale os pacotes de AI do Unity Catalog 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()
Inicialize o Cliente de Funções do Databricks
Inicialize o Cliente de Função do Databricks.
from unitycatalog.ai.core.base import get_uc_function_client
client = get_uc_function_client()
Defina a lógica da ferramenta
Crie uma função do Unity Catalog 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'
Envolva a função usando o UCFunctionToolKit
Encapsule a função usando UCFunctionToolkit para torná-la acessível às bibliotecas de autoria de agente. O kit de ferramentas garante a consistência entre diferentes bibliotecas e adiciona recursos úteis como 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
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 autoria de agente visto em Crie um agente e o implante no Databricks Apps.
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?"})
Use o Databricks Unity Catalog para integrar funções SQL e Python como ferramentas em fluxos de trabalho do LlamaIndex. Essa integração combina a governança do Unity Catalog com os recursos do LlamaIndex para indexar e query grandes datasets para LLMs.
-
Instale o pacote de integração Databricks Unity Catalog para LlamaIndex.
Python%pip install unitycatalog-llamaindex[databricks]
dbutils.library.restartPython() -
Crie uma instância do cliente de funções do Unity Catalog.
Pythonfrom unitycatalog.ai.core.base import get_uc_function_client
client = get_uc_function_client() -
Crie uma função do Unity Catalog escrita em Python.
PythonCATALOG = "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
) -
Crie uma instância da função do Unity Catalog como um kit de ferramentas, e execute-a para verificar se a ferramenta se comporta corretamente.
Pythonfrom 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"} -
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.
Pythonfrom 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)`")
Use o Databricks Unity Catalog para integrar funções SQL e Python como ferramentas em fluxos de trabalho da OpenAI. Esta integração combina a governança do Unity Catalog com a OpenAI para criar aplicativos de AI poderosos.
-
Instale o pacote de integração do Databricks Unity Catalog para OpenAI.
Python%pip install unitycatalog-openai[databricks]
%pip install mlflow -U
dbutils.library.restartPython() -
Crie uma instância do cliente de funções do Unity Catalog.
Pythonfrom unitycatalog.ai.core.base import get_uc_function_client
client = get_uc_function_client() -
Crie uma função do Unity Catalog escrita em Python.
PythonCATALOG = "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
) -
Crie uma instância da função do Unity Catalog como um kit de ferramentas e verifique se a ferramenta se comporta corretamente executando a função.
Pythonfrom 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] -
Envie a solicitação ao modelo OpenAI junto com as ferramentas.
Pythonimport 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) -
Depois que o OpenAI retorna uma resposta, invoque a chamada de função do Unity Catalog para gerar a resposta de volta ao OpenAI.
Pythonimport 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) -
Uma vez que a resposta tenha sido retornada, você pode construir o payload da resposta para chamadas subsequentes à OpenAI.
Python# 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"]
)
utilidades
Para simplificar o processo de elaboração da resposta da ferramenta, o ucai-openai pacote possui uma utilidade, generate_tool_call_messages, que converte as mensagens de resposta do OpenAI ChatCompletion para que possam ser usadas para a 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)
Se a resposta contiver entradas de múltipla escolha, você pode 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.
Utilize o Databricks Unity Catalog para integrar funções SQL e Python como ferramentas em chamadas de LLM do SDK da Anthropic. Esta integração combina a governança do Unity Catalog com modelos da Anthropic para criar aplicativos de AI poderosos.
A integração da Anthropic requer Databricks Runtime 15.0 e acima.
-
Instale o pacote de integração do Databricks Unity Catalog para a Anthropic.
Python%pip install unitycatalog-anthropic[databricks]
dbutils.library.restartPython() -
Crie uma instância do cliente de funções do Unity Catalog.
Pythonfrom unitycatalog.ai.core.base import get_uc_function_client
client = get_uc_function_client() -
Crie uma função do Unity Catalog escrita em Python.
PythonCATALOG = "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
) -
Crie uma instância da função do Unity Catalog como um kit de ferramentas.
Pythonfrom unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit
# Create an instance of the toolkit
toolkit = UCFunctionToolkit(function_names=[func_name], client=client) -
Use uma chamada de ferramenta no Anthropic.
Pythonimport 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) -
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.
Pythonfrom 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 pacote unitycatalog.ai-anthropic inclui uma utilidade de manipulador de mensagens para simplificar a análise e o tratamento de uma chamada para a função do Unity Catalog. A utilidade faz o seguinte:
- Detecta os requisitos de chamada de ferramenta.
- Extrai informações de chamada de ferramenta da query.
- Realiza a chamada para a função do Unity Catalog.
- Analisa a resposta da função do Unity Catalog.
- Crie o próximo formato de mensagem para continuar a conversa com Claude.
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 de entrada do usuário) e todas as respostas subsequentes geradas por LLM e resultados de chamada de ferramenta em várias etapas.