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.

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 d’exécuter uniquement du code approuvé et d’implémenter des garde-fous et des autorisations appropriées pour empêcher l’accès involontaire aux données.

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
• Version de Python : 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 exécuter des fonctions de catalogue Unity en tant qu’outils d’agent IA en production. Consultez les exigences en matière de calcul serverless.
  • L’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 d’agent

Dans cet exemple, vous créez un outil de catalogue Unity, testez ses fonctionnalités et ajoutez-le à un agent. Exécutez le code suivant dans un notebook Databricks.

Installer des dépendances

Installez les packages UNITY Catalog AI avec l’élément [databricks] supplémentaire, puis installez le package d’intégration Databricks-LangChain.

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.

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

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

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.

Vous pouvez créer des fonctions de catalogue Unity à l’aide de l’une des deux API suivantes :

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

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

Pour rendre un appelant Python reconnaissable au modèle de données des fonctions Unity Catalog, votre fonction doit répondre aux exigences suivantes :

• Indicateurs de type : la signature de fonction doit définir des indicateurs 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é des types : 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.

Les extraits de code suivants utilisent create_python_function pour enregistrer la fonction 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'

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 les exemples ResponsesAgent.

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écution de 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, ce qui 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 catalogue Unity exécutées en tant qu’outils d’agent IA nécessitent un calcul générique serverless (Spark Connect serverless), et non pas des entrepôts SQL serverless. Les tentatives d’exécution d’outils sans calcul générique serverless 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. L’exécution DatabricksFunctionClient envoie l’exécution en tant qu’UDF au calcul générique serverless.

Mode local pour le développement

Le mode local exécute des fonctions Python dans un sous-processus local au lieu d’effectuer des requêtes à un calcul générique serverless. 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 des fonctions catalogue Python Unity.

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 appelante Python, met en cache l’appelant localement et valide les noms et types de 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 basées sur 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.

Étapes suivantes

• Ajoutez des outils de catalogue Unity aux agents par programmation. Voir ResponsesAgent exemples.

• Ajoutez des outils de catalogue Unity aux agents à l’aide de l’interface utilisateur AI Playground. Consultez Prototype d’agent d’appel d’outil dans AI Playground.

• Gérer les fonctions de catalogue Unity à l’aide du client de fonction. Consultez la documentation de Unity Catalog – Client fonction