Integrieren von Unity-Katalogtools in generative KI-Frameworks von Drittanbietern

Unity Catalog AI Agent Tools können in gängigen Bibliotheken für generative KI wie LangChain, LlamaIndex, OpenAI und Anthropic verwendet werden. Diese Integrationen kombinieren die Tool-Governance von Unity Catalog mit den Funktionen von Drittanbieter-Agentenentwicklungs-Frameworks. Beispiel:

• In LangChain können Unity-Katalogfunktionen Teil des Workflows eines Agents sein, um Aufgaben wie Abfragen oder Transformieren von Daten auszuführen.
• In OpenAI- oder Anthropic-Integrationen werden die Funktionen während der Ausführung direkt vom KI-Modell aufgerufen.

Wählen Sie Ihr Framework auf den folgenden Registerkarten aus, um ein Unity-Katalogtool zu erstellen und mit diesem Framework zu verwenden. Führen Sie den Code in einem Azure Databricks Notizbuch oder Python Skript aus.

Anforderungen

• Installieren Sie Python 3.10 oder höher.

LangChain

Verwenden Sie Azure Databricks Unity-Katalog, um SQL- und Python-Funktionen als Tools in LangChain- und LangGraph-Workflows zu integrieren. Diese Integration kombiniert die Governance des Unity-Katalogs mit LangChain-Funktionen, um leistungsstarke LLM-basierte Anwendungen zu erstellen.

In diesem Beispiel erstellen Sie ein Unity Catalog-Tool, testen seine Funktionalität und fügen es einem Agent hinzu.

Abhängigkeiten installieren

Installieren Sie Unity Catalog AI-Pakete mit den optionalen Databricks-Paketen, und installieren Sie das LangChain-Integrationspaket.

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

# Install Databricks Langchain integration package
%pip install databricks-langchain
dbutils.library.restartPython()

Initialisieren des Databricks-Funktionsclients

Initialisieren Sie den Databricks-Funktionsclient.

from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

Definieren der Logik des Tools

Erstellen Sie eine Unity-Katalogfunktion, die die Logik des Tools enthält.

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
)

Testen der Funktion

Testen Sie Ihre Funktion so, dass sie wie erwartet funktioniert:

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'

Die Funktion mithilfe des UCFunctionToolKit umschließen

Umschließen Sie die Funktion mithilfe des UCFunctionToolkit, um sie für die Agenterstellungsbibliotheken zugänglich zu machen. Das Toolkit sorgt für Konsistenz in verschiedenen Bibliotheken und fügt hilfreiche Funktionen wie das automatische Tracking für Retriever hinzu.

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

Verwenden des Tools in einem Agent

Fügen Sie das Tool einem LangChain-Agent mithilfe der Eigenschaft tools aus dem UCFunctionToolkit hinzu.

In diesem Beispiel wird der Einfachheit halber ein einfacher Agent mit der LangChain-API AgentExecutor erstellt. Verwenden Sie für Produktionsworkloads den Agenterstellungsworkflow, der in "Author an AI"-Agent angezeigt wird, und stellen Sie ihn in Databricks-Apps bereit.

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 (replace with your LLM of choice, if desired)
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?"})

LlamaIndex

Verwenden Sie Azure Databricks Unity-Katalog, um SQL- und Python-Funktionen als Tools in LlamaIndex-Workflows zu integrieren. Diese Integration kombiniert Unity Catalog Governance mit den Funktionen von LlamaIndex zum Indizieren und Abfragen großer Datasets für LLMs.

1. Installieren Sie das Integrationspaket "Databricks Unity Catalog" für LlamaIndex.

   %pip install unitycatalog-llamaindex[databricks]
   dbutils.library.restartPython()
   

2. Erstellen Sie eine Instanz des Unity Catalog-Funktionsclients.

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

3. Erstellen Sie eine Unity Catalog-Funktion, die in Python geschrieben wurde.

   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. Erstellen Sie eine Instanz des Unity-Katalogs als Toolkit, und führen Sie sie aus, um zu überprüfen, ob sich das Tool ordnungsgemäß verhält.

   from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit
   import mlflow
   
   # Enable traces
   mlflow.llama_index.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
   python_exec_tool = tools[0]
   
   # Run the tool directly
   result = python_exec_tool.call(code="print(1 + 1)")
   print(result)  # Outputs: {"format": "SCALAR", "value": "2\n"}
   

5. Verwenden Sie das Tool in einem LlamaIndex ReActAgent, indem Sie die Unity-Katalogfunktion als Teil einer LlamaIndex-Toolsammlung definieren. Überprüfen Sie dann, ob sich der Agent ordnungsgemäß verhält, indem Sie die LlamaIndex-Toolsammlung aufrufen.

   from llama_index.llms.openai import OpenAI
   from llama_index.core.agent import ReActAgent
   
   llm = OpenAI()
   
   agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)
   
   agent.chat("Please run the following python code: `print(1 + 1)`")
   

OpenAI

Verwenden Sie Azure Databricks Unity-Katalog, um SQL- und Python-Funktionen als Tools in OpenAI-Workflows zu integrieren. Diese Integration kombiniert die Governance des Unity-Katalogs mit OpenAI, um leistungsstarke KI-Apps der Generation zu erstellen.

