将 OpenAI 与 Databricks Unity 目录工具集成

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

要求

  • 安装 Python 3.10 或更高版本。

将 Unity 目录工具与 OpenAI 集成

在笔记本或 Python 脚本中运行以下代码,以创建 Unity 目录工具,并在调用 OpenAI 模型时使用。

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

公用事业

为了简化工具响应的创建过程,包 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参数,以选择要利用的选择项。 目前不支持处理多个选择项。