跳到主要内容

教程：使用 ChatModel 定制 GenAI 模型

注意

从 MLflow 3.0.0 开始，我们推荐使用 ResponsesAgent 而非 ChatModel。更多详情请参阅 ResponsesAgent 简介。

快速发展的生成式人工智能 (GenAI) 领域带来了令人兴奋的机会和集成挑战。为了有效利用最新的 GenAI 进展，开发人员需要一个能够平衡灵活性和标准化特性的框架。MLflow 通过在2.11.0 版本中引入的 mlflow.pyfunc.ChatModel 类来满足这一需求，为 GenAI 应用程序提供了一致的接口，同时简化了部署和测试。

ChatModel 和 PythonModel 的选择​

在 MLflow 中构建 GenAI 应用程序时，至关重要的是选择正确的模型抽象，以平衡易用性与所需的定制级别。MLflow 提供两个主要类来实现此目的：mlflow.pyfunc.ChatModel 和 mlflow.pyfunc.PythonModel。每个类都有其优点和缺点，因此理解哪个最适合您的用例至关重要。

ChatModelPythonModel
何时使用当您想开发和部署一个具有与 OpenAI 规范兼容的标准聊天模式的会话模型时使用。当您想对模型的接口进行完全控制或定制模型行为的各个方面时使用。
接口固定为 OpenAI 的聊天模式。对模型的输入和输出模式完全控制。
设置快速。开箱即用，适用于会话应用程序，具有预定义的模型签名和输入示例。自定义。您需要自己定义模型签名或输入示例。
复杂性低。标准化的接口简化了模型的部署和集成。高。部署和集成自定义 PythonModel 可能并不直接。例如，模型需要处理 Pandas DataFrames，因为 MLflow 在将输入数据传递给 PythonModel 之前会将其转换为 DataFrames。

本教程的目的​

本教程将指导您完成使用 MLflow 的 mlflow.pyfunc.ChatModel 类创建自定义聊天代理的过程。

在本教程结束时，您将能够

  • 将MLflow 跟踪集成到自定义的 mlflow.pyfunc.ChatModel 实例中。
  • 在 mlflow.pyfunc.log_model() 中使用 model_config 参数来自定义您的模型。
  • 利用标准化的签名接口简化部署。
  • 识别并避免扩展 mlflow.pyfunc.ChatModel 类时常见的陷阱。

先决条件​

  • 熟悉 MLflow 日志记录 API 和 GenAI 概念。
  • 已安装 MLflow 2.11.0 或更高版本以使用 mlflow.pyfunc.ChatModel。
  • 已安装 MLflow 2.14.0 或更高版本以使用MLflow 跟踪。

本教程将Databricks Foundation Model API 仅作为与外部服务交互的示例。您可以轻松地将提供商示例替换为任何托管的 LLM 托管服务（如Amazon Bedrock、Azure AI Studio、OpenAI、Anthropic 等等）。

核心概念​

GenAI 的跟踪定制​

MLflow 跟踪允许您监控和记录模型方法的执行情况，在调试和性能优化期间提供宝贵的见解。

在我们的 BasicAgent 示例实现中，我们使用两种独立的 API 来启动跟踪跨度：装饰器 API 和流畅 API。

装饰器 API​

@mlflow.trace
def _get_system_message(self, role: str) -> Dict:
    if role not in self.models:
        raise ValueError(f"Unknown role: {role}")

    instruction = self.models[role]["instruction"]
    return ChatMessage(role="system", content=instruction).to_dict()

使用 @mlflow.trace 跟踪装饰器是将跟踪功能添加到函数和方法的最简单方法。默认情况下，由此装饰器生成的跨度将使用函数名作为跨度名。可以通过以下方式覆盖此命名以及与跨度相关的其他参数：

@mlflow.trace(name="custom_span_name", attributes={"key": "value"}, span_type="func")
def _get_system_message(self, role: str) -> Dict:
    if role not in self.models:
        raise ValueError(f"Unknown role: {role}")

    instruction = self.models[role]["instruction"]
    return ChatMessage(role="system", content=instruction).to_dict()
提示

始终建议为生成的任何跨度设置一个人类可读的名称，特别是如果您正在对私有或通用命名的函数或方法进行仪器化。MLflow Trace UI 默认会显示函数或方法的名称，如果函数和方法的命名含糊不清，这可能会令人困惑。

流畅 API​

启动跨度的上下文处理器的流畅 API 实现，在您需要完全控制跨度数据的每个方面的日志记录时非常有用。

下面是我们的应用程序中确保在通过 load_context 方法加载模型时捕获设置的参数的示例。我们从实例属性 self.models_config 和 self.models 中提取数据来设置跨度的属性。

