使用 mlflow.pyfunc.ChatModel 构建工具调用模型

下载此笔记本

欢迎来到关于使用 mlflow.pyfunc.ChatModel 包装器构建简单工具调用模型的笔记本教程。ChatModel 是 MLflow 高度可定制的 PythonModel 的子类，它专门设计用于简化 GenAI 工作流的创建。

简而言之，使用 ChatModel 的一些好处是：

1. 无需定义复杂的签名！聊天模型通常接受具有多级嵌套的复杂输入，这自行定义起来可能很麻烦。
2. 支持 JSON / dict 输入（无需包装输入或转换为 Pandas DataFrame）
3. 包含使用 Dataclasses 定义预期输入/输出，以简化开发体验

要更深入地了解 ChatModel，请参阅 详细指南。

在本教程中，我们将构建一个简单的 OpenAI 包装器，该包装器利用工具调用支持（在 MLflow 2.17.0 中发布）。

环境设置

首先，让我们设置环境。我们将需要 OpenAI Python SDK，以及 MLflow >= 2.17.0。我们还需要设置 OpenAI API 密钥才能使用 SDK。

%pip install 'mlflow>=2.17.0' 'openai>=1.0' -qq

Note: you may need to restart the kernel to use updated packages.

import os
from getpass import getpass

os.environ["OPENAI_API_KEY"] = getpass("Enter your OpenAI API key: ")

步骤 1：创建工具定义

让我们开始定义我们的模型！如引言中所述，我们将继承 mlflow.pyfunc.ChatModel。在此示例中，我们将构建一个玩具模型，该模型使用工具来检索给定城市的天气。

第一步是创建我们可以传递给 OpenAI 的工具定义。我们通过使用 mlflow.types.llm.FunctionToolDefinition 来描述我们的工具接受的参数。此 dataclass 的格式与 OpenAI 规范一致。

import mlflow
from mlflow.types.llm import (
  FunctionToolDefinition,
  ParamProperty,
  ToolParamsSchema,
)


class WeatherModel(mlflow.pyfunc.ChatModel):
  def __init__(self):
      # a sample tool definition. we use the `FunctionToolDefinition`
      # class to describe the name and expected params for the tool.
      # for this example, we're defining a simple tool that returns
      # the weather for a given city.
      weather_tool = FunctionToolDefinition(
          name="get_weather",
          description="Get weather information",
          parameters=ToolParamsSchema(
              {
                  "city": ParamProperty(
                      type="string",
                      description="City name to get weather information for",
                  ),
              }
          ),
          # make sure to call `to_tool_definition()` to convert the `FunctionToolDefinition`
          # to a `ToolDefinition` object. this step is necessary to normalize the data format,
          # as multiple types of tools (besides just functions) might be available in the future.
      ).to_tool_definition()

      # OpenAI expects tools to be provided as a list of dictionaries
      self.tools = [weather_tool.to_dict()]

步骤 2：实现工具

现在我们有了工具的定义，我们需要实际实现它。在本教程中，我们只是模拟一个响应，但实现可以是任意的——例如，您可以调用实际的天气服务 API。

class WeatherModel(mlflow.pyfunc.ChatModel):
  def __init__(self):
      weather_tool = FunctionToolDefinition(
          name="get_weather",
          description="Get weather information",
          parameters=ToolParamsSchema(
              {
                  "city": ParamProperty(
                      type="string",
                      description="City name to get weather information for",
                  ),
              }
          ),
      ).to_tool_definition()

      self.tools = [weather_tool.to_dict()]

      def get_weather(self, city: str) -> str:
          # in a real-world scenario, the implementation might be more complex
          return f"It's sunny in {city}, with a temperature of 20C"

步骤 3：实现 predict 方法

接下来我们需要做的是定义一个 predict() 函数，该函数接受以下参数：

1. context: PythonModelContext (本教程未使用)
2. messages: List[ChatMessage]。这是模型用于生成的聊天输入。
3. params: ChatParams。这些是常用的参数，用于配置聊天模型，例如 temperature、max_tokens 等。工具规范可以在这里找到。

这是推理期间最终将被调用的函数。

在实现中，我们将简单地将用户的输入转发给 OpenAI，并提供 get_weather 工具作为 LLM 可选的工具。如果我们收到工具调用请求，我们将调用 get_weather() 函数并将响应返回给 OpenAI。我们需要使用我们在前两个步骤中定义的来完成此操作。

import json

from openai import OpenAI

import mlflow
from mlflow.types.llm import (
  ChatMessage,
  ChatParams,
  ChatResponse,
)


