MLflow AI 网关

警告

MLflow AI 网关不支持 Windows 系统。

MLflow AI 网关是一个强大的工具，旨在简化组织内部各种大型语言模型（LLM）提供商（如 OpenAI 和 Anthropic）的使用和管理。它提供了一个高级接口，通过提供统一的端点来处理特定的 LLM 相关请求，从而简化了与这些服务的交互。

使用 MLflow AI 网关的一个主要优势是其对 API 密钥的集中管理。通过将这些密钥存储在一个安全的集中位置，组织可以显著增强其安全态势，最大限度地减少敏感 API 密钥在整个系统中的暴露。这也有助于防止在代码中暴露这些密钥，或要求最终用户安全地管理密钥。

网关服务器设计灵活且适应性强，能够通过更新配置文件轻松定义和管理端点。这使得系统可以轻松集成新的 LLM 提供商或提供商的 LLM 类型，而无需更改与网关服务器交互的应用程序。这种适应性使得 MLflow AI 网关服务在需要敏捷和快速响应变化的环境中成为一个不可或缺的工具。

这种简化和集中化的语言模型交互，再加上 API 密钥管理的额外安全层，使得 MLflow AI 网关成为那些经常使用 LLM 的组织的理想选择。

教程和指南

如果您有兴趣立即开始分步指南，以尽快使用 MLflow AI 网关，以下指南将是您的最佳首站。

查看网关服务器入门指南

快速入门

以下指南将帮助您快速启动并运行，使用包含 3 个端点的配置连接到 OpenAI 服务的聊天、补全和嵌入功能。

步骤 1：安装 MLflow AI 网关

首先，您需要在您的机器上安装 MLflow AI 网关。您可以从 PyPI 或 MLflow 仓库使用 pip 进行安装。

从 PyPI 安装

pip install 'mlflow[genai]'

步骤 2：为每个提供商设置 OpenAI API 密钥

网关服务器需要与 OpenAI API 通信。为此，它需要一个 API 密钥。您可以从 OpenAI 面板创建 API 密钥。

在此示例中，我们仅连接到 OpenAI。如果配置中包含其他提供商，也需要设置这些提供商的密钥。

获取密钥后，您可以在终端中将其设置为环境变量

export OPENAI_API_KEY=your_api_key_here

这将设置一个临时的、基于会话的环境变量。对于生产用例，建议将此密钥存储在 .bashrc 或 .zshrc 文件中，以便在系统重启时无需重新输入密钥。

步骤 3：创建网关服务器配置文件

接下来，您需要创建一个网关服务器配置文件。这是一个 YAML 文件，您可以在其中指定 MLflow AI 网关应该暴露的端点。让我们创建一个包含三个使用 OpenAI 作为提供商的端点文件：补全、聊天和嵌入。

有关配置文件的参数详细信息（包括 OpenAI 以外的其他提供商的参数），请参见下面的AI 网关服务器配置详情部分。

endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        openai_api_key: $OPENAI_API_KEY
    limit:
      renewal_period: minute
      calls: 10

  - name: chat
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        openai_api_key: $OPENAI_API_KEY

  - name: embeddings
    endpoint_type: llm/v1/embeddings
    model:
      provider: openai
      name: text-embedding-ada-002
      config:
        openai_api_key: $OPENAI_API_KEY

将此文件保存到将运行 MLflow AI 网关的系统上的某个位置。

步骤 4：启动网关服务器

您现在已经准备好启动网关服务器了！

使用 MLflow AI 网关的 start-server 命令并指定配置文件的路径

mlflow gateway start --config-path config.yaml --port {port} --host {host} --workers {worker count}

配置文件的路径也可以使用 MLFLOW_DEPLOYMENTS_CONFIG 环境变量来设置

export MLFLOW_DEPLOYMENTS_CONFIG=/path/to/config.yaml

如果不指定 host，将使用 localhost 地址。

如果不指定 port，将使用端口 5000。

gunicorn 的 worker 数量默认为 2 个 worker。

步骤 5：访问交互式 API 文档

MLflow AI 网关提供了一个交互式 API 文档端点，您可以使用它来探索和测试暴露的端点。在浏览器中导航到 http://{host}:{port}/（或 http://{host}:{port}/docs）即可访问。

docs 端点允许直接与端点交互，并通过点击端点定义条目中的“try it now”选项来向提供商服务提交实际请求。

步骤 6：使用客户端 API 发送请求

有关更多信息，请参见客户端 API 部分。

步骤 7：通过 REST API 向端点发送请求

您现在可以向暴露的端点发送请求了。有关请求格式的指导，请参见REST 示例。

步骤 8：比较提供商模型

这是一个从提供商添加新模型以确定哪个模型实例更适合给定用例的示例。

首先，使用附加的端点定义更新MLflow AI 网关配置文件 YAML 文件以进行测试

endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        openai_api_key: $OPENAI_API_KEY
  - name: completions-gpt4
    endpoint_type: llm/v1/completions
    model:
      provider: openai
      name: gpt-4
      config:
        openai_api_key: $OPENAI_API_KEY

此更新的配置添加了一个新的补全端点 completions-gpt4，同时保留了使用 gpt-4o-mini 模型配置的原始 completions 端点。

配置文件更新后，只需保存更改即可。网关服务器将自动创建新端点，无需停机。

