将 Unity 目录工具与第三方生成 AI 框架集成

Unity 目录 AI 代理工具可用于常用的 Gen AI 库,如 LangChain、LlamaIndex、OpenAI 和 Anthropic。 这些集成将 Unity 目录工具治理与第三方代理创作框架的功能相结合。 例如:

  • 在 LangChain 中,Unity 目录函数可以是代理工作流的一部分,用于执行查询或转换数据等任务。
  • 在 OpenAI 或人类集成中,函数在执行过程中由 AI 模型直接调用。

在以下选项卡中选择框架以创建 Unity 目录工具并将其用于该框架。 在Azure Databricks笔记本或Python脚本中运行代码。

Requirements

  • 安装 Python 3.10 或更高版本。

LangChain

使用 Azure Databricks Unity 目录将 SQL 和 Python 函数作为 LangChain 和 LangGraph 工作流中的工具进行集成。 此集成将 Unity 目录的治理与 LangChain 功能相结合,以构建功能强大的基于 LLM 的应用程序。

在此示例中,你将创建 Unity 目录工具,测试其功能,并将其添加到代理。

安装依赖项

使用 Databricks 可选安装 Unity 目录 AI 包并安装 LangChain 集成包。

# 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()

初始化 Databricks 函数客户端

初始化 Databricks 功能客户端。

from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

定义工具的逻辑

创建包含工具逻辑的 Unity 目录函数。


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
)

测试函数

测试函数以检查其是否按预期工作:

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'

使用 UCFunctionToolKit 封装函数

使用 UCFunctionToolkit 函数包装函数,使代理创作库可以访问该函数。 该工具包可确保不同库的一致性,并添加了有用的功能,例如检索器的自动跟踪。

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

在代理中使用该工具

使用 tools 属性从 UCFunctionToolkit 中将该工具添加到 LangChain 代理。

此示例使用 LangChain 的 AgentExecutor API 创建一个简单的代理,以便于简单。 对于生产工作负荷,请使用 在 Author an AI 代理中看到的代理创作工作流,并将其部署到 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 (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

使用 Azure Databricks Unity 目录将 SQL 和 Python 函数作为 LlamaIndex 工作流中的工具进行集成。 此集成将 Unity Catalog 治理与 LlamaIndex 的功能相结合,以索引和查询 LLM 的大型数据集。

  1. 安装适用于 LlamaIndex 的 Databricks Unity Catalog 集成包。

    %pip install unitycatalog-llamaindex[databricks]
    dbutils.library.restartPython()
    
  2. 创建 Unity 目录函数客户端的实例。

    from unitycatalog.ai.core.base import get_uc_function_client
    
    client = get_uc_function_client()
    
  3. 创建用 Python 编写的 Unity 目录函数。

    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. 创建 Unity 目录函数的实例作为工具包,并运行它以验证该工具是否正常运行。

    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. 通过将 Unity Catalog 函数定义为 LlamaIndex 工具集合的一部分,在 LlamaIndex ReActAgent 中使用该工具。 然后通过调用 LlamaIndex 工具集合来验证代理是否行为正常。

    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

使用 Azure Databricks Unity 目录将 SQL 和 Python 函数作为 OpenAI 工作流中的工具进行集成。 此集成结合了 Unity 目录与 OpenAI 的治理,以创建功能强大的 Gen AI 应用。

  1. 安装用于 OpenAI 的 Databricks Unity 目录集成包。

    %pip install unitycatalog-openai[databricks]
    %pip install mlflow -U
    dbutils.library.restartPython()
    
  2. 创建 Unity 目录函数客户端的实例。

    from unitycatalog.ai.core.base import get_uc_function_client
    
    client = get_uc_function_client()
    
  3. 创建用 Python 编写的 Unity 目录函数。

    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. 创建 Unity 目录函数的实例作为工具包,并通过运行函数来验证该工具的行为是否正确。

    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. 将请求与工具一起提交到 OpenAI 模型。

    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. OpenAI 返回响应后,调用 Unity 目录函数生成响应并返回给 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. 返回答案后,可以为对 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"]
    )
    

Utilities

为了简化工具响应的创建过程,包 ucai-openai 具有一个实用工具, generate_tool_call_messages用于转换 OpenAI ChatCompletion 响应消息,以便它们可用于响应生成。

from unitycatalog.ai.openai.utils import generate_tool_call_messages

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

注释

如果响应包含多个选择项,则可以在调用generate_tool_call_messages时传递choice_index参数,以选择要利用的选择项。 目前不支持处理多个选择项。

Anthropic

使用 Azure Databricks Unity 目录将 SQL 和 Python 函数作为Anthropic SDK LLM 调用中的工具进行集成。 此集成将 Unity Catalog 的治理与 Anthropic 模型相结合,以创建强大的 gen AI 应用程序。

注释

Anthropic 集成需要 Databricks Runtime 15.0 及更高版本。

  1. 安装适用于 Anthropic 的 Databricks Unity Catalog 集成包。

    %pip install unitycatalog-anthropic[databricks]
    dbutils.library.restartPython()
    
  2. 创建 Unity 目录函数客户端的实例。

    from unitycatalog.ai.core.base import get_uc_function_client
    
    client = get_uc_function_client()
    
  3. 创建用 Python 编写的 Unity 目录函数。

    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. 将 Unity Catalog 函数的实例创建为工具包。

    from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit
    
    # Create an instance of the toolkit
    toolkit = UCFunctionToolkit(function_names=[func_name], client=client)
    
  5. 在 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. 构建工具响应。 如果需要调用工具,则来自 Claude 模型的响应包含一个工具请求元数据块。

    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)
    

unitycatalog.ai-anthropic 包包括一个消息处理程序实用程序,用于简化对 Unity Catalog 函数的调用的解析和处理。 该实用程序执行以下作:

  1. 检测工具调用需求。
  2. 从查询中提取工具调用信息。
  3. 执行对 Unity Catalog 函数的调用。
  4. 分析来自 Unity Catalog 函数的响应。
  5. 制作下一个消息格式以继续与 Claude 的对话。

注释

必须在 conversation_history API 的 generate_tool_call_messages 参数中提供整个对话历史记录。 Claude 模型需要初始化对话(原始用户输入问题)以及所有后续 LLM 生成的响应和多轮次工具调用结果。