class WeatherModel(mlflow.pyfunc.ChatModel):
  def __init__(self):
      weather_tool = FunctionToolDefinition(
          name="get_weather",
          description="Get weather information",
          parameters=ToolParamsSchema(
              {
                  "city": ParamProperty(
                      type="string",
                      description="City name to get weather information for",
                  ),
              }
          ),
      ).to_tool_definition()

      self.tools = [weather_tool.to_dict()]

  def get_weather(self, city: str) -> str:
      return "It's sunny in {}, with a temperature of 20C".format(city)

  # the core method that needs to be implemented. this function
  # will be called every time a user sends messages to our model
  def predict(self, context, messages: list[ChatMessage], params: ChatParams):
      # instantiate the OpenAI client
      client = OpenAI()

      # convert the messages to a format that the OpenAI API expects
      messages = [m.to_dict() for m in messages]

      # call the OpenAI API
      response = client.chat.completions.create(
          model="gpt-4o-mini",
          messages=messages,
          # pass the tools in the request
          tools=self.tools,
      )

      # if OpenAI returns a tool_calling response, then we call
      # our tool. otherwise, we just return the response as is
      tool_calls = response.choices[0].message.tool_calls
      if tool_calls:
          print("Received a tool call, calling the weather tool...")

          # for this example, we only provide the model with one tool,
          # so we can assume the tool call is for the weather tool. if
          # we had more, we'd need to check the name of the tool that
          # was called
          city = json.loads(tool_calls[0].function.arguments)["city"]
          tool_call_id = tool_calls[0].id

          # call the tool and construct a new chat message
          tool_response = ChatMessage(
              role="tool", content=self.get_weather(city), tool_call_id=tool_call_id
          ).to_dict()

          # send another request to the API, making sure to append
          # the assistant's tool call along with the tool response.
          messages.append(response.choices[0].message)
          messages.append(tool_response)
          response = client.chat.completions.create(
              model="gpt-4o-mini",
              messages=messages,
              tools=self.tools,
          )

      # return the result as a ChatResponse, as this
      # is the expected output of the predict method
      return ChatResponse.from_dict(response.to_dict())

步骤 4（可选，但推荐）：为模型启用跟踪

此步骤是可选的，但强烈建议用于提高应用程序的可观测性。我们将使用 MLflow Tracing 来记录我们模型内部函数的输入和输出，以便在出现问题时可以轻松调试。类似代理的工具调用模型在单个请求的生命周期中可能会进行多层函数调用，因此跟踪对于帮助我们理解每个步骤中的情况非常有价值。

集成跟踪很简单，我们只需使用 @mlflow.trace 装饰我们感兴趣的函数（get_weather() 和 predict()）！MLflow Tracing 还与许多流行的 GenAI 框架（如 LangChain、OpenAI、LlamaIndex 等）集成。有关完整列表，请查看此 文档页面。在本教程中，我们使用 OpenAI SDK 进行 API 调用，因此可以通过调用 mlflow.openai.autolog() 来为该调用启用跟踪。

要查看 UI 中的跟踪，请在单独的终端 shell 中运行 mlflow ui，并在下面使用模型进行推理后导航到 Traces 选项卡。

from mlflow.entities.span import (
  SpanType,
)

# automatically trace OpenAI SDK calls
mlflow.openai.autolog()


class WeatherModel(mlflow.pyfunc.ChatModel):
  def __init__(self):
      weather_tool = FunctionToolDefinition(
          name="get_weather",
          description="Get weather information",
          parameters=ToolParamsSchema(
              {
                  "city": ParamProperty(
                      type="string",
                      description="City name to get weather information for",
                  ),
              }
          ),
      ).to_tool_definition()

      self.tools = [weather_tool.to_dict()]

  @mlflow.trace(span_type=SpanType.TOOL)
  def get_weather(self, city: str) -> str:
      return "It's sunny in {}, with a temperature of 20C".format(city)

  @mlflow.trace(span_type=SpanType.AGENT)
  def predict(self, context, messages: list[ChatMessage], params: ChatParams):
      client = OpenAI()

      messages = [m.to_dict() for m in messages]

      response = client.chat.completions.create(
          model="gpt-4o-mini",
          messages=messages,
          tools=self.tools,
      )

      tool_calls = response.choices[0].message.tool_calls
      if tool_calls:
          print("Received a tool call, calling the weather tool...")

          city = json.loads(tool_calls[0].function.arguments)["city"]
          tool_call_id = tool_calls[0].id

          tool_response = ChatMessage(
              role="tool", content=self.get_weather(city), tool_call_id=tool_call_id
          ).to_dict()

          messages.append(response.choices[0].message)
          messages.append(tool_response)
          response = client.chat.completions.create(
              model="gpt-4o-mini",
              messages=messages,
              tools=self.tools,
          )

      return ChatResponse.from_dict(response.to_dict())

