Partager via


Créer des outils d’agent IA à l’aide des fonctions de catalogue Unity

Utilisez les fonctions catalogue Unity 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 fonctionnalités de LLMs au-delà de la génération de langage.

Quand utiliser des fonctions de catalogue Unity et des serveurs MCP

Databricks recommande d’utiliser des fonctions de catalogue Unity en tant qu’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. Consultez Connecter des agents à des données structurées.

Dans la plupart des autres cas d’usage, Databricks recommande les serveurs MCP ou la définition de la logique directement dans le code de l’agent pour accélérer l’exécution, la prise en charge de l’authentification par utilisateur et une flexibilité supplémentaire.

Spécifications

Pour créer et utiliser des fonctions de catalogue Unity en tant qu’outils d’agent IA, vous avez besoin des éléments suivants :

  • Databricks Runtime : Utiliser Databricks Runtime 15.0 et versions ultérieures
  • Python version : installer Python 3.10 ou version ultérieure

Pour exécuter des fonctions de catalogue Unity :

  • Le calcul serverless doit être activé dans votre espace de travail pour utiliser les fonctions de Unity Catalog en tant qu'outils des agents IA en production. Consultez les exigences en matière de calcul serverless.
    • Exécution en mode local pour les fonctions Python ne nécessite pas l’exécution du calcul générique serverless, mais le mode local est destiné uniquement à des fins de développement et de test.

Pour créer des fonctions de catalogue Unity :

  • Le calcul générique serverless doit être activé dans votre espace de travail pour créer des fonctions à l’aide des instructions Databricks Workspace Client ou SQL Body.
    • Les fonctions Python peuvent être créées sans calcul serverless.

Créer un outil pour les fonctions du catalogue Unity

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

Installer des dépendances

Installez les packages IA du catalogue Unity avec l'ajout [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 la création, la gestion et l’exécution de fonctions de catalogue Unity dans Databricks.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()

Définir la logique de l’outil

Les outils Unity Catalog ne sont vraiment que des fonctions définies par l'utilisateur (UDF) du catalogue Unity en arrière-plan. Lorsque vous définissez un outil de catalogue Unity, vous inscrivez une fonction dans Unity Catalog. Pour en savoir plus sur les fonctions définies par l’utilisateur du catalogue Unity, consultez les fonctions définies par l’utilisateur dans le catalogue Unity.

Avertissement

L’exécution d’un 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 du code approuvé 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 de catalogue Unity à l’aide de l’une des deux API suivantes :

  • create_python_function accepte un objet Python appelable.
  • create_function accepte une instruction de fonction de création de corps SQL. Consultez Create Python functions.

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

Pour rendre une fonction Python appelable reconnaissable par le modèle de données des fonctions Unity Catalog, votre fonction doit satisfaire aux exigences suivantes :

  • Indications de type : la signature de fonction doit définir des indications de type Python valides. Les arguments nommés et la valeur de retour doivent avoir leurs types définis.
  • N’utilisez pas d’arguments de variable : les arguments de variable tels que *args et **kwargs ne sont pas pris en charge. Tous les arguments doivent être définis explicitement.
  • compatibilité Type : tous les types Python ne sont pas pris en charge dans SQL. Consultez les types de données pris en charge par Spark.
  • Documents descriptifs : le kit de ressources de fonctions de catalogue Unity lit, analyse et extrait des informations importantes de votre docstring.
    • Les docstrings doivent être mis en forme en fonction de la syntaxe google docstring.
    • Écrivez des descriptions claires pour votre fonction et ses arguments pour aider le LLM à comprendre comment et quand utiliser la fonction.
  • Importations de dépendances : il faut importer les bibliothèques dans le corps de la fonction. Les importations en dehors de la fonction ne sont pas résolues lors de l’exécution de l’outil.

L'extrait de code suivant utilise le create_python_function pour enregistrer l'objet 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
)

Tester la fonction

Testez votre fonction pour vérifier qu’elle fonctionne comme prévu. Spécifiez un nom de fonction complet 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 de catalogue Unity à votre agent

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

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

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

L’URL MCP gérée pour les fonctions catalogue Unity est la suivante : 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 de catalogue Unity via MCP. Remplacez <catalog> et <schema> par l’emplacement de vos fonctions.

SDK des Agents OpenAI (Applications)

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 Catalogue Unity 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'

LangGraph (Applications)

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 Catalogue Unity 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'

Utilisation de modèles

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 IA génératives (Model Serving). Pour plus d’informations sur la journalisation des agents avec des ressources MCP, consultez Utiliser des serveurs MCP gérés par Databricks.

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 Intégrer les outils de catalogue Unity à des infrastructures d’IA générative tierces.

Installer des dépendances supplémentaires

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

%pip install unitycatalog-langchain[databricks]

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

dbutils.library.restartPython()

Encapsuler la fonction à l’aide de l’UCFunctionToolKit

Encapsulez la fonction en utilisant le UCFunctionToolkit pour la rendre accessible aux bibliothèques de création d’agents. Le kit de ressources garantit la cohérence entre différentes bibliothèques d'IA générative et ajoute des fonctionnalités utiles telles que le suivi 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

Utiliser l’outil dans un agent

