Aller au contenu principal

Intégrer les outils Unity Catalog avec des frameworks d'IA générative tiers

Les outils d’agent IA de Unity Catalog peuvent être utilisés dans des bibliothèques d’IA générative populaires comme LangChain, LlamaIndex, OpenAI et Anthropic. Ces intégrations combinent la gouvernance des outils de Unity Catalog avec les capacités des frameworks de création d’agents tiers. Par exemple :

  • Dans LangChain, les fonctions Unity Catalog peuvent faire partie du workflow d'un agent pour effectuer des tâches comme l'interrogation ou la transformation de données.
  • Dans les intégrations OpenAI ou Anthropic, les fonctions sont appelées directement par le modèle d'IA pendant l'exécution.

Sélectionnez votre framework dans les onglets suivants pour créer un outil Unity Catalog et l'utiliser avec ce framework. Exécutez le code dans un notebook Databricks ou un script Python.

Exigences

  • Installez Python 3.10 ou version ultérieure.

Utiliser Databricks Unity Catalog pour intégrer les fonctions SQL et Python comme outils dans les workflows LangChain et LangGraph. Cette intégration combine la gouvernance d'Unity Catalog avec les capacités de LangChain pour créer de puissantes applications basées sur les LLM.

Dans cet exemple, vous créez un outil Unity Catalog, testez sa fonctionnalité et l'ajoutez à un agent.

Installer les dépendances

Installez les packages d'IA du Unity Catalog avec les options Databricks et installez le package d'intégration LangChain.

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

Initialisez le client de fonctions Databricks

Initialisez le client de fonction Databricks.

Python
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 Unity Catalog contenant la logique de l'outil.

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
)

Testez la fonction

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

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'

Enveloppez la fonction à l’aide du UCFunctionToolKit

Encapsulez la fonction à l'aide du UCFunctionToolkit pour la rendre accessible aux bibliothèques de création d'agents. La boîte à outils assure la cohérence entre les différentes bibliothèques et ajoute des fonctionnalités utiles comme l'auto-traçage pour les récupérateurs.

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

Utiliser l'outil dans un agent

Ajoutez l'outil à un agent LangChain en utilisant la propriété tools de UCFunctionToolkit.

Cet exemple crée un agent simple à l'aide de l'API AgentExecutor de LangChain par souci de simplicité. Pour les charges de travail en production, utilisez le workflow de création d'agent présenté dans Créer un agent d'IA et le déployer sur Databricks Apps.

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