Aller au contenu principal

Créez des outils d'agent IA à l'aide des fonctions Unity Catalog

Utilisez les fonctions Unity Catalog pour créer des outils d’agent IA qui exécutent une logique personnalisée et effectuent des tâches spécifiques qui étendent les capacités des LLM au-delà de la génération linguistique.

Quand utiliser les fonctions Unity Catalog par rapport aux serveurs MCP

Databricks recommande d'utiliser les fonctions Unity Catalog comme outils d'agent spécifiquement pour les outils de récupération de données structurées lorsque la requête est connue à l'avance et que l'agent fournit les paramètres. Voir Connecter des agents à des données structurées.

Dans la plupart des autres cas d'utilisation, Databricks recommande les serveurs MCP ou de définir la logique directement dans le code de l'agent pour une exécution plus rapide, la prise en charge de l'authentification par utilisateur et une flexibilité supplémentaire.

Exigences

Pour créer et utiliser les fonctions Unity Catalog comme outils d'agent IA, vous avez besoin des éléments suivants :

  • Databricks Runtime : utilisez Databricks Runtime 15.0 et versions ultérieures
  • **Version de Python** : installez Python 3.10 ou une version supérieure.

Pour exécuter les fonctions Unity Catalog

  • Le compute Serverless doit être activé dans votre Workspace pour exécuter les fonctions Unity Catalog en tant qu'outils d'agents IA en production. Consultez les exigences de compute Serverless.
    • L'exécution en mode local pour les fonctions Python ne nécessite pas de compute générique Serverless pour s'exécuter, cependant le mode local est uniquement destiné aux fins de développement et de test.

Pour créer des fonctions Unity Catalog :

  • Le compute générique serverless doit être activé dans votre Workspace pour créer des fonctions à l'aide du client Databricks Workspace ou des instructions de corps SQL.
    • Les fonctions Python peuvent être créées sans compute serverless.

Créer un outil de fonction Unity Catalog

Les étapes suivantes expliquent comment créer et tester une fonction Unity Catalog. Exécutez le code suivant dans un notebook Databricks.

prompt

Demandez à Genie Code (mode Agent) de le faire pour vous :

Create a Unity Catalog Python function that an AI agent can use as a tool. It should take two floating point numbers and return their sum, with type hints and a Google-style docstring. Register it using the Databricks Function Client, then test calling it.

Installer les dépendances

Installez les packages d'IA Unity Catalog avec l'extra [databricks].

Python
# Install Unity Catalog AI integration packages with the Databricks extra
%pip install unitycatalog-ai[databricks]

dbutils.library.restartPython()

Initialisez le client de fonctions Databricks

Initialisez le client de fonctions Databricks, qui est une interface spécialisée pour créer, gérer et exécuter des fonctions Unity Catalog dans Databricks.

Python
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Définir la logique de l'outil

Les outils Unity Catalog ne sont en réalité que des fonctions définies par l'utilisateur (UDF) de Unity Catalog. Lorsque vous définissez un outil Unity Catalog, vous enregistrez une fonction dans Unity Catalog. Pour en savoir plus sur les UDF de Unity Catalog, consultez les fonctions définies par l'utilisateur (UDF) SQL et Python dans Unity Catalog.

attention

L’exécution de code arbitraire dans un outil d’agent peut exposer des informations sensibles ou privées auxquelles l’agent a accès. Les clients sont responsables de l'exécution de code fiable uniquement et de la configuration des garde-fous et des autorisations appropriées pour empêcher l'accès involontaire aux données.

Vous pouvez créer des fonctions Unity Catalog à l'aide de l'une des deux APIs :

  • create_python_function accepte un callable Python.
  • create_function accepte une instruction SQL de création de corps de fonction. Consultez Créer des fonctions Python.

Utilisez l'API create_python_function pour créer la fonction.

