旧式输入和输出代理架构

AI 代理必须遵循特定的输入和输出架构要求，才能与 Azure Databricks 上的其他功能兼容。 本文介绍如何使用旧代理创作签名和接口： ChatModel 接口、 SplitChatMessageRequest 输入架构和 StringResponse 输出架构。

创作旧版 ChatModel 代理

ChatModel 是 MLflow 中的旧代理创作接口，它扩展了 OpenAI 的 ChatCompletion 架构，允许你在添加自定义功能时保持与支持 ChatCompletion 标准的平台的兼容性。 有关更多详细信息 ，请参阅 MLflow：ChatModel 入门 。

将代理编写为 mlflow.pyfunc.ChatModel 的子类具有以下优势：

• 在调用已部署代理时，启用流式传输代理输出（在请求正文中绕过 {stream: true}）。
• 在代理被服务时，会自动启用 AI 网关推理表，从而提供访问增强的请求日志元数据（例如请求者名称）的权限。
• 允许使用类型化的 Python 类编写与 ChatCompletion 架构兼容的代理代码。
• MLflow 在记录代理时自动推断聊天完成兼容的签名，即使没有 input_example。 这简化了注册和部署代理的过程。 请参阅在日志记录期间推断模型签名。

以下代码最好在 Databricks 笔记本中执行。 笔记本提供了一个方便的环境，用于开发、测试和迭代代理。

MyAgent 类扩展 mlflow.pyfunc.ChatModel，实现所需的 predict 方法。 这可确保与马赛克 AI 代理框架兼容。

该类还包括用于处理流输出的可选方法 _create_chat_completion_chunk 和 predict_stream。

# Install the latest version of mlflow
%pip install -U mlflow
dbutils.library.restartPython()

import re
from typing import Optional, Dict, List, Generator
from mlflow.pyfunc import ChatModel
from mlflow.types.llm import (
  # Non-streaming helper classes
  ChatCompletionRequest,
  ChatCompletionResponse,
  ChatCompletionChunk,
  ChatMessage,
  ChatChoice,
  ChatParams,
  # Helper classes for streaming agent output
  ChatChoiceDelta,
  ChatChunkChoice,
)

class MyAgent(ChatModel):
  """
  Defines a custom agent that processes ChatCompletionRequests
  and returns ChatCompletionResponses.
  """
  def predict(self, context, messages: list[ChatMessage], params: ChatParams) -> ChatCompletionResponse:
    last_user_question_text = messages[-1].content
    response_message = ChatMessage(
      role="assistant",
      content=(
        f"I will always echo back your last question. Your last question was: {last_user_question_text}. "
      )
    )
    return ChatCompletionResponse(
      choices=[ChatChoice(message=response_message)]
    )

  def _create_chat_completion_chunk(self, content) -> ChatCompletionChunk:
    """Helper for constructing a ChatCompletionChunk instance for wrapping streaming agent output"""
    return ChatCompletionChunk(
      choices=[ChatChunkChoice(
        delta=ChatChoiceDelta(
          role="assistant",
          content=content
        )
      )]
    )

  def predict_stream(
    self, context, messages: List[ChatMessage], params: ChatParams
  ) -> Generator[ChatCompletionChunk, None, None]:
    last_user_question_text = messages[-1].content
    yield self._create_chat_completion_chunk(f"Echoing back your last question, word by word.")
    for word in re.findall(r"\S+\s*", last_user_question_text):
      yield self._create_chat_completion_chunk(word)

agent = MyAgent()
model_input = ChatCompletionRequest(
  messages=[ChatMessage(role="user", content="What is Databricks?")]
)
response = agent.predict(context=None, messages=model_input.messages, params=None)
print(response)

在一个笔记本中定义代理类 MyAgent 时，应创建单独的驱动程序笔记本。 驱动程序笔记本将代理记录到模型注册表，并使用 Model Service 部署代理。

此分离遵循 Databricks 推荐的工作流，使用 MLflow 的 "Models from Code" 方法来记录模型。

SplitChatMessageRequest 输入架构 （已弃用）

SplitChatMessagesRequest 允许将当前查询和历史记录单独作为代理输入传递。

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

StringResponse 输出架构（已弃用）

StringResponse 允许以具有单个字符串 content 字段的对象的形式返回代理的响应：

{"content": "This is an example string response"}