如果您不再需要某个端点，可以从配置 YAML 文件中删除它并保存更改。网关服务器将自动删除该端点。

步骤 9：使用网关服务器端点进行模型开发

现在您已经创建了几个网关服务器端点，您可以创建查询这些端点的 MLflow 模型，使用 prompt engineering 等技术构建特定于应用程序的逻辑。有关更多信息，请参见网关服务器与 MLflow 模型。

概念

MLflow AI 网关 API、配置定义、示例和文档中引用了几个概念。熟悉这些术语将有助于简化新端点的配置和 MLflow AI 网关 API 的使用。

提供商

MLflow AI 网关设计用于支持各种模型提供商。提供商代表机器学习模型的来源，例如 OpenAI、Anthropic 等。每个提供商都有其特定的特征和配置，这些特征和配置被封装在 MLflow AI 网关中端点的模型部分。

支持的提供商

下表列出了 MLflow AI 网关中每个 LLM 提供商支持的相应端点类型。请注意，✅ 标记并不意味着提供商的所有模型都与端点类型兼容。例如，OpenAI 提供商支持所有三种端点类型，但模型 gpt-4 仅与 llm/v1/chat 端点类型兼容。

提供商	llm/v1/completions	llm/v1/chat	llm/v1/embeddings
OpenAI §	✅	✅	✅
Azure OpenAI	✅	✅	✅
MosaicML	✅	✅	✅
Anthropic	✅	✅	❌
Cohere	✅	✅	✅
PaLM	✅	✅	✅
MLflow	✅*	✅*	✅**
HuggingFace TGI	❌	✅	❌
AI21 Labs	✅	❌	❌
Amazon Bedrock	✅	✅	❌
Mistral	✅	✅	✅
TogetherAI	✅	✅	✅

§ 有关 OpenAI 的完整兼容性参考，请参阅OpenAI 模型兼容性矩阵。

在配置文件的每个模型块中，provider 字段用于指定该模型的提供商名称。这是一个字符串值，需要与 MLflow AI 网关支持的提供商相对应。

注意

* MLflow Model Serving 仅当输出返回是端点兼容格式时才适用于聊天或补全。响应必须符合 {"predictions": str} 或 {"predictions": {"candidates": str}} 的输出格式。任何不符合这些结构的复杂模型返回类型将在查询时引发异常。

** 嵌入支持仅适用于响应签名符合 {"predictions": List[float]} 或 {"predictions": List[List[float]]} 结构化格式的模型。任何其他返回类型将在查询时引发异常。transformers 中的 FeatureExtractionPipeline 和使用 sentence_transformers 风格的模型将为嵌入端点返回正确的数据结构。

以下是一个端点内的提供商配置示例

endpoints:
  - name: chat
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4
      config:
        openai_api_key: $OPENAI_API_KEY
    limit:
      renewal_period: minute
      calls: 10

在上述配置中，openai 是模型的提供商。

截至目前，MLflow AI 网关支持以下提供商

• mosaicml：用于MosaicML 提供的模型。
• openai：用于OpenAI 提供的模型以及 Azure OpenAI 和使用 AAD 的 Azure OpenAI 的Azure 集成。
• anthropic：用于Anthropic 提供的模型。
• cohere：用于Cohere 提供的模型。
• palm：用于PaLM 提供的模型。
• huggingface text generation inference：用于使用Huggingface Text Generation Inference 部署的模型。
• ai21labs：用于AI21 Labs 提供的模型。
• bedrock：用于Amazon Bedrock 提供的模型。
• mistral：用于Mistral 提供的模型。
• togetherai：用于TogetherAI 提供的模型。

更多提供商正在持续添加中。请查看 MLflow AI 网关文档的最新版本，获取支持提供商的最新列表。

如果您想使用上述提供商未提供的 LLM 模型，或者您想集成私有 LLM 模型，您可以创建一个提供商插件来与 MLflow AI 网关集成。

端点

端点是 MLflow AI 网关如何工作的核心。每个端点都充当用户的代理端点，将请求转发到配置文件中指定的底层模型和提供商。

MLflow AI 网关中的一个端点包含以下字段

• name：这是端点的唯一标识符。在使用 MLflow AI 网关进行 API 调用时，这将是 URL 的一部分。
• type：端点的类型对应于您期望的语言模型交互类型。例如，文本补全操作使用 llm/v1/completions，文本嵌入使用 llm/v1/embeddings，聊天操作使用 llm/v1/chat。
• model：定义此端点将请求转发到的模型。模型包含以下详细信息
  • provider：指定此模型的提供商名称。例如，OpenAI 的 GPT-4o 模型使用 openai。
  • name：要使用的模型的名称。例如，OpenAI 的 GPT-4o-Mini 模型使用 gpt-4o-mini。
  • config：包含模型所需的任何附加配置详细信息。这包括指定 API 基础 URL 和 API 密钥。
• limit：指定此端点将遵循的速率限制设置。limit 字段包含以下字段
  • renewal_period：速率限制的时间单位，取值范围为 [second|minute|hour|day|month|year]。
  • calls：此端点在指定时间单位内接受的调用次数。

以下是一个端点配置示例

endpoints:
  - name: completions
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        openai_api_key: $OPENAI_API_KEY
    limit:
      renewal_period: minute
      calls: 10