with mlflow.start_span("Audit Agent") as root_span:
    root_span.set_inputs(messages)
    attributes = {**params.to_dict(), **self.models_config, **self.models}
    root_span.set_attributes(attributes)
    # More span manipulation...

MLflow UI 中的跟踪​

在运行包含这些组合使用模式的跟踪跨度生成和仪器化示例后，

Traces in the MLflow UI for the Agent example

我们示例中的关键类和方法​

  • BasicAgent：我们的自定义聊天代理类，它扩展了 ChatModel。
  • _get_system_message：检索特定角色的系统消息配置。
  • get_agent_response：向端点发送消息并检索响应。
  • _call_agent：管理代理角色之间的对话流程。
  • prepare_message_list：准备要发送的消息列表。
  • load_context：初始化模型上下文和配置。
  • predict：处理聊天模型的预测逻辑。

在上面列出的这些方法中，load_context 和 predict 方法覆盖了 ChatModel 的基本抽象实现。为了定义 ChatModel 的子类，您必须（至少）实现 predict 方法。load_context 方法仅在使用（如下面的示例）自定义加载逻辑时使用，其中需要加载静态配置才能使模型对象工作，或者需要执行其他依赖逻辑才能使对象实例化正确运行。

自定义 ChatModel 示例​

在下面的完整示例中，我们通过子类化 mlflow.pyfunc.ChatModel 来创建一个自定义聊天代理。此代理名为 BasicAgent，它利用了几个重要功能，有助于简化 GenAI 应用程序的开发、部署和跟踪。通过子类化 ChatModel，我们确保了处理会话代理的一致接口，同时避免了更通用模型相关的常见陷阱。

以下实现强调了以下关键方面：

  • 跟踪：我们利用 MLflow 的跟踪功能来跟踪和记录关键操作，同时使用装饰器和流畅 API 上下文处理器方法。
    • 装饰器 API：用于轻松跟踪 _get_agent_response 和 _call_agent 等方法，以实现自动跨度创建。
    • 流畅 API：在代理交互期间审计关键输入和输出，如图 predict 方法所示，它提供了对跨度创建的细粒度控制。
    • 提示：我们确保了人类可读的跨度名称，以便在 MLflow Trace UI 中进行调试以及通过客户端 API 获取已记录的跟踪。
  • 自定义配置:
    • 模型配置：通过在模型日志记录期间传递自定义配置（使用 model_config 参数），我们将模型行为与硬编码值分离开。这允许快速测试不同的代理配置，而无需修改源代码。
    • load_context 方法：确保在运行时加载配置，并使用必要的设置初始化代理，防止因缺少配置而导致的运行时失败。
    • 提示：我们避免在 load_context 中直接设置未定义的实例属性。相反，所有属性都在类构造函数中用默认值初始化，以确保我们的模型正确加载。
  • 对话管理:
    • 我们使用 _get_system_message、_get_agent_response 和 _call_agent 等方法实现了多步代理交互模式。这些方法管理多个代理之间的通信流程，例如“预言家”和“裁判”角色，每个角色都配置有特定的指令和参数。
    • 静态输入/输出结构：通过遵循 ChatModel 所需的输入（List[ChatMessage]）和输出（ChatCompletionResponse）格式，我们消除了与转换 JSON 或表格数据相关的复杂性，而这在 PythonModel 等更通用的模型中很常见。
  • 避免常见陷阱:
    • 通过输入示例进行模型验证：我们在模型日志记录期间提供了一个输入示例，允许 MLflow 验证输入接口并及早捕获结构性问题，从而减少部署期间的调试时间。
  • 来自代码的模型:
    • MLflow 建议在编写 GenAI 代理或应用程序时使用来自代码的模型，以实现强大的日志记录和代理（包含任意 Python 代码）的直接部署。
import mlflow
from mlflow.types.llm import ChatCompletionResponse, ChatMessage, ChatParams, ChatChoice
from mlflow.pyfunc import ChatModel
from mlflow import deployments
from typing import List, Optional, Dict