Pour qu'un callable Python soit reconnaissable par le modèle de données des fonctions Unity Catalog, votre fonction doit répondre aux exigences suivantes :

  • Annotations de type : La signature de la fonction doit définir des annotations de type Python valides. Les arguments nommés et la valeur de retour doivent avoir leurs types définis.

  • N’utilisez pas d’arguments variables : Les arguments variables tels que *args et **kwargs ne sont pas pris en charge. Tous les arguments doivent être explicitement définis.

  • Compatibilité des types : tous les types Python ne sont pas pris en charge dans SQL. Voir types de données pris en charge par Spark.

  • Docstrings descriptives : La boîte à outils de fonctions Unity Catalog lit, analyse et extrait des informations importantes de vos docstrings.

    • Les docstrings doivent être formatées selon la syntaxe Google des docstrings.
    • Rédigez des descriptions claires pour votre fonction et ses arguments afin d'aider le LLM à comprendre comment et quand utiliser la fonction.
  • Importations de dépendances : Les bibliothèques doivent être importées dans le corps de la fonction. Les importations en dehors de la fonction ne seront pas résolues lors de l'exécution de l'outil.

Les extraits de code suivants utilisent create_python_function pour enregistrer l'objet appelable Python add_numbers:

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. Spécifiez un nom de fonction entièrement qualifié dans l'API execute_function pour exécuter la fonction :

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'

Ajouter des fonctions Unity Catalog à votre agent

Une fois que vous avez créé et testé votre fonction Unity Catalog, choisissez l'une des approches suivantes pour l'ajouter à votre agent.

Icône MCP. Utilisation de MCP (recommandé)

Utilisation de MCP (recommandé)

Databricks recommande d'utiliser des serveurs MCP pour ajouter des fonctions Unity Catalog à votre agent. L'approche MCP offre une intégration plus simple avec la découverte automatique d'outils et la prise en charge de l'authentification intégrée.

L'URL MCP gérée pour les fonctions Unity Catalog est : https://<workspace-hostname>/api/2.0/mcp/functions/{catalog}/{schema}. Vous pouvez éventuellement spécifier une fonction spécifique en ajoutant /{function_name}.

Les exemples suivants montrent comment connecter votre agent aux fonctions Unity Catalog via MCP. Remplacez <catalog> et <schema> par l'emplacement de vos fonctions.

Python
from agents import Agent, Runner
from databricks.sdk import WorkspaceClient
from databricks_openai.agents import McpServer

workspace_client = WorkspaceClient()

async with McpServer.from_uc_function(
catalog="<catalog>",
schema="<schema>",
workspace_client=workspace_client,
name="uc-functions",
) as uc_server:
agent = Agent(
name="Tool-using agent",
instructions="You are a helpful assistant. Use the available tools to answer questions.",
model="databricks-claude-sonnet-4-5",
mcp_servers=[uc_server],
)
result = await Runner.run(agent, "Look up customer info for Acme Corp")
print(result.final_output)

Accordez à l'application l'accès à la fonction Unity Catalog dans databricks.yml:

YAML
resources:
apps:
my_agent_app:
resources:
- name: 'my_uc_function'
uc_securable:
securable_full_name: '<catalog>.<schema>.<function-name>'
securable_type: 'FUNCTION'
permission: 'EXECUTE'

Icône de la fonction. Utilisation de UCFunctionToolkit

Utilisation de UCFunctionToolkit

Cet exemple utilise LangChain, mais une approche similaire peut être appliquée à d'autres bibliothèques. Découvrez l'intégration de l'outil Unity Catalog.

Installer des dépendances supplémentaires

Installez les packages d'intégration LangChain pour UCFunctionToolkit.

Python
%pip install unitycatalog-langchain[databricks]==0.2.0

# Install the Databricks LangChain integration package
%pip install databricks-langchain==0.5.0

dbutils.library.restartPython()

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 différentes bibliothèques d'IA générative et ajoute des fonctionnalités utiles telles que le traçage automatique 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.

remarque

This example uses LangChain. However you can integrate Unity Catalog tools with other frameworks such as LlamaIndex, OpenAI, Anthropic, and more. See Unity Catalog tool integration.

Cet exemple crée un agent simple à l'aide de l'API LangChain AgentExecutor par souci de simplicité. Pour les charges de travail de production, utilisez le workflow de création d'agents présenté dans Créez un agent d'IA et déployez-le 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 (optional: replace with your LLM of choice)
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?"})

Améliorer l'appel d'outil avec une documentation claire

