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.
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].
# 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.
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.
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_functionaccepte un appelable Python.create_functionaccepte 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:
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 :
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.
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.
- OpenAI Agents SDK (Apps)
- LangGraph (Apps)
- Model Serving
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:
resources:
apps:
my_agent_app:
resources:
- name: 'my_uc_function'
uc_securable:
securable_full_name: '<catalog>.<schema>.<function-name>'
securable_type: 'FUNCTION'
permission: 'EXECUTE'
from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient
from langgraph.prebuilt import create_react_agent
workspace_client = WorkspaceClient()
host = workspace_client.config.host
mcp_client = DatabricksMultiServerMCPClient([
DatabricksMCPServer(
name="uc-functions",
url=f"{host}/api/2.0/mcp/functions/<catalog>/<schema>",
workspace_client=workspace_client,
),
])
async with mcp_client:
tools = await mcp_client.get_tools()
agent = create_react_agent(
ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
tools=tools,
)
result = await agent.ainvoke(
{"messages": [{"role": "user", "content": "Look up customer info for Acme Corp"}]}
)
print(result["messages"][-1].content)
Accordez à l'application l'accès à la fonction Unity Catalog dans databricks.yml:
resources:
apps:
my_agent_app:
resources:
- name: 'my_uc_function'
uc_securable:
securable_full_name: '<catalog>.<schema>.<function-name>'
securable_type: 'FUNCTION'
permission: 'EXECUTE'
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient
import mlflow
workspace_client = WorkspaceClient()
host = workspace_client.config.host
# Connect to the UC functions MCP server
mcp_client = DatabricksMCPClient(
server_url=f"{host}/api/2.0/mcp/functions/<catalog>/<schema>",
workspace_client=workspace_client,
)
# List available tools
tools = mcp_client.list_tools()
# Log the agent with the required resources for deployment
mlflow.pyfunc.log_model(
"agent",
python_model=my_agent,
resources=mcp_client.get_databricks_resources(),
)
Pour déployer l'agent, consultez Déployer un agent pour les applications d'IA générative (Model Serving). Pour plus de détails sur la journalisation des agents avec les ressources MCP, consultez Serveurs MCP gérés par Databricks.
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.
%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.
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.
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.
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
COMMENTpour 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.
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 :
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.
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.
# 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 :
DatabricksFunctionClientenvoie 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.- Le
DatabricksFunctionClientextrait la définition de fonction et valide les noms et types des paramètres. - Le
DatabricksFunctionClientsoumet 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 :
- 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.
- Extrait la définition de l’appelable Python, met en cache l’appelable localement et valide les noms et les types des paramètres.
- 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.
# 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 |
|---|---|---|
|
| Durée maximale autorisée de l'exécution du processeur (mode local uniquement). |
|
| Allocation maximale autorisée de mémoire virtuelle pour le processus (mode local uniquement). |
|
| Durée maximale totale de l'horloge murale (mode local uniquement). |
|
| Le nombre maximal de tentatives pour réessayer d'actualiser le client de session en cas d'expiration du jeton. |
|
| Le nombre maximal de lignes à renvoyer lors de l'exécution de fonctions à l'aide de compute serverless et de |
Appeler des APIs externes avec http_request (hérité)
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 :
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).
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
-
Ajoutez des outils Unity Catalog aux agents par programme. Consultez Créez un agent d'IA et déployez-le sur Databricks Apps.
-
Ajoutez les outils Unity Catalog aux agents à l'aide de l'interface utilisateur de l'AI Playground. Consultez Get start: interroger des LLM et prototyper des agents d'IA sans code.
-
Gérer les fonctions Unity Catalog à l'aide du client de fonction. Consultez la documentation Unity Catalog - client de fonction
-
Connectez les agents à des outils tiers avec les services MCP pour un aperçu de toutes les approches pour connecter les agents à des services externes.