Ajoutez l’outil à un agent LangChain à l’aide de la propriété tools à partir de UCFunctionToolkit.

Remarque

Cet exemple utilise LangChain. Toutefois, vous pouvez intégrer des outils de catalogue Unity à d’autres frameworks tels que LlamaIndex, OpenAI, Anthropic, etc. Consultez Intégrer les outils de catalogue Unity à des infrastructures d’IA générative tierces.

Cet exemple crée un agent simple à l’aide de l’API LangChain AgentExecutor pour plus de simplicité. Pour les charges de travail de production, utilisez le flux de travail de création d’agent vu dans Créer un agent IA et déployez-le 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éliorer l’appel d’outils 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 du catalogue Unity, utilisez la clause COMMENT pour décrire les fonctionnalités et les paramètres de l’outil.
  • Définissez clairement les entrées et sorties attendues.
  • Écrivez des descriptions significatives pour faciliter l’utilisation des outils pour les agents et les humains.

Exemple : Documentation efficace de l’outil

L’exemple suivant montre des chaînes claires COMMENT 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 sur les outils inefficaces

L’exemple suivant ne contient pas de détails importants, ce qui complique l’utilisation efficace de l’outil par les agents :

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 à l’aide du mode serverless ou local

Lorsqu’un service IA de génération 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 la ressource qui exécute la fonction.

Mode serverless pour la production

Le mode serverless est l’option par défaut et recommandée pour les cas d’utilisation de production lors de l’exécution de fonctions de catalogue Unity en tant qu’outils d’agent IA. Ce mode utilise le calcul 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 libre des risques liés à l’exécution de code arbitraire localement.

Remarque

Les fonctions du catalogue Unity exécutées en tant qu’outils d’agent IA nécessitent un traitement générique sans serveur (Spark Connect serverless), et non pas des entrepôts SQL serverless. Les tentatives d’exécution d’outils sans utilisation du calcul générique sans serveur produisent des erreurs telles que PERMISSION_DENIED: Cannot access Spark Connect.

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

Lorsque votre agent demande une exécution d’outil en mode serverless , voici ce qui suit :

  1. Envoie DatabricksFunctionClient une demande au catalogue Unity pour récupérer la définition de fonction si la définition n’a pas été mise en cache localement.
  2. Le DatabricksFunctionClient extrait la définition de la fonction et valide les noms et les types des paramètres.
  3. DatabricksFunctionClient soumet l'exécution en tant qu’UDF au calcul générique sans serveur.

Mode local pour le développement

Le mode local lance des fonctions Python dans un sous-processus local au lieu d’effectuer des requêtes à un calcul générique sans serveur. Cela vous permet de résoudre les problèmes d’appels d’outils plus efficacement en fournissant des traces de pile locales. Il est conçu pour développer et déboguer les fonctions du catalogue Unity pour Python.

Lorsque votre agent demande l’exécution d’un outil en mode local , procédez DatabricksFunctionClient comme suit :

  1. Envoie une demande au catalogue Unity 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 d'une fonction Python pouvant être appelée, met en cache la fonction localement et valide les noms et types des paramètres.
  3. Appelle le joignable avec les paramètres spécifiés dans un sous-processus restreint avec une protection du délai d’expiration.
# Defaults to serverless if `execution_mode` is not specified
client = DatabricksFunctionClient(execution_mode="local")

L’exécution en "local" mode fournit les fonctionnalités suivantes :

  • Limite de temps processeur : Limite le runtime d’UC total pour l’exécution pouvant être appelée afin d’éviter des charges de calcul excessives.

    La limite de temps processeur est basée sur l’utilisation réelle du processeur, et non sur l’heure de l’horloge murale. En raison de la planification système et des processus simultanés, le temps processeur peut dépasser le temps Horloge dans les scénarios réels.

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

  • Protection du délai d’expiration : applique un délai d’expiration Horloge total pour les fonctions en cours d’exécution.

Personnalisez ces limites à l’aide de variables d’environnement (en lecture plus loin).

Limitations du mode local

  • Fonctions Python uniquement : les fonctions SQL ne sont pas prises en charge en mode local.
  • Considérations relatives à la sécurité pour le code non approuvé : alors que le mode local exécute des 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 IA. Il s'agit principalement d'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 la bibliothèque : les versions de la bibliothèque peuvent différer entre les environnements d’exécution serverless et local, ce qui peut entraîner un comportement de fonction différent.

Variables d'environnement

Configurez la façon dont les fonctions s’exécutent dans 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 d’exécution du processeur autorisée (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 maximale totale de l’horloge murale (mode local uniquement).
UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS 5 Nombre maximal de tentatives d’actualisation du client de session en cas d’expiration du jeton.
UCAI_DATABRICKS_SERVERLESS_EXECUTION_RESULT_ROW_LIMIT 100 Nombre maximal de lignes à retourner lors de l’exécution de fonctions à l’aide du calcul sans serveur et databricks-connect.

Exemples de notebooks

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

Outil d’agent de messagerie Slack

Obtenir un ordinateur portable

Outil d'agent de l’API Microsoft Graph

Obtenir un ordinateur portable

outil de l’agent Recherche Azure AI

Obtenir un ordinateur portable

Étapes suivantes