class BasicAgent(ChatModel):
    def __init__(self):
        """Initialize the BasicAgent with placeholder values."""
        self.deploy_client = None
        self.models = {}
        self.models_config = {}
        self.conversation_history = []

    def load_context(self, context):
        """Initialize the connectors and model configurations."""
        self.deploy_client = deployments.get_deploy_client("databricks")
        self.models = context.model_config.get("models", {})
        self.models_config = context.model_config.get("configuration", {})

    def _get_system_message(self, role: str) -> Dict:
        """
        Get the system message configuration for the specified role.

        Args:
            role (str): The role of the agent (e.g., "oracle" or "judge").

        Returns:
            dict: The system message for the given role.
        """
        if role not in self.models:
            raise ValueError(f"Unknown role: {role}")

        instruction = self.models[role]["instruction"]
        return ChatMessage(role="system", content=instruction).to_dict()

    @mlflow.trace(name="Raw Agent Response")
    def _get_agent_response(
        self, message_list: List[Dict], endpoint: str, params: Optional[dict] = None
    ) -> Dict:
        """
        Call the agent endpoint to get a response.

        Args:
            message_list (List[Dict]): List of messages for the agent.
            endpoint (str): The agent's endpoint.
            params (Optional[dict]): Additional parameters for the call.

        Returns:
            dict: The response from the agent.
        """
        response = self.deploy_client.predict(
            endpoint=endpoint, inputs={"messages": message_list, **(params or {})}
        )
        return response["choices"][0]["message"]

    @mlflow.trace(name="Agent Call")
    def _call_agent(
        self, message: ChatMessage, role: str, params: Optional[dict] = None
    ) -> Dict:
        """
        Prepares and sends the request to a specific agent based on the role.

        Args:
            message (ChatMessage): The message to be processed.
            role (str): The role of the agent (e.g., "oracle" or "judge").
            params (Optional[dict]): Additional parameters for the call.

        Returns:
            dict: The response from the agent.
        """
        system_message = self._get_system_message(role)
        message_list = self._prepare_message_list(system_message, message)

        # Fetch agent response
        agent_config = self.models[role]
        response = self._get_agent_response(
            message_list, agent_config["endpoint"], params
        )

        # Update conversation history
        self.conversation_history.extend([message.to_dict(), response])
        return response

    @mlflow.trace(name="Assemble Conversation")
    def _prepare_message_list(
        self, system_message: Dict, user_message: ChatMessage
    ) -> List[Dict]:
        """
        Prepare the list of messages to send to the agent.

        Args:
            system_message (dict): The system message dictionary.
            user_message (ChatMessage): The user message.

        Returns:
            List[dict]: The complete list of messages to send.
        """
        user_prompt = {
            "role": "user",
            "content": self.models_config.get(
                "user_response_instruction", "Can you make the answer better?"
            ),
        }
        if self.conversation_history:
            return [system_message, *self.conversation_history, user_prompt]
        else:
            return [system_message, user_message.to_dict()]

    def predict(
        self, context, messages: List[ChatMessage], params: Optional[ChatParams] = None
    ) -> ChatCompletionResponse:
        """
        Predict method to handle agent conversation.

        Args:
            context: The MLflow context.
            messages (List[ChatMessage]): List of messages to process.
            params (Optional[ChatParams]): Additional parameters for the conversation.

        Returns:
            ChatCompletionResponse: The structured response object.
        """
        # Use the fluent API context handler to have added control over what is included in the span
        with mlflow.start_span(name="Audit Agent") as root_span:
            # Add the user input to the root span
            root_span.set_inputs(messages)

            # Add attributes to the root span
            attributes = {**params.to_dict(), **self.models_config, **self.models}
            root_span.set_attributes(attributes)

            # Initiate the conversation with the oracle
            oracle_params = self._get_model_params("oracle")
            oracle_response = self._call_agent(messages[0], "oracle", oracle_params)

            # Process the response with the judge
            judge_params = self._get_model_params("judge")
            judge_response = self._call_agent(
                ChatMessage(**oracle_response), "judge", judge_params
            )

            # Reset the conversation history and return the final response
            self.conversation_history = []

            output = ChatCompletionResponse(
                choices=[ChatChoice(index=0, message=ChatMessage(**judge_response))],
                usage={},
                model=judge_params.get("endpoint", "unknown"),
            )

            root_span.set_outputs(output)

        return output

    def _get_model_params(self, role: str) -> dict:
        """
        Retrieves model parameters for a given role.

        Args:
            role (str): The role of the agent (e.g., "oracle" or "judge").

        Returns:
            dict: A dictionary of parameters for the agent.
        """
        role_config = self.models.get(role, {})

        return {
            "temperature": role_config.get("temperature", 0.5),
            "max_tokens": role_config.get("max_tokens", 500),
        }


# IMPORTANT: specifies the Python ChatModel instance to use for inference requests when
# the model is loaded back
agent = BasicAgent()
mlflow.models.set_model(agent)

上面的代码片段将我们的代理定义为 ChatModel 的子类。使用来自代码的模型方法，我们调用

mlflow.models.set_model，传入我们 BasicAgent 的实例，以指示在代理重新加载时使用哪个模型对象进行推理。

