Integrar ferramentas do Unity Catalog com frameworks de AI generativa de terceiros
As ferramentas de agente de AI do Unity Catalog podem ser usadas em bibliotecas populares de AI generativa 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:
- 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.
- Nas integrações OpenAI ou Anthropic, as funções são chamadas diretamente pelo modelo de AI durante a execução.
Selecione sua estrutura nas seguintes abas para criar uma ferramenta do Unity Catalog e usá-la com essa estrutura. Execute o código em um Notebook Databricks ou script Python.
Requisitos
- Instale o Python 3.10 ou acima.
- LangChain
- LlamaIndex
- OpenAI
- Anthropic
Use o Databricks Unity Catalog 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 Unity Catalog com os recursos da LangChain para construir aplicativos poderosos baseados em LLM.
Neste exemplo, você cria uma ferramenta do Unity Catalog, testa sua funcionalidade e a adiciona a um agente.
Instalar dependências
Instale pacotes de AI do Unity Catalog com o opcional do Databricks e instale o pacote de integração da 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ção do Databricks
Inicializar 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 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
)
Teste a função
Teste sua função para verificar se funciona como 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
Empacote a função usando UCFunctionToolkit para torná-la acessível às bibliotecas de criação de agentes. O kit de ferramentas garante a consistência em diferentes bibliotecas 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
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 do LangChain por simplicidade. Para cargas de trabalho de produção, use o fluxo de trabalho de criação de agentes visto em Crie um agente de AI e seja implantado nos 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 Unity Catalog do Databricks para integrar funções SQL e Python como ferramentas em fluxos de trabalho do LlamaIndex. Esta integração combina a governança do Unity Catalog com os recursos do LlamaIndex para indexar e consultar grandes datasets para LLMs.
-
Instale o pacote de integração do 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 realize a execução dela 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 do Unity Catalog como parte de uma coleção de ferramentas LlamaIndex. Em seguida, verifique se o agente se comporta adequadamente 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 do OpenAI. Esta integração combina a governança do Unity Catalog com o OpenAI para criar poderosos aplicativos de AI generativa.
-
Instale o pacote de integração do Unity Catalog da Databricks 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 para o modelo OpenAI juntamente 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 retornar 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) -
Assim que a resposta for retornada, você pode construir o payload de 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 criar a resposta da ferramenta, o pacote ucai-openai tem uma utilidade, generate_tool_call_messages, que converte 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ê 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 processar entradas de múltipla escolha.
Use o Databricks Unity Catalog para integrar funções SQL e Python como ferramentas em chamadas de LLM do Anthropic SDK. Esta integração combina a governança do Unity Catalog com modelos Anthropic para criar aplicativos de IA generativa poderosos.
A integração da Anthropic requer o Databricks Runtime 15.0 e acima.
-
Instale o pacote de integração do Databricks Unity Catalog para 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) -
Construir uma resposta da ferramenta. A resposta do modelo Claude contém um bloco de metadados de solicitação de ferramenta caso uma ferramenta precise 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 utilidades de manipulador de mensagens para simplificar a análise e o tratamento de uma chamada para a função do Unity Catalog. O utilitário executa o seguinte:
- Detecta requisitos de chamada de ferramenta.
- Extrai informação 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.
- Elabore o próximo formato de mensagem para continuar a conversa com Claude.
A história completa 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 os resultados de chamadas de ferramenta de vários turnos.