在上面的示例中，发送到 completions 端点的请求将被转发到由 openai 提供的 gpt-4o-mini 模型。

配置文件的端点可以随时更新，MLflow AI 网关将自动更新可用端点，无需重启。此功能使您可以灵活地根据需求添加、删除或修改端点。它支持端点的“热插拔”，为与 MLflow AI 网关交互的任何应用程序或服务提供无缝体验。

在配置文件中定义端点时，请确保每个名称都是唯一的，以防止冲突。重复的端点名称将引发 MlflowException。

模型

endpoint 部分内的 model 部分指定用于生成响应的模型。此配置块需要包含一个 name 字段，用于指定要使用的确切模型实例。此外，需要指定一个提供商，并且您需要拥有其经过身份验证的 API 密钥。

不同的端点类型通常与特定模型关联。例如，llm/v1/chat 和 llm/v1/completions 端点通常与对话模型关联，而 llm/v1/embeddings 端点通常与嵌入或 Transformer 模型关联。您选择的模型应适合指定的端点类型。

以下是一个端点内模型名称配置示例

endpoints:
  - name: embeddings
    endpoint_type: llm/v1/embeddings
    model:
      provider: openai
      name: text-embedding-ada-002
      config:
        openai_api_key: $OPENAI_API_KEY

在上述配置中，text-embedding-ada-002 是用于嵌入端点的模型。

指定模型时，提供商必须支持您请求的模型至关重要。例如，openai 作为提供商支持 text-embedding-ada-002 等模型，但其他提供商可能不支持。如果提供商不支持该模型，当尝试将请求路由到该模型时，MLflow AI 网关将返回 HTTP 4xx 错误。

重要

始终查看指定提供商的最新文档，以确保您要使用的模型受配置的端点类型支持。

请记住，您选择的模型直接影响您从 API 调用中获得的响应结果。因此，选择一个适合您的用例需求的模型。例如，对于生成对话响应，您通常会选择聊天模型。相反，对于生成文本嵌入，您会选择嵌入模型。

配置网关服务器

MLflow AI 网关依赖于用户提供的 YAML 格式配置文件，该文件定义了可供服务器使用的端点和提供商。配置文件决定了网关服务器如何与各种语言模型提供商交互，并决定了用户可以访问的端点。

AI 网关服务器配置

配置文件包含一系列部分，每个部分代表一个唯一的端点。每个端点部分都有一个名称、一个类型和一个模型规范，其中包括模型提供商、名称和配置详细信息。配置部分通常包含 API 的基础 URL 和 API 密钥的环境变量。

以下是单端点配置示例

endpoints:
  - name: chat
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        openai_api_key: $OPENAI_API_KEY
    limit:
      renewal_period: minute
      calls: 10

在此示例中，我们定义了一个名为 chat 的端点，对应于 llm/v1/chat 类型，它将使用 OpenAI 的 gpt-4o-mini 模型返回 OpenAI 服务的查询响应，并接受每分钟最多 10 个请求。

MLflow AI 网关配置非常易于更新。只需编辑配置文件并保存更改，MLflow AI 网关将自动更新端点，无需任何中断或停机。这使您可以在保持应用程序稳定可靠的同时尝试新的提供商或模型类型。

为了定义给定提供商的 API 密钥，有三种主要选项

1. 直接将其包含在 YAML 配置文件中。
2. 使用环境变量存储 API 密钥并在 YAML 配置文件中引用它。
3. 将 API 密钥定义在一个文件中，并在 YAML 配置文件中引用该包含密钥文件的位置。

如果您选择直接包含 API 密钥，请将 YAML 文件中的 $OPENAI_API_KEY 替换为您的实际 API 密钥。

警告

MLflow AI 网关提供对收费外部 LLM 服务的直接访问。强烈建议限制对此服务器的访问。有关指导，请参见安全性部分。

如果您更喜欢使用环境变量（推荐），您可以在 shell 环境中定义它。例如

export OPENAI_API_KEY="your_openai_api_key"

注意： 将“your_openai_api_key”替换为您的实际 OpenAI API 密钥。

AI 网关服务器配置详情

MLflow AI 网关依赖于用户提供的配置文件。它定义了网关服务器如何与各种语言模型提供商交互，并决定了用户可以访问的端点。

配置文件采用 YAML 格式编写，包含一系列部分，每个部分代表一个唯一的端点。每个端点部分都有一个名称、一个类型和一个模型规范，其中包括提供商、模型名称和提供商特定的配置详细信息。

以下是每个配置参数的详细信息

通用配置参数

• endpoints：这是一个端点配置列表。每个端点代表一个唯一的端点，映射到特定的语言模型服务。

每个端点具有以下配置参数

• name：这是端点的名称。它需要是唯一的名称，不能包含空格或连字符和下划线以外的非字母数字字符。
• endpoint_type：这指定了此端点提供的服务类型。它决定了端点输入的接口和返回的输出。当前支持的端点类型有
  • "llm/v1/completions"
  • "llm/v1/chat"
  • "llm/v1/embeddings"