1. Installieren Sie das Integrationspaket "Databricks Unity Catalog" für OpenAI.

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

2. Erstellen Sie eine Instanz des Unity Catalog-Funktionsclients.

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

3. Erstellen Sie eine Unity Catalog-Funktion, die in Python geschrieben wurde.

   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. Erstellen Sie eine Instanz des Unity-Katalogs als Toolkit, und stellen Sie sicher, dass sich das Tool ordnungsgemäß verhält, indem Sie die Funktion ausführen.

   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. Übermitteln Sie die Anforderung zusammen mit den Tools an das OpenAI-Modell.

   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. Nachdem OpenAI eine Antwort zurückgegeben hat, rufen Sie den Unity Catalog-Funktionsaufruf auf, um die Antwortantwort zurück zu OpenAI zu generieren.

   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. Nachdem die Antwort zurückgegeben wurde, können Sie die Antwortnutzlast für nachfolgende Aufrufe an OpenAI erstellen.

   # 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"]
   )
   

Versorgungsdienste

Um das Erstellen der Toolantwort zu vereinfachen, verfügt das ucai-openai Paket über ein Dienstprogramm, generate_tool_call_messagesdas OpenAI ChatCompletion-Antwortnachrichten konvertiert, sodass sie für die Antwortgenerierung verwendet werden können.

from unitycatalog.ai.openai.utils import generate_tool_call_messages

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

Anthropic

Verwenden Sie Azure Databricks Unity-Katalog, um SQL- und Python-Funktionen als Tools in Anthropic SDK LLM-Aufrufe zu integrieren. Diese Integration kombiniert die Governance von Unity Catalog mit anthropischen Modellen, um leistungsstarke KI-Apps für die Generation zu erstellen.

1. Installieren Sie das Databricks Unity Catalog-Integrationspaket für Anthropic.

   %pip install unitycatalog-anthropic[databricks]
   dbutils.library.restartPython()
   

2. Erstellen Sie eine Instanz des Unity Catalog-Funktionsclients.

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

3. Erstellen Sie eine Unity Catalog-Funktion, die in Python geschrieben wurde.

   CATALOG = "your_catalog"
   SCHEMA = "your_schema"
   
   func_name = f"{CATALOG}.{SCHEMA}.weather_function"
   
   def weather_function(location: str) -> str:
     """
     Fetches the current weather from a given location in degrees Celsius.
   
     Args:
       location (str): The location to fetch the current weather from.
     Returns:
       str: The current temperature for the location provided in Celsius.
     """
     return f"The current temperature for {location} is 24.5 celsius"
   
   client.create_python_function(
     func=weather_function,
     catalog=CATALOG,
     schema=SCHEMA,
     replace=True
   )
   

4. Erstellen Sie eine Instanz der Unity Catalog-Funktion als Toolkit.

   from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit
   
   # Create an instance of the toolkit
   toolkit = UCFunctionToolkit(function_names=[func_name], client=client)
   

5. Verwenden Sie einen Werkzeugaufruf in Anthropic.

   import anthropic
   
   # Initialize the Anthropic client with your API key
   anthropic_client = anthropic.Anthropic(api_key="YOUR_ANTHROPIC_API_KEY")
   
   # User's question
   question = [{"role": "user", "content": "What's the weather in New York City?"}]
   
   # Make the initial call to Anthropic
   response = anthropic_client.messages.create(
     model="claude-3-5-sonnet-20240620",  # Specify the model
     max_tokens=1024,  # Use 'max_tokens' instead of 'max_tokens_to_sample'
     tools=toolkit.tools,
     messages=question  # Provide the conversation history
   )
   
   # Print the response content
   print(response)
   

6. Erstellen Sie eine Tool-Antwort. Die Antwort des Claude-Modells enthält einen Metadatenblock für die Werkzeuganforderung, wenn ein Werkzeug aufgerufen werden muss.

   from unitycatalog.ai.anthropic.utils import generate_tool_call_messages
   
   # Call the UC function and construct the required formatted response
   tool_messages = generate_tool_call_messages(
     response=response,
     client=client,
     conversation_history=question
   )
   
   # Continue the conversation with Anthropic
   tool_response = anthropic_client.messages.create(
     model="claude-3-5-sonnet-20240620",
     max_tokens=1024,
     tools=toolkit.tools,
     messages=tool_messages,
   )
   
   print(tool_response)
   

Das unitycatalog.ai-anthropic Paket enthält ein Nachrichtenhandler-Hilfsprogramm, um das Analysieren und Verarbeiten eines Aufrufs der Unity Catalog-Funktion zu vereinfachen. Das Dienstprogramm führt folgende Schritte aus:

1. Erkennt die Anforderungen an Tool-Aufrufe.
2. Extrahiert Informationen zum Werkzeugaufruf aus der Abfrage.
3. Führt den Aufruf der Unity Catalog-Funktion aus.
4. Analysiert die Antwort aus der Unity Catalog-Funktion.
5. Erstellen Sie das nächste Nachrichtenformat, um das Gespräch mit Claude fortzusetzen.