将代理代码保存在 Python 文件中，例如 basic_agent.py。这是来自代码的模型的一个关键部分——它允许记录包含模型代码的文件，而不是序列化的模型对象，从而绕过序列化问题。

用我们的模型在代码文件中定义，在我们可以记录它之前只剩最后一步：我们需要定义配置以初始化我们的模型。这是通过定义我们的 model_config 配置来完成的。

设置我们的 model_config 值​

在记录模型之前，我们需要定义控制我们模型代理行为的配置。这种配置与模型核心逻辑的分离使我们能够轻松测试和比较不同的代理行为，而无需修改模型实现。通过使用灵活的配置系统，我们可以有效地试验不同的设置，从而大大简化迭代和微调我们的模型。

为何分离配置？​

在生成式人工智能 (GenAI) 的背景下，代理的行为会因给每个代理的指令集和参数（例如 temperature 或 max_tokens）而大不相同。如果我们直接将这些配置硬编码到模型逻辑中，每个新测试都需要更改模型源代码，从而导致：

  • 效率低下：为每次测试更改源代码会减慢实验过程。
  • 错误风险增加：不断修改源代码会增加引入错误或意外副作用的可能性。
  • 缺乏可重现性：如果没有代码和配置之间的清晰分离，跟踪和重现特定结果所使用的确切配置将变得具有挑战性。

通过在外部通过 model_config 参数设置这些值，我们使模型能够灵活适应不同的测试场景。这种方法还与 MLflow 的评估工具（如 mlflow.genai.evaluate()）无缝集成，该工具允许您系统地比较不同配置下的模型输出。

定义模型配置​

配置包含两个主要部分：

  1. Models：此部分定义了特定于代理的配置，例如本示例中的 judge 和 oracle 角色。每个代理都有：

    • 一个 endpoint：指定此代理使用的模型类型或服务。
    • 一个 instruction：定义代理的角色和职责（例如，回答问题、评估响应）。
    • Temperature 和 Max Tokens：控制生成的可变性（temperature）和响应的 token 限制。

  2. General Configuration：模型整体行为的其他设置，例如如何将用户响应框定给后续代理。

注意

有两种可用的选项可用于设置模型配置：直接在日志记录代码中（如下所示），或将配置写入本地位置的 yaml 格式的配置文件，其路径可以在日志记录期间定义 model_config 参数时指定。要了解有关 model_config 参数如何使用的更多信息，请参阅关于 model_config 用法的指南。

这是我们设置代理配置的方式：

model_config = {
    "models": {
        "judge": {
            "endpoint": "databricks-meta-llama-3-1-405b-instruct",
            "instruction": (
                "You are an evaluator of answers provided by others. Based on the context of both the question and the answer, "
                "provide a corrected answer if it is incorrect; otherwise, enhance the answer with additional context and explanation."
            ),
            "temperature": 0.5,
            "max_tokens": 2000,
        },
        "oracle": {
            "endpoint": "databricks-mixtral-8x7b-instruct",
            "instruction": (
                "You are a knowledgeable source of information that excels at providing detailed, but brief answers to questions. "
                "Provide an answer to the question based on the information provided."
            ),
            "temperature": 0.9,
            "max_tokens": 5000,
        },
    },
    "configuration": {
        "user_response_instruction": "Can you evaluate and enhance this answer with the provided contextual history?"
    },
}

外部配置的好处​

  • 灵活性：解耦的配置允许我们轻松切换或调整模型行为，而无需修改核心逻辑。例如，我们可以更改模型的指令或调整 temperature 以测试不同程度的创造性响应。
  • 可扩展性：随着系统中代理的增加或新角色的引入，我们可以扩展此配置，而不会使模型代码混乱。这种分离使代码库更整洁、更易于维护。
  • 可重现性和比较：通过将配置保持在外部，我们可以使用 MLflow 记录每次运行所使用的具体设置。这使得重现结果和比较不同实验更加容易，从而确保一个稳健的评估和裁定过程，以选择性能最佳的配置。

配置到位后，我们就可以记录模型并使用这些设置运行实验了。通过利用 MLflow 强大的跟踪和日志记录功能，我们将能够高效地管理实验，并从代理的响应中提取有价值的见解。

定义输入示例​

在记录模型之前，提供一个演示如何与模型交互的 input_example 很重要。此示例具有几个关键目的：

  • 日志记录时的验证：包含 input_example 允许 MLflow 在日志记录过程中使用此示例执行 predict 方法。这有助于验证您的模型是否可以处理预期的输入格式，并及早捕获任何问题。
  • UI 表示：input_example 显示在 MLflow UI 的模型工件下。这为用户提供了一个方便的参考，以了解在与已部署模型交互时预期的输入结构。