• model：这定义了语言模型的提供商特定详细信息。它包含以下字段
  • provider：这表示 AI 模型的提供商。它接受以下值
    • "openai"
    • "mosaicml"
    • "anthropic"
    • "cohere"
    • "palm"
    • "azure" / "azuread"
    • "mlflow-model-serving"
    • "huggingface-text-generation-inference"
    • "ai21labs"
    • "bedrock"
    • "mistral"
    • "togetherai"
  • name：这是一个可选字段，用于指定模型的名称。
  • config：这包含提供商特定的配置详细信息。

提供商特定的配置参数

OpenAI

配置参数	必需	默认值	描述
openai_api_key	是		这是 OpenAI 服务的 API 密钥。
openai_api_type	否		这是一个可选字段，用于指定要使用的 OpenAI API 类型。
openai_api_base	否	https://api.openai.com/v1	这是 OpenAI API 的基础 URL。
openai_api_version	否		这是一个可选字段，用于指定 OpenAI API 版本。
openai_organization	否		这是一个可选字段，用于指定 OpenAI 中的组织。

MosaicML

配置参数	必需	默认值	描述
mosaicml_api_key	是	不适用	这是 MosaicML 服务的 API 密钥。

Cohere

配置参数	必需	默认值	描述
cohere_api_key	是	不适用	这是 Cohere 服务的 API 密钥。

HuggingFace Text Generation Inference

配置参数	必需	默认值	描述
hf_server_url	是	不适用	这是 Huggingface TGI 服务器的 url。

PaLM

配置参数	必需	默认值	描述
palm_api_key	是	不适用	这是 PaLM 服务的 API 密钥。

AI21 Labs

配置参数	必需	默认值	描述
ai21labs_api_key	是	不适用	这是 AI21 Labs 服务的 API 密钥。

Anthropic

配置参数	必需	默认值	描述
anthropic_api_key	是	不适用	这是 Anthropic 服务的 API 密钥。

Amazon Bedrock

Amazon Bedrock 端点的顶级模型配置必须是以下两种支持的身份验证模式之一：基于密钥或基于角色。

配置参数	必需	默认值	描述
aws_config	否		一个包含基于密钥或基于角色模式的对象。

要使用基于密钥的身份验证，请定义一个 Amazon Bedrock 端点，并填写以下必需字段。

注意

如果配置的端点仅用于开发或测试，建议使用 IAM 用户角色或临时的短期标准 IAM 角色；而对于生产部署，建议使用标准的长期 IAM 角色，以确保端点能够处理长时间的身份验证。如果身份验证过期且需要提供新的密钥，则必须重新创建端点以持久化新密钥。

配置参数	必需	默认值	描述
aws_region	否	AWS_REGION/AWS_DEFAULT_REGION	用于访问 Bedrock 的 AWS 区域。
aws_secret_access_key	是		用于授权使用 Bedrock 的 IAM 用户/角色的 AWS 秘密访问密钥。
aws_access_key_id	是		用于授权使用 Bedrock 的 IAM 用户/角色的 AWS 访问密钥 ID。
aws_session_token	否	无	可选的会话令牌，如果需要。

或者，对于基于角色的身份验证，可以定义并初始化一个 Amazon Bedrock 端点，并提供一个授权访问 Bedrock 的 IAM 角色 ARN。MLflow AI 网关将尝试使用标准凭证提供商链假设此角色，并在角色凭证过期时续订。

配置参数	必需	默认值	描述
aws_region	否	AWS_REGION/AWS_DEFAULT_REGION	用于访问 Bedrock 的 AWS 区域。
aws_role_arn	是		一个授权使用 Bedrock 的 AWS 角色。标准凭证提供商链必须能够找到授权 assume 此角色的凭证。
session_length_seconds	否	900	请求的会话长度（秒）。

MLflow Model Serving

配置参数	必需	默认值	描述
model_server_url	是	不适用	这是 MLflow Model Server 的 url。

请注意，对于 MLflow model serving，model 定义的 name 参数不用于验证，仅用于参考目的。此别名对于理解用于引用已部署模型的特定版本或端点定义非常有用。您可以选择任何您想要的名称，前提是它可以进行 JSON 序列化。

Azure OpenAI

Azure 提供了两种不同的机制与 OpenAI 集成，每种机制对应一种不同的安全验证类型。一种依赖于访问令牌进行验证，称为 azure，另一种使用 Azure Active Directory (Azure AD) 集成进行身份验证，称为 azuread。

为了匹配您用户的交互和安全访问要求，请调整 openai_api_type 参数以表示首选的安全验证模型。这将确保您的 Azure-OpenAI 集成实现无缝交互和可靠安全。

配置参数	必需	默认值	描述
openai_api_key	是		这是 Azure OpenAI 服务的 API 密钥。
openai_api_type	是		此字段必须是 azure 或 azuread，具体取决于安全访问协议。
openai_api_base	是		这是 Azure 提供的 Azure OpenAI API 服务的基础 URL。
openai_api_version	是		要使用的 Azure OpenAI 服务版本，由日期指定。
openai_deployment_name	是		这是 Azure OpenAI 服务的部署资源的名称。
openai_organization	否		这是一个可选字段，用于指定 OpenAI 中的组织。

Mistral

配置参数	必需	默认值	描述
mistral_api_key	是	不适用	这是 Mistral 服务的 API 密钥。

TogetherAI

配置参数	必需	默认值	描述
togetherai_api_key	是	不适用	这是 TogetherAI 服务的 API 密钥。

Azure OpenAI 的配置示例如下

endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: openai
      name: gpt-35-turbo
      config:
        openai_api_type: "azuread"
        openai_api_key: $AZURE_AAD_TOKEN
        openai_deployment_name: "{your_deployment_name}"
        openai_api_base: "https://{your_resource_name}-azureopenai.openai.azure.com/"
        openai_api_version: "2023-05-15"
    limit:
      renewal_period: minute
      calls: 10

注意

与直接的 OpenAI 服务相比，Azure OpenAI 具有明显的特点。有关概述，请参阅比较文档。

对于指定 API 密钥，有三种选项

1. （首选）使用环境变量存储 API 密钥并在 YAML 配置文件中引用它。这由环境变量名称前的 $ 符号表示。
2. （首选）将 API 密钥定义在一个文件中，并在 YAML 配置文件中引用该包含密钥文件的位置。
3. 直接将其包含在 YAML 配置文件中。

重要

为了更好的安全实践，建议使用环境变量或文件密钥。如果 API 密钥直接包含在配置文件中，应确保文件安全存储并进行适当的访问控制。请确保配置文件存储在安全位置，因为它包含敏感的 API 密钥。

查询 AI 网关服务器

配置并启动 MLflow AI 网关后，它就可以接收来自用户的流量了。

标准查询参数

MLflow AI 网关为聊天、补全和嵌入定义了标准参数，这些参数可以在查询任何端点时使用，无论其提供商如何。每个参数都有一个标准范围和默认值。查询特定提供商的端点时，MLflow AI 网关会根据提供商在该参数的值范围自动调整参数值。

补全

类型为 llm/v1/completions 的补全端点的标准参数包括

查询参数	类型	必需	默认值	描述
prompt	string	是	不适用	用于生成补全的 prompt。
n	integer	否	1	为指定的 prompt 生成的补全数量，介于 1 和 5 之间。
temperature	float	否	0.0	要使用的采样温度，介于 0 和 1 之间。值越高，输出越随机；值越低，输出越确定。
max_tokens	integer	否	无	最大补全长度，介于 1 和无穷大（无限制）之间。
stop	array[string]	否	无	模型应停止生成 token 并返回补全的序列。

聊天

类型为 llm/v1/chat 的聊天端点的标准参数包括

查询参数	类型	必需	默认值	描述
messages	array[message]	是	不适用	对话中的消息列表，用于生成新消息（聊天补全）。有关消息结构的信息，请参见消息。
n	integer	否	1	为指定的 prompt 生成的聊天补全数量，介于 1 和 5 之间。
temperature	float	否	0.0	要使用的采样温度，介于 0 和 1 之间。值越高，输出越随机；值越低，输出越确定。
max_tokens	integer	否	无	最大补全长度，介于 1 和无穷大（无限制）之间。
stop	array[string]	否	无	模型应停止生成 token 并返回聊天补全的序列。

消息

每条聊天消息都是一个字符串字典，包含以下字段

字段名称	必需	默认值	描述
role	是	不适用	发送消息的对话参与者的角色。必须是以下之一："system"、"user" 或 "assistant"。
content	是	不适用	消息内容。

嵌入

类型为 llm/v1/embeddings 的补全端点的标准参数包括

查询参数	类型	必需	默认值	描述
input	string or array[string]	是	不适用	用于生成嵌入的字符串或字符串列表。

附加查询参数

除了标准查询参数之外，您可以在查询中传递端点提供商支持的任何附加参数。例如

• logit_bias（支持 OpenAI、Cohere）
• top_k（支持 MosaicML、Anthropic、PaLM、Cohere）
• frequency_penalty（支持 OpenAI、Cohere、AI21 Labs）
• presence_penalty（支持 OpenAI、Cohere、AI21 Labs）
• stream（支持 OpenAI、Cohere）

以下是使用附加参数向 MLflow AI 网关端点提交查询请求的示例

from mlflow.deployments import get_deploy_client

client = get_deploy_client("http://my.deployments:8888")

data = {
    "prompt": (
        "What would happen if an asteroid the size of "
        "a basketball encountered the Earth traveling at 0.5c? "
        "Please provide your answer in .rst format for the purposes of documentation."
    ),
    "temperature": 0.5,
    "max_tokens": 1000,
    "n": 1,
    "frequency_penalty": 0.2,
    "presence_penalty": 0.2,
}

client.predict(endpoint="completions-gpt4", inputs=data)

查询结果如下

{
    "id": "chatcmpl-8Pr33fsCAtD2L4oZHlyfOkiYHLapc",
    "object": "text_completion",
    "created": 1701172809,
    "model": "gpt-4-0613",
    "choices": [
        {
            "index": 0,
            "text": "If an asteroid the size of a basketball ...",
        }
    ],
    "usage": {
        "prompt_tokens": 43,
        "completion_tokens": 592,
        "total_tokens": 635,
    },
}

流式处理

一些提供商支持流式响应。当您希望在响应生成后立即接收它们，而不是等待整个响应生成后再接收时，流式响应非常有用。以下提供商支持流式响应

提供商	端点
	llm/v1/completions	llm/v1/completions
OpenAI	✓	✓
Cohere	✓	✓
Anthropic	✘	✓

要启用流式响应，请在请求中将 stream 参数设置为 true。例如

