Aller au contenu principal

Créer 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 de langage.

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 query est connue à l'avance et que l'agent fournit les parameters. Voir Connecter les agents aux 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é accrue.

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 Python : installez Python 3.10 ou version ultérieure

Pour exécuter des fonctions Unity Catalog :

  • **Serverless compute** doit être activé dans votre Workspace pour exécuter les fonctions Unity Catalog en tant qu'outils d'agent IA en production. Consultez Exigences en matière de Serverless compute.
    • 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 montrent 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()

Initialiser le client de fonction Databricks

Initialisez le client de fonction 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 sont en réalité juste des fonctions définies par l'utilisateur (UDF) 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 uniquement de code fiable et de la configuration de garde-fous et d'autorisations appropriées afin d'éviter tout 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 appelable Python.
  • create_function accepte une instruction de création de fonction de corps SQL. Voir Créer des fonctions Python.

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

Pour rendre un callable Python 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. Aussi bien les arguments nommés que 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 en SQL. Voir les types de données pris en charge par Spark.

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

    • Les docstrings doivent être formatées conformément à la syntaxe de docstring Google.
    • 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 le create_python_function pour enregistrer l'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 vous 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 facultativement 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 fonction. Utilisation de UCFunctionToolkit

Utilisation de UCFunctionToolkit

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

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

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

Utilisez 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 en utilisant l'API LangChain AgentExecutor pour des raisons 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 (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éliorez l'appel d'outils grâce à 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é et les parameters de l'outil.
  • Définissez clairement les entrées et les sorties attendues.
  • Rédigez des descriptions significatives pour faciliter l'utilisation des outils par les agents et les humains.

Exemple : documentation d'outil efficace

L'exemple suivant montre COMMENT chaînes 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 plus difficile pour les agents d'utiliser l'outil efficacement :

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écuter des 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'API DatabricksFunctionClient.execute_function.

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 IA. Ce mode utilise un compute serverless générique (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 d'IA nécessitent un compute générique Serverless (Spark Connect Serverless), et non des SQL Warehouses 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 , voici ce qui se passe :

  1. DatabricksFunctionClient envoie une requête à Unity Catalog pour récupérer la définition de fonction si celle-ci n'a pas été mise en cache localement.
  2. Le DatabricksFunctionClient extrait la définition de fonction et valide les noms et 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'envoyer des requêtes au compute générique Serverless. Cela vous permet de dépanner les appels d'outil plus efficacement en fournissant des traces de la pile locales. Il est conçu pour le développement et le debugging des fonctions Python Unity Catalog.

Lorsqu'un 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 en cache l’appelable 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épassement de délai.
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 :** Limite la durée d'exécution totale du CPU pour l'exécution des appelables 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 la durée prévue. En raison de la planification système et des processus concurrents, la durée CPU peut dépasser la durée prévue dans des scénarios réels.

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

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

Personnalisez ces limites à l'aide de variables d'environnement (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 fiable : 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 des systèmes d'IA. Ceci est principalement une préoccupation lorsque les fonctions exécutent du code Python généré dynamiquement qui n'a pas été examiné.
  • **Différences de version de bibliothèque** : Les versions de bibliothèque peuvent différer entre les environnements d'exécution serverless et locaux, ce qui pourrait entraîner un comportement de fonction différent.

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

Durée maximale autorisée de l'exécution du processeur (mode local uniquement).

EXECUTOR_MAX_MEMORY_LIMIT

100 Mo

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

EXECUTOR_TIMEOUT

20 secondes

Durée maximale totale de l'horloge murale (mode local uniquement).

UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS

5

Le nombre maximal de tentatives pour réessayer d'actualiser le 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 à l'aide de compute serverless et de databricks-connect.

Variable d'environnement

Valeur par défaut

Description

EXECUTOR_MAX_CPU_TIME_LIMIT

10 secondes

Durée maximale autorisée de l'exécution du processeur (mode local uniquement).

EXECUTOR_MAX_MEMORY_LIMIT

100 Mo

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

EXECUTOR_TIMEOUT

20 secondes

Durée maximale totale de l'horloge murale (mode local uniquement).

UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS

5

Le nombre maximal de tentatives pour réessayer d'actualiser le 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 à l'aide de compute serverless et de databricks-connect.

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

remarque

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

Vous pouvez créer une fonction Unity Catalog qui enveloppe 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 d'enregistrement dynamique de client. Utilisez plutôt le SDK Databricks pour Python.

Exemples de Notebooks

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

Outil d'agent de messagerie Slack

Outil d'agent de l'API Microsoft Graphe

Outil d'agent Azure AI Search

Étapes suivantes