步骤 5：记录模型

最后，我们需要记录模型。这会将模型作为 artifact 保存在 MLflow Tracking 中，并允许我们稍后加载和部署它。

（注意：这是 MLflow 中的一个基本模式。要了解更多信息，请查看 快速入门指南！）

为了做到这一点，我们需要做几件事：

1. 定义一个输入示例，告知用户我们期望的输入
2. 实例化模型
3. 使用上述参数调用 mlflow.pyfunc.log_model()

注意单元格末尾打印的模型 URI——稍后部署模型时我们会用到它！

# messages to use as input examples
messages = [
  {"role": "system", "content": "Please use the provided tools to answer user queries."},
  {"role": "user", "content": "What's the weather in Singapore?"},
]

input_example = {
  "messages": messages,
}

# instantiate the model
model = WeatherModel()

# log the model
with mlflow.start_run():
  model_info = mlflow.pyfunc.log_model(
      name="weather-model",
      python_model=model,
      input_example=input_example,
  )

  print("Successfully logged the model at the following URI: ", model_info.model_uri)

2024/10/29 09:30:14 INFO mlflow.pyfunc: Predicting on input example to validate output

Received a tool call, calling the weather tool...

Downloading artifacts:   0%|          | 0/7 [00:00<?, ?it/s]

Received a tool call, calling the weather tool...
Successfully logged the model at the following URI:  runs:/8051850efa194a3b8b2450c4c9f4d42f/weather-model

使用模型进行推理

模型已记录后，我们的工作基本就完成了！为了使用模型进行推理，让我们使用 mlflow.pyfunc.load_model() 将其加载回来。

import mlflow

# Load the previously logged ChatModel
tool_model = mlflow.pyfunc.load_model(model_info.model_uri)

system_prompt = {
  "role": "system",
  "content": "Please use the provided tools to answer user queries.",
}

messages = [
  system_prompt,
  {"role": "user", "content": "What's the weather in Singapore?"},
]

# Call the model's predict method
response = tool_model.predict({"messages": messages})
print(response["choices"][0]["message"]["content"])

messages = [
  system_prompt,
  {"role": "user", "content": "What's the weather in San Francisco?"},
]

# Generating another response
response = tool_model.predict({"messages": messages})
print(response["choices"][0]["message"]["content"])

2024/10/29 09:30:27 WARNING mlflow.tracing.processor.mlflow: Creating a trace within the default experiment with id '0'. It is strongly recommended to not use the default experiment to log traces due to ambiguous search results and probable performance issues over time due to directory table listing performance degradation with high volumes of directories within a specific path. To avoid performance and disambiguation issues, set the experiment for your environment using `mlflow.set_experiment()` API.

Received a tool call, calling the weather tool...
The weather in Singapore is sunny, with a temperature of 20°C.
Received a tool call, calling the weather tool...
The weather in San Francisco is sunny, with a temperature of 20°C.

部署模型

MLflow 还允许您使用 mlflow models serve CLI 工具来部署模型。在另一个终端 shell 中，在与此笔记本相同的文件夹中运行以下命令：

$ export OPENAI_API_KEY=<YOUR OPENAI API KEY>
$ mlflow models serve -m <MODEL_URI>

这将启动模型在 http://127.0.0.1:5000 上部署，并且可以通过 POST 请求到 /invocations 路由来查询模型。

import requests

messages = [
  system_prompt,
  {"role": "user", "content": "What's the weather in Tokyo?"},
]

response = requests.post("http://127.0.0.1:5000/invocations", json={"messages": messages})
response.raise_for_status()
response.json()

{'choices': [{'index': 0,
 'message': {'role': 'assistant',
  'content': 'The weather in Tokyo is sunny, with a temperature of 20°C.'},
 'finish_reason': 'stop'}],
'usage': {'prompt_tokens': 100, 'completion_tokens': 16, 'total_tokens': 116},
'id': 'chatcmpl-ANVOhWssEiyYNFwrBPxp1gmQvZKsy',
'model': 'gpt-4o-mini-2024-07-18',
'object': 'chat.completion',
'created': 1730165599}

结论

在本教程中，我们介绍了如何使用 MLflow 的 ChatModel 类来创建一个方便的、支持工具调用的 OpenAI 包装器。尽管用例很简单，但这里涵盖的概念可以轻松扩展以支持更复杂的功能。

如果您想更深入地了解构建高质量的 GenAI 应用，您可能还对查看 MLflow Tracing 感兴趣，它是一个可用于跟踪任意函数执行（例如您的工具调用）的可观测性工具。