curl -X POST http://my.deployments:8888/endpoints/chat/invocations \
  -H "Content-Type: application/json" \
  -d '{"messages": [{"role": "user", "content": "hello"}], "stream": true}'

查询结果遵循OpenAI schema。

聊天

data: {"choices": [{"delta": {"content": null, "role": "assistant"}, "finish_reason": null, "index": 0}], "created": 1701161926, "id": "chatcmpl-8PoDWSiVE8MHNsUZF2awkW5gNGYs3", "model": "gpt-35-turbo", "object": "chat.completion.chunk"}

data: {"choices": [{"delta": {"content": "Hello", "role": null}, "finish_reason": null, "index": 0}], "created": 1701161926, "id": "chatcmpl-8PoDWSiVE8MHNsUZF2awkW5gNGYs3", "model": "gpt-35-turbo", "object": "chat.completion.chunk"}

data: {"choices": [{"delta": {"content": " there", "role": null}, "finish_reason": null, "index": 0}], "created": 1701161926, "id": "chatcmpl-8PoDWSiVE8MHNsUZF2awkW5gNGYs3", "model": "gpt-35-turbo", "object": "chat.completion.chunk"}

data: {"choices": [{"delta": {"content": null, "role": null}, "finish_reason": "stop", "index": 0}], "created": 1701161926, "id": "chatcmpl-8PoDWSiVE8MHNsUZF2awkW5gNGYs3", "model": "gpt-35-turbo", "object": "chat.completion.chunk"}

补全

data: {"choices": [{"delta": {"role": null, "content": null}, "finish_reason": null, "index": 0}], "created": 1701161629, "id": "chatcmpl-8Po8jVXzljc245k1Ah4UsAcm2zxQ2", "model": "gpt-35-turbo", "object": "text_completion_chunk"}

data: {"choices": [{"delta": {"role": null, "content": "If"}, "finish_reason": null, "index": 0}], "created": 1701161629, "id": "chatcmpl-8Po8jVXzljc245k1Ah4UsAcm2zxQ2", "model": "gpt-35-turbo", "object": "text_completion_chunk"}

data: {"choices": [{"delta": {"role": null, "content": " an"}, "finish_reason": null, "index": 0}], "created": 1701161629, "id": "chatcmpl-8Po8jVXzljc245k1Ah4UsAcm2zxQ2", "model": "gpt-35-turbo", "object": "text_completion_chunk"}

data: {"choices": [{"delta": {"role": null, "content": " asteroid"}, "finish_reason": null, "index": 0}], "created": 1701161629, "id": "chatcmpl-8Po8jVXzljc245k1Ah4UsAcm2zxQ2", "model": "gpt-35-turbo", "object": "text_completion_chunk"}

data: {"choices": [{"delta": {"role": null, "content": null}, "finish_reason": "length", "index": 0}], "created": 1701161629, "id": "chatcmpl-8Po8jVXzljc245k1Ah4UsAcm2zxQ2", "model": "gpt-35-turbo", "object": "text_completion_chunk"}

FastAPI 文档 ("/docs")

FastAPI 是用于构建 MLflow AI 网关的框架，它提供了一个自动交互式 API 文档界面，可通过 "/docs" 端点访问（例如，http://my.deployments:9000/docs）。此交互式界面对于探索和测试可用的 API 端点非常方便。

为方便起见，访问根 URL（例如，http://my.deployments:9000）会重定向到此 "/docs" 端点。

MLflow Python 客户端 API

MlflowDeploymentClient 是用于与 MLflow AI 网关交互的面向用户的客户端 API。它通过简单易用的 Python API 抽象了对网关服务器的 HTTP 请求。

客户端 API

要使用 MlflowDeploymentClient API，请参见以下示例以了解可用的 API 方法

1. 创建一个 MlflowDeploymentClient

from mlflow.deployments import get_deploy_client

client = get_deploy_client("http://my.deployments:8888")

2. 列出所有端点

   list_endpoints() 方法返回所有端点的列表。

   endpoint = client.list_endpoints()
   for endpoint in endpoints:
       print(endpoint)

3. 查询端点

   predict() 方法将查询提交到已配置的提供商端点。您在查询中发送的数据结构取决于端点。

   response = client.predict(
       endpoint="chat",
       inputs={"messages": [{"role": "user", "content": "Tell me a joke about rabbits"}]},
   )
   print(response)

LangChain 集成

LangChain 支持与 MLflow Deployments 的集成。此集成使用户能够在网关服务器中使用 prompt engineering、检索增强生成和其他技术处理 LLM。

示例

import mlflow
from langchain import LLMChain, PromptTemplate
from langchain.llms import Mlflow

llm = Mlflow(target_uri="http://127.0.0.1:5000", endpoint="completions")
llm_chain = LLMChain(
    llm=llm,
    prompt=PromptTemplate(
        input_variables=["adjective"],
        template="Tell me a {adjective} joke",
    ),
)
result = llm_chain.run(adjective="funny")
print(result)

with mlflow.start_run():
    model_info = mlflow.langchain.log_model(llm_chain, "model")

model = mlflow.pyfunc.load_model(model_info.model_uri)
print(model.predict([{"adjective": "funny"}]))

MLflow 模型

