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 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 telles que 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 tabs suivantes 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.

Utilisez 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 ses fonctionnalités et l'ajoutez à un agent.

Installer les dépendances

Installez les packages IA d'Unity Catalog avec l'option 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()

Initialiser le client de fonction 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'

Encapsulez la fonction à l'aide du UCFunctionToolKit.

Encapsulez la fonction à l'aide de 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

Utilisez 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 utilisant l'API AgentExecutor de LangChain pour plus de simplicité. Pour les charges de travail de production, utilisez le workflow de création d'agent décrit dans Créer un agent 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?"})