Intégrer LangChain aux outils de catalogue Databricks Unity

Utilisez databricks Unity Catalog pour intégrer des fonctions SQL et Python en tant qu’outils dans les flux de travail LangChain et LangGraph. Cette intégration combine la gouvernance du catalogue Unity avec les fonctionnalités LangChain pour créer des applications LLM puissantes.

Spécifications

• Installez Python 3.10 et versions ultérieures.

Intégrer LangChain au catalogue Databricks Unity

Dans cet exemple, vous créez un outil de catalogue Unity, testez ses fonctionnalités et ajoutez-le à un agent. Exécutez le code suivant dans un notebook Databricks.

Installer des dépendances

Installez les packages UNITY Catalog AI avec databricks facultatif et installez le package d’intégration LangChain.

Cet exemple utilise LangChain, mais une approche similaire peut être appliquée à d’autres bibliothèques. Consultez Intégrer les outils de catalogue Unity à des infrastructures d’IA générative tierces.

# 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()

Initialiser le client de fonction Databricks

Initialisez le client de fonction Databricks.

from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

Définir la logique de l’outil

Créez une fonction de catalogue Unity contenant la logique de l’outil.

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
)

Tester la fonction

Testez votre fonction pour vérifier qu’elle fonctionne comme prévu :

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'

Encapsuler la fonction à l’aide de l’UCFunctionToolKit

Encapsulez la fonction en utilisant le UCFunctionToolkit pour la rendre accessible aux bibliothèques de création d’agents. Le kit de ressources garantit la cohérence entre différentes bibliothèques et ajoute des fonctionnalités utiles telles que le suivi automatique pour les récupérateurs.

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

Utiliser l’outil dans un agent

Ajoutez l’outil à un agent LangChain à l’aide de la propriété tools à partir de UCFunctionToolkit.

Cet exemple crée un agent simple à l’aide de l’API AgentExecutor LangChain pour plus de simplicité. Pour les charges de travail de production, utilisez le flux de travail de création d’agent vu dans les exemples ResponsesAgent.

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