与 MLflow 模型交互有两种方式。通过使用自定义 PyFunc 模型，可以直接向网关服务器端点发出查询，并在模型内部更广泛的上下文中使用。数据可以进行增强、操作或用于专家混合范式。另一种利用 MLflow AI 网关和 MLflow 模型的方法是直接将已提供服务的 MLflow 模型定义为网关服务器内部的端点。

使用网关服务器查询已提供服务的 MLflow 模型

有关使用 MLflow serving 集成通过 MLflow AI 网关直接查询模型的完整演练和示例，请参阅完整示例。在该指南中，您将看到从不同服务器提供多个模型并配置 MLflow AI 网关实例以提供处理查询的单一统一端点的整个端到端过程。

使用 MLflow 模型查询网关服务器

您还可以构建和部署调用 MLflow AI 网关的 MLflow 模型。以下示例演示了如何在自定义 pyfunc 模型中使用网关服务器。

注意

以下示例中显示的自定义 Model 使用环境变量获取网关服务器的 uri。这些值也可以在定义中手动设置，或者在设置 uri 后通过mlflow.deployments.get_deployments_target() 应用。对于以下示例，MLFLOW_DEPLOYMENTS_TARGET 的值为 http://127.0.0.1:5000/。对于实际部署用例，此值将设置为已配置的生产部署服务器。

import os
import pandas as pd
import mlflow


def predict(data):
    from mlflow.deployments import get_deploy_client

    client = get_deploy_client(os.environ["MLFLOW_DEPLOYMENTS_TARGET"])

    payload = data.to_dict(orient="records")
    return [
        client.predict(endpoint="completions", inputs=query)["choices"][0]["text"]
        for query in payload
    ]


input_example = pd.DataFrame.from_dict(
    {"prompt": ["Where is the moon?", "What is a comet made of?"]}
)
signature = mlflow.models.infer_signature(
    input_example, ["Above our heads.", "It's mostly ice and rocks."]
)

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model=predict,
        registered_model_name="anthropic_completions",
        artifact_path="anthropic_completions",
        input_example=input_example,
        signature=signature,
    )

df = pd.DataFrame.from_dict(
    {
        "prompt": ["Tell me about Jupiter", "Tell me about Saturn"],
        "temperature": 0.6,
        "max_records": 500,
    }
)

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

print(loaded_model.predict(df))

此自定义 MLflow 模型可以像任何其他 MLflow 模型一样使用。它可以在 spark_udf 中使用，与mlflow.evaluate() 一起使用，或像任何其他模型一样部署。

REST API

REST API 允许您直接向 MLflow AI 网关发送 HTTP 请求。如果您不使用 Python 或者喜欢直接使用 HTTP 与网关服务器交互，这将很有用。

以下是一些如何使用 curl 与 MLflow AI 网关交互的示例

1. 获取特定端点的信息：GET /api/2.0/endpoints/{name}

   此路由返回端点数据结构的序列化表示。它提供了请求端点的名称、类型以及模型详细信息。

   curl -X GET http://my.deployments:8888/api/2.0/endpoints/embeddings

2. 列出所有端点：GET /api/2.0/endpoints/

   此路由返回所有端点的列表。

   curl -X GET http://my.deployments:8888/api/2.0/endpoints/

3. 查询端点：POST /endpoints/{name}/invocations

   此路由允许您向已配置的提供商端点提交查询。您在查询中发送的数据结构取决于端点。以下是“completions”、“chat”和“embeddings”端点的示例

   • 补全

     curl -X POST http://my.deployments:8888/endpoints/completions/invocations \
       -H "Content-Type: application/json" \
       -d '{"prompt": "Describe the probability distribution of the decay chain of U-235"}'

   • 聊天

     curl -X POST http://my.deployments:8888/endpoints/chat/invocations \
       -H "Content-Type: application/json" \
       -d '{"messages": [{"role": "user", "content": "Can you write a limerick about orange flavored popsicles?"}]}'

   • 嵌入

     curl -X POST http://my.deployments:8888/endpoints/embeddings/invocations \
       -H "Content-Type: application/json" \
       -d '{"input": ["I would like to return my shipment of beanie babies, please", "Can I please speak to a human now?"]}'

注意： 请记住将 my.deployments:8888 替换为您实际的 MLflow AI 网关的 URL。

插件式 LLM 提供商（实验性）

注意

此功能正在积极开发中，被标记为实验性功能。它在将来的版本中可能会发生变化，恕不另行通知。

MLflow AI 网关支持通过使用插件来使用自定义语言模型提供商。插件是一个 Python 包，提供语言模型提供商的自定义实现。这允许用户将其自己的语言模型服务与 MLflow AI 网关集成。

要创建自定义插件，您需要实现一个继承自 mlflow.gateway.providers.BaseProvider 的提供商类，以及一个继承自 mlflow.gateway.base_models.ConfigModel 的配置类。

示例

import os
from typing import AsyncIterable

from mlflow.utils.pydantic_utils import field_validator
from mlflow.gateway.base_models import ConfigModel
from mlflow.gateway.config import RouteConfig
from mlflow.gateway.providers import BaseProvider
from mlflow.gateway.schemas import chat, completions, embeddings


class MyLLMConfig(ConfigModel):
    # This model defines the configuration for the provider such as API keys
    my_llm_api_key: str

    @field_validator("my_llm_api_key", mode="before")
    def validate_my_llm_api_key(cls, value):
        return os.environ[value.lstrip("$")]