Une bonne documentation aide vos agents à savoir quand et comment utiliser chaque outil. Suivez ces bonnes pratiques pour documenter vos outils :

  • Pour les fonctions Unity Catalog, utilisez la clause COMMENT pour décrire la fonctionnalité de l'outil et les paramètres.
  • Définissez clairement les entrées et les sorties attendues.
  • Écrivez des descriptions significatives pour rendre les outils plus faciles à utiliser pour les agents et les humains.

Exemple : documentation d'outil efficace

L'exemple suivant montre des chaînes COMMENT claires pour un outil qui interroge une table structurée.

SQL
CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
customer_name STRING COMMENT 'Name of the customer whose info to look up.'
)
RETURNS STRING
COMMENT 'Returns metadata about a specific customer including their email and ID.'
RETURN SELECT CONCAT(
'Customer ID: ', customer_id, ', ',
'Customer Email: ', customer_email
)
FROM main.default.customer_data
WHERE customer_name = customer_name
LIMIT 1;

Exemple : Documentation d'outil inefficace

L'exemple suivant manque de détails importants, ce qui rend l'utilisation de l'outil plus difficile pour les agents :

SQL
CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
customer_name STRING COMMENT 'Name of the customer.'
)
RETURNS STRING
COMMENT 'Returns info about a customer.'
RETURN SELECT CONCAT(
'Customer ID: ', customer_id, ', ',
'Customer Email: ', customer_email
)
FROM main.default.customer_data
WHERE customer_name = customer_name
LIMIT 1;

Exécutez les fonctions en mode Serverless ou local.

Lorsqu'un service d'IA générative détermine qu'un appel d'outil est nécessaire, les packages d'intégration (UCFunctionToolkit instances) exécutent l'DatabricksFunctionClient.execute_function API.

L'appel execute_function peut exécuter des fonctions dans deux modes d'exécution : serverless ou local. Ce mode détermine quelle ressource exécute la fonction.

Mode Serverless pour la production

Le mode Serverless est l'option par default et recommandée pour les cas d'utilisation en production lors de l'exécution de fonctions Unity Catalog en tant qu'outils d'agent d'IA. Ce mode utilise un compute générique Serverless (Spark Connect Serverless) pour exécuter des fonctions à distance, et Lakeguard garantit que le processus de votre agent reste sécurisé et exempt des risques liés à l'exécution de code arbitraire localement.

remarque

Les fonctions Unity Catalog exécutées en tant qu’outils d’agent IA nécessitent un compute générique Serverless (Spark Connect Serverless), et non des SQL Warehouse Serverless. Les tentatives d’exécution d’outils sans compute générique serverless produiront des erreurs comme PERMISSION_DENIED: Cannot access Spark Connect.

Python
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="serverless")

Lorsque votre agent demande l'exécution d'un outil en mode serverless , les éléments suivants se produisent :

  1. Le DatabricksFunctionClient envoie une requête à Unity Catalog pour récupérer la définition de la fonction si la définition n'a pas été mise en cache localement.
  2. Le DatabricksFunctionClient extrait la définition de fonction et valide les noms et les types des paramètres.
  3. Le DatabricksFunctionClient soumet l'exécution en tant qu'UDF au compute générique serverless.

Mode local pour le développement

Le mode local exécute les fonctions Python dans un sous-processus local au lieu d’effectuer des requêtes vers le compute générique serverless. Cela vous permet de dépanner plus efficacement les appels d’outils en fournissant des traces de pile locales. Il est conçu pour le développement et le debugging des fonctions Python Unity Catalog.

Lorsque votre agent demande l'exécution d'un outil en mode local , le DatabricksFunctionClient effectue les opérations suivantes :

  1. Envoie une requête à Unity Catalog pour récupérer la définition de fonction si la définition n'a pas été mise en cache localement.
  2. Extrait la définition de l'appelable Python, met l'appelable en cache localement et valide les noms et les types des paramètres.
  3. Appelle l'objet appelable avec les paramètres spécifiés dans un sous-processus restreint avec protection contre le délai d'expiration.
Python
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="local")

L'exécution en mode "local" offre les fonctionnalités suivantes :

  • Limite de temps CPU : restreint le temps d'exécution CPU total pour l'exécution appelable afin d'éviter les charges de calcul excessives.

    La limite de temps CPU est basée sur l'utilisation réelle du CPU, et non sur le temps réel. En raison de la planification du système et des processus concurrents, le temps CPU peut dépasser le temps réel dans des scénarios réels.

  • Limite de mémoire : limite la mémoire virtuelle allouée au processus.

  • Protection contre le délai d'expiration : applique un délai d'exécution maximal pour les fonctions en cours d'exécution.