通过提供输入示例，您可以确保模型使用真实数据进行测试，从而增加模型按预期运行的信心。

提示

当使用 mlflow.pyfunc.ChatModel 定义 GenAI 应用程序时，如果未提供输入示例，将使用默认的占位符输入示例。如果您在 MLflow UI 的工件查看器中注意到不熟悉或通用的输入示例，那很可能是系统分配的默认占位符。为避免这种情况，请确保在保存模型时指定自定义输入示例。

这是我们将使用的输入示例：

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is a good recipe for baking scones that doesn't require a lot of skill?",
        }
    ]
}

此示例代表了一个用户请求简单的烤饼食谱。它与我们的 BasicAgent 模型预期的输入结构一致，该模型处理一个消息列表，其中每条消息都包含一个 role 和 content。

提供输入示例的好处

  • 执行和验证：MLflow 在日志记录期间会将此 input_example 传递给模型的 predict 方法，以确保它可以处理输入而不会出错。任何输入处理问题，如数据类型不正确或字段缺失，都将在此时捕获，从而节省您稍后的调试时间。
  • 用户界面显示：input_example 将在 MLflow UI 的模型工件视图部分中可见。这有助于用户理解模型期望的输入数据格式，从而更轻松地与部署后的模型进行交互。
  • 部署信心：通过提前使用示例输入验证模型，您可以获得额外的保证，即模型将在生产环境中正确运行，从而降低部署后意外行为的风险。

包含 input_example 是一个简单而强大的步骤，用于验证模型是否已准备好部署，并且在接收用户输入时会按预期运行。

记录和加载我们的自定义代理​

要使用 MLflow 记录和加载模型，请执行以下操作：

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="model",
        # If needed, update `python_model` to the Python file containing your agent code
        python_model="basic_agent.py",
        model_config=model_config,
        input_example=input_example,
    )

loaded = mlflow.pyfunc.load_model(model_info.model_uri)

response = loaded.predict(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is the best material to make a baseball bat out of?",
            }
        ]
    }
)

结论​

在本教程中，您已通过 MLflow 的 mlflow.pyfunc.ChatModel 类探索了创建自定义 GenAI 聊天代理的过程。我们演示了如何实现一种灵活、可扩展且标准化的方法来管理 GenAI 应用程序的部署，使您能够利用最新的 AI 进步，即使对于那些尚未通过命名 flavor 在 MLflow 中得到原生支持的库和框架也是如此。

通过使用 ChatModel 而不是更通用的 PythonModel，您可以利用在所有已部署 GenAI 接口中一致的不可变签名接口的优势，避免部署 GenAI 时遇到的许多常见陷阱，从而通过提供一致的体验来简化所有解决方案的使用。

本教程的关键要点包括：

  • 跟踪和监控：通过将跟踪直接集成到模型中，您可以深入了解应用程序的内部工作原理，使调试和优化更加简单。装饰器和流畅 API 方法都提供了管理关键操作跟踪的通用方式。
  • 灵活的配置管理：将配置与模型代码解耦可确保您能够快速测试和迭代，而无需修改源代码。这种方法不仅简化了实验，而且随着应用程序的演进，还增强了可重现性和可扩展性。
  • 标准化的输入和输出结构：利用 ChatModel 的静态签名简化了部署和提供 GenAI 模型的复杂性。通过遵循既定标准，您可以减少集成和验证输入/输出格式时通常遇到的摩擦。
  • 避免常见陷阱：在整个实现过程中，我们强调了最佳实践以避免常见问题，例如妥善处理秘密、验证输入示例以及理解加载上下文的细微差别。遵循这些实践可确保您的模型在生产环境中保持安全、健壮和可靠。
  • 验证和部署准备情况：在部署之前验证模型的重要性怎么强调都不为过。通过使用 mlflow.models.validate_serving_input() 等工具，您可以及早发现并解决潜在的部署问题，从而在生产部署过程中节省时间和精力。

随着生成式 AI 领域的不断发展，构建适应性强且标准化的模型将是利用未来几个月和几年中将解锁的令人兴奋且强大的功能的关键。本教程中涵盖的方法为您提供了在 MLflow 中集成和管理 GenAI 技术的强大框架，使您能够轻松开发、跟踪和部署复杂的 AI 解决方案。

我们鼓励您扩展和定制此基础示例以满足您的特定需求，并探索进一步的增强功能。通过利用 MLflow 不断增长的功能，您可以继续优化您的 GenAI 模型，确保它们在任何应用程序中都能提供有影响力的可靠结果。