class MyLLMProvider(BaseProvider):
    # Define the provider name. This will be displayed in log and error messages.
    NAME = "my_llm"
    # Define the config model for the provider.
    # This must be a subclass of ConfigModel.
    CONFIG_TYPE = MyLLMConfig

    def __init__(self, config: RouteConfig) -> None:
        super().__init__(config)
        if config.model.config is None or not isinstance(
            config.model.config, MyLLMConfig
        ):
            raise TypeError(f"Unexpected config type {config.model.config}")
        self.my_llm_config: MyLLMConfig = config.model.config

    # You can implement one or more of the following methods
    # depending on the capabilities of your provider.
    # Implementing `completions`, `chat` and `embeddings` will enable the respective endpoints.
    # Implementing `completions_stream` and `chat_stream` will enable the `stream=True`
    # option for the respective endpoints.
    # Unimplemented methods will return a 501 Not Implemented HTTP response upon invocation.
    async def completions_stream(
        self, payload: completions.RequestPayload
    ) -> AsyncIterable[completions.StreamResponsePayload]:
        ...

    async def completions(
        self, payload: completions.RequestPayload
    ) -> completions.ResponsePayload:
        ...

    async def chat_stream(
        self, payload: chat.RequestPayload
    ) -> AsyncIterable[chat.StreamResponsePayload]:
        ...

    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
        ...

    async def embeddings(
        self, payload: embeddings.RequestPayload
    ) -> embeddings.ResponsePayload:
        ...

然后，您需要创建一个包含插件实现的 Python 包。您必须在 mlflow.gateway.providers 组下指定一个入口点，以便 MLflow 可以检测到您的插件。入口点应采用 <name> = <module>:<class> 的格式。

pyproject.toml

[project]
name = "my_llm"
version = "1.0"

[project.entry-points."mlflow.gateway.providers"]
my_llm = "my_llm.providers:MyLLMProvider"

[tool.setuptools.packages.find]
include = ["my_llm*"]
namespaces = false

如果您的包中有多个提供商，可以在同一个包中指定多个入口点。请注意，入口点名称必须全局唯一。如果两个插件指定了相同的入口点名称，MLflow 将在启动时引发错误。

MLflow 默认提供了许多提供商。您的插件名称不能与其中任何一个相同。有关默认提供商的完整列表，请参见AI 网关服务器配置详情。

最后，您需要在与 MLflow AI 网关相同的环境中安装插件包。

重要

仅安装您信任来源的插件包。使用插件提供商启动服务器将执行插件包中定义的任何任意代码。

然后，您可以根据入口点名称在 MLflow AI 网关配置文件中指定插件提供商。

endpoints:
  - name: chat
    endpoint_type: llm/v1/chat
    model:
      provider: my_llm
      name: my-model-0.1.2
      config:
        my_llm_api_key: $MY_LLM_API_KEY

示例

一个工作示例可以在 MLflow 仓库的examples/gateway/plugin 中找到。

MLflow AI 网关 API 文档

API 文档

OpenAI 兼容性

MLflow AI 网关与 OpenAI API 兼容，并支持 chat、completions 和 embeddings API。可以使用 OpenAI 客户端查询服务器，如下例所示

1. 创建配置文件

endpoints:
  - name: my-chat
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        openai_api_key: $OPENAI_API_KEY

2. 使用配置文件启动服务器

mlflow gateway start --config-path /path/to/config.yaml --port 7000

3. 服务器启动并运行后，使用 OpenAI 客户端查询服务器

from openai import OpenAI

client = OpenAI(base_url="http://localhost:7000/v1")
completion = client.chat.completions.create(
    model="my-chat",
    messages=[{"role": "user", "content": "Hello"}],
)
print(completion.choices[0].message.content)

Unity Catalog 集成

有关如何将 MLflow AI 网关与 Unity Catalog 集成的信息，请参见Unity Catalog 集成。

网关服务器安全注意事项

请记住确保对运行 MLflow AI 网关的系统进行安全访问，以保护对这些密钥的访问。

保护网关服务器的一种有效方法是将其置于反向代理之后。这将允许反向代理处理传入请求并将其转发到 MLflow AI 网关。反向代理有效地保护您的应用程序免受直接暴露于互联网流量。

反向代理的一个常用选择是 Nginx。除了处理到您的应用程序的流量外，Nginx 还可以提供静态文件，如果您有多个应用程序实例正在运行，还可以对流量进行负载均衡。

此外，为了确保客户端和服务器之间的数据完整性和机密性，强烈建议在反向代理上启用 HTTPS。

除了反向代理之外，还建议在请求到达 MLflow AI 网关之前添加一个认证层。这可以是 HTTP 基本认证、OAuth 或任何适合您需求的任何其他方法。

例如，以下是 Nginx 的基本认证简单配置

http {
    server {
        listen 80;

        location / {
            auth_basic "Restricted Content";
            auth_basic_user_file /etc/nginx/.htpasswd;

            proxy_pass http://localhost:5000;  # Replace with the MLflow AI Gateway port
        }
    }
}

在此示例中，/etc/nginx/.htpasswd 是一个包含认证用户名和密码的文件。

这些措施，连同适当的网络设置，可以显著提高系统的安全性，并确保只有授权用户才能提交请求到您的 LLM 服务。