Personnalisez ces limites à l'aide de variables d'environnement (pour en savoir plus).

Limitations du mode local

  • **Fonctions Python uniquement** : Les fonctions basées sur SQL ne sont pas prises en charge en mode local.
  • Considérations de sécurité pour le code non approuvé : Bien que le mode local exécute les fonctions dans un sous-processus pour l'isolation des processus, il existe un risque de sécurité potentiel lors de l'exécution de code arbitraire généré par les systèmes d'IA. Ceci est principalement préoccupant lorsque des fonctions exécutent du code Python généré dynamiquement qui n'a pas été examiné.
  • Différences de version des bibliothèques : les versions des bibliothèques peuvent différer entre les environnements d'exécution Serverless et locaux, ce qui peut entraîner un comportement différent des fonctions.

Variables d'environnement

Configurez la manière dont les fonctions s’exécutent dans le DatabricksFunctionClient à l’aide des variables d’environnement suivantes :

Variable d'environnement

Valeur par défaut

Description

EXECUTOR_MAX_CPU_TIME_LIMIT

10 secondes

Temps d'exécution CPU maximal autorisé (mode local uniquement).

EXECUTOR_MAX_MEMORY_LIMIT

100 Mo

Allocation maximale de mémoire virtuelle autorisée pour le processus (mode local uniquement).

EXECUTOR_TIMEOUT

20 secondes

Durée prévue totale maximale (mode local uniquement).

UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS

5

Le nombre maximal de tentatives de nouvelle actualisation du client de session en cas d’expiration du jeton.

UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT

100

Le nombre maximal de lignes à renvoyer lors de l'exécution de fonctions utilisant le compute Serverless et databricks-connect.

Variable d'environnement

Valeur par défaut

Description

EXECUTOR_MAX_CPU_TIME_LIMIT

10 secondes

Temps d'exécution CPU maximal autorisé (mode local uniquement).

EXECUTOR_MAX_MEMORY_LIMIT

100 Mo

Allocation maximale de mémoire virtuelle autorisée pour le processus (mode local uniquement).

EXECUTOR_TIMEOUT

20 secondes

Durée prévue totale maximale (mode local uniquement).

UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS

5

Le nombre maximal de tentatives de nouvelle actualisation du client de session en cas d’expiration du jeton.

UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT

100

Le nombre maximal de lignes à renvoyer lors de l'exécution de fonctions utilisant le compute Serverless et databricks-connect.

Appeler les APIs externes avec http_request (hérité)

remarque

Pour connecter des agents à des services externes, Databricks recommande les services MCP ou le proxy de connexions Unity Catalog. Les outils de fonctions UC qui encapsulent http_request restent pris en charge, mais ne sont plus l'approche recommandée.

Vous pouvez créer une fonction Unity Catalog qui encapsule http_request() pour appeler des services externes. Cette approche est utile pour les définitions d’outils basées sur SQL.

L'exemple suivant crée un outil de fonction Unity Catalog qui publie un message sur Slack :

SQL
CREATE OR REPLACE FUNCTION main.default.slack_post_message(
text STRING COMMENT 'message content'
)
RETURNS STRING
COMMENT 'Sends a Slack message by passing in the message and returns the response received from the external service.'
RETURN (http_request(
conn => 'test_sql_slack',
method => 'POST',
path => '/api/chat.postMessage',
json => to_json(named_struct(
'channel', "C032G2DAH3",
'text', text
))
)).text

Voir CREATE FUNCTION (SQL, Python, Scala et Java).

remarque

L'accès SQL avec http_request est bloqué pour les types de connexion Utilisateur-à-machine par utilisateur et Enregistrement dynamique des clients. Utilisez plutôt le SDK Python de Databricks.

Exemples de notebooks

Les notebooks suivants démontrent la création d’outils d’agent IA qui se connectent aux services externes à l’aide des fonctions Unity Catalog.

Outil d'agent de messagerie Slack

Outil d'agent API Microsoft Graphe

Outil d'agent Azure AI Search

Étapes suivantes