通过

定义代理的输入和输出架构

  • 项目

MLflow 模型签名 定义 AI 代理的输入和输出架构要求。 模型签名告知内部和外部组件如何与代理交互。 模型签名是验证检查，以确保输入符合架构要求。

例如，若要使用代理评估评审应用，代理必须遵循 代理评估输入架构。

支持的输入架构

马赛克 AI 代理框架支持以下输入架构。

OpenAI 聊天完成架构

注意

Databricks 建议 OpenAI 聊天完成架构，因为它广泛使用并与许多代理框架和应用程序互操作。 如果 OpenAI 聊天完成架构不符合你的需求，则可以定义自己的架构。 请参阅 自定义代理架构。

  • （推荐） Databricks 建议使用 OpenAI 聊天完成架构。 OpenAI 聊天完成架构应具有一个对象数组作为 messages 参数。 此格式最适合 RAG 应用程序。

    question = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-Augmented Generation?",
        },
        {
            "role": "assistant",
            "content": "RAG, or Retrieval Augmented Generation, is a generative AI design pattern that combines a large language model (LLM) with external knowledge retrieval. This approach allows for real-time data connection to generative AI applications, improving their accuracy and quality by providing context from your data to the LLM during inference. Databricks offers integrated tools that support various RAG scenarios, such as unstructured data, structured data, tools & function calling, and agents.",
        },
        {
            "role": "user",
            "content": "How to build RAG for unstructured data",
        },
    ]
}

SplitChatMessageRequest

SplitChatMessagesRequest 建议用于多轮次聊天应用程序，尤其是在想要单独管理当前查询和历史记录时。

  question = {
      "query": "What is MLflow",
      "history": [
          {
              "role": "user",
              "content": "What is Retrieval-augmented Generation?"
          },
          {
              "role": "assistant",
              "content": "RAG is"
          }
      ]
  }

Langchain 表达式语言

如果代理使用 LangChain，则可以使用 LangChain 表达式语言编写链。 在链定义代码中，可以使用 itemgetter/> 获取消息或queryhistory对象，具体取决于输入格式。

支持的输出架构

马赛克 AI 代理框架支持以下输出架构。

ChatCompletionResponse

（建议） 建议使用 OpenAI 响应格式互操作性的客户使用 ChatCompletionResponse 。

LangChain - ChatCompletionsOutputParser

如果代理使用 LangChain，请使用 ChatCompletionsOutputParser() MLflow 作为最终链步骤。 这会将 LangChain AI 消息的格式设置为与代理兼容的格式。


  from mlflow.langchain.output_parsers import ChatCompletionsOutputParser

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(DatabricksChat)
      | ChatCompletionsOutputParser()
  )

PyFunc - 批注输入和输出类

如果你使用 PyFunc，Databricks 建议使用类型提示来批注 predict() 函数，其中输入和输出数据类是 mlflow.models.rag_signatures 中定义的类的子类。

可以从数据类内部 predict()构造输出对象。 返回的对象必须转换成字典表示形式，以确保可以序列化。


  from mlflow.models.rag_signatures import ChatCompletionRequest, ChatCompletionResponse, ChainCompletionChoice, Message

  class RAGModel(PythonModel):
    ...
      def predict(self, context, model_input: ChatCompletionRequest) -> ChatCompletionResponse:
        ...
        return asdict(ChatCompletionResponse(
            choices=[ChainCompletionChoice(message=Message(content=text))]
        ))

显式签名和推断签名

MLflow 可以在运行时推断代理的输入和输出架构，并自动创建签名。 如果使用支持的输入和输出架构，则推断的签名与代理框架兼容。 有关支持的架构的信息，请参阅 支持的输入架构。

但是，如果使用自定义代理架构，则必须根据自定义代理架构中的说明显式定义模型签名。

自定义代理架构

可以通过创建支持的输入/输出架构的子类来自定义代理的架构，以便传入和从代理传递和返回其他字段。 然后，添加额外的键 custom_inputs 并 custom_outputs 包含其他字段。 请参阅 Pyfunc 和 Langchain 的代码示例，以及使用自定义输入的基于 UI 的方法。

若要使用 databricks-agents SDK、Databricks 客户端 UI（如 AI Playground 和 Review App）和其他马赛克 AI 代理框架功能，代理的架构必须满足以下要求：

  • 代理必须使用 mlflow 版本 2.17.1 或更高版本。
  • 在代理笔记本中，将子类 Optional 中添加的其他字段标记为并分配默认值。
  • 在驱动程序笔记本中，使用子类的实例构造一个 ModelSignature using infer_signature 。
  • 在驱动程序笔记本中，通过调用 asdict 子类来构造输入示例。

PyFunc 自定义架构

除了上述要求外，基于 PyFunc 的代理还必须满足以下要求才能与马赛克 AI 代理功能进行交互。

PyFunc 自定义架构要求

在代理笔记本中，预测和预测流函数必须满足以下要求：

  • 为输入子类提供类型提示。
  • 使用点表示法访问数据类字段（例如，使用 model_input.custom_input.id 而不是 model_input["custom_inputs"]使用）。
  • 返回 dictionary。 可以调用 asdict 子类的实例，以将返回的格式设置为字典。

以下笔记本显示了使用 PyFunc 的自定义架构示例。

PyFunc 自定义架构代理笔记本

获取笔记本

PyFunc 自定义架构驱动程序笔记本

获取笔记本

Langchain 自定义架构

以下笔记本显示了使用 LangChain 的自定义架构示例。 可以修改笔记本中的wrap_output函数，以便从消息流中分析和提取信息。

Langchain 自定义架构代理笔记本

获取笔记本

Langchain 自定义架构驱动程序笔记本

获取笔记本

在 AI Playground 和代理评审应用中提供custom_inputs

如果使用字段定义具有其他输入 custom_inputs 的自定义代理架构，则可以在 AI Playground 和 代理评审应用中手动提供这些输入。 如果未提供自定义输入，代理将使用架构中指定的默认值。

  1. 在 AI 操场或代理评审应用中，选择齿轮图标 齿轮图标。

  2. 启用 custom_inputs。

  3. 提供与代理定义的输入架构匹配的 JSON 对象。

    在 AI 操场中提供custom_inputs。

JSON 对象必须与代理的输入架构匹配。 例如，如果有一个 custom_inputs 定义如下的数据类：

@dataclass
class CustomInputs():
  id: int = 0
  user: str = "default"

然后，在custom_inputs字段中输入的 JSON 字符串必须提供其值iduser，如以下示例所示：

{
  "id": 123
  "user": "dev_test",
}