Pular para o conteúdo principal

Integrar o LangChain com as ferramentas do Databricks Unity Catalog

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

Requisitos

  • Instale o Python 3.10 e o acima.

Integrar o LangChain com o Databricks Unity Catalog

Neste exemplo, o senhor cria uma ferramenta do Unity Catalog, testa sua funcionalidade e a adiciona a um agente. Execute o seguinte código em um Notebook Databricks.

Instalar dependências

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

Este exemplo usa o site LangChain, mas uma abordagem semelhante pode ser aplicada a outras bibliotecas. Consulte Integrar ferramentas Unity Catalog com estruturas geradoras de AI de terceiros.

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

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

Python

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 ela funciona conforme o esperado:

Python
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 endereço UCFunctionToolkit para torná-la acessível à biblioteca de autoria do agente. O kit de ferramentas garante a consistência entre diferentes bibliotecas e acrescenta recursos úteis, como o rastreamento automático para recuperadores.

Python
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 criação de agentes visto nos exemplos do siteChatAgent.

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