Intégrer OpenAI aux outils de catalogue Databricks Unity

Utilisez databricks Unity Catalog pour intégrer des fonctions SQL et Python en tant qu’outils dans les flux de travail OpenAI. Cette intégration combine la gouvernance du catalogue Unity avec OpenAI pour créer des applications IA de génération puissantes.

Spécifications

• Installez Python 3.10 ou version ultérieure.

Intégrer des outils de catalogue Unity à OpenAI

Exécutez le code suivant dans un notebook ou un script Python pour créer un outil de catalogue Unity et l’utiliser lors de l’appel d’un modèle OpenAI.

1. Installez le package d’intégration databricks Unity Catalog pour OpenAI.

   %pip install unitycatalog-openai[databricks]
   %pip install mlflow -U
   dbutils.library.restartPython()
   

2. Créez une instance du client de fonctions de catalogue Unity.

   from unitycatalog.ai.core.base import get_uc_function_client
   
   client = get_uc_function_client()
   

3. Créez une fonction de catalogue Unity écrite en Python.

   CATALOG = "your_catalog"
   SCHEMA = "your_schema"
   
   func_name = f"{CATALOG}.{SCHEMA}.code_function"
   
   def code_function(code: str) -> str:
     """
     Runs Python code.
   
     Args:
       code (str): The python code to run.
     Returns:
       str: The result of running the Python code.
     """
     import sys
     from io import StringIO
     stdout = StringIO()
     sys.stdout = stdout
     exec(code)
     return stdout.getvalue()
   
   client.create_python_function(
     func=code_function,
     catalog=CATALOG,
     schema=SCHEMA,
     replace=True
   )
   

4. Créez une instance de la fonction Catalogue Unity en tant que kit de ressources et vérifiez que l’outil se comporte correctement en exécutant la fonction.

   from unitycatalog.ai.openai.toolkit import UCFunctionToolkit
   import mlflow
   
   # Enable tracing
   mlflow.openai.autolog()
   
   # Create a UCFunctionToolkit that includes the UC function
   toolkit = UCFunctionToolkit(function_names=[func_name])
   
   # Fetch the tools stored in the toolkit
   tools = toolkit.tools
   client.execute_function = tools[0]
   

5. Envoyez la demande au modèle OpenAI avec les outils.

   import openai
   
   messages = [
     {
       "role": "system",
       "content": "You are a helpful customer support assistant. Use the supplied tools to assist the user.",
     },
     {"role": "user", "content": "What is the result of 2**10?"},
   ]
   response = openai.chat.completions.create(
     model="gpt-4o-mini",
     messages=messages,
     tools=tools,
   )
   # check the model response
   print(response)
   

6. Une fois qu'OpenAI a retourné une réponse, appelez la fonction Unity Catalog pour retourner la réponse à OpenAI.

   import json
   
   # OpenAI sends only a single request per tool call
   tool_call = response.choices[0].message.tool_calls[0]
   # Extract arguments that the Unity Catalog function needs to run
   arguments = json.loads(tool_call.function.arguments)
   
   # Run the function based on the arguments
   result = client.execute_function(func_name, arguments)
   print(result.value)
   

7. Une fois la réponse retournée, vous pouvez construire la charge utile de réponse pour les appels suivants à OpenAI.

   # Create a message containing the result of the function call
   function_call_result_message = {
     "role": "tool",
     "content": json.dumps({"content": result.value}),
     "tool_call_id": tool_call.id,
   }
   assistant_message = response.choices[0].message.to_dict()
   completion_payload = {
     "model": "gpt-4o-mini",
     "messages": [*messages, assistant_message, function_call_result_message],
   }
   
   # Generate final response
   openai.chat.completions.create(
     model=completion_payload["model"], messages=completion_payload["messages"]
   )
   

Services publics

Pour simplifier le processus d’élaboration de la réponse de l’outil, le ucai-openai package a un utilitaire, generate_tool_call_messagesqui convertit les messages de réponse OpenAI ChatCompletion afin qu’ils puissent être utilisés pour la génération de réponse.

from unitycatalog.ai.openai.utils import generate_tool_call_messages

messages = generate_tool_call_messages(response=response, client=client)
print(messages)