MLflow AI Gateway

警告

MLflow AI Gateway 不支持 Windows。

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

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

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

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

教程和指南

如果您有兴趣立即深入分步指南，以尽快启动并运行 MLflow AI Gateway，下面的指南将是您的最佳首选。

查看网关服务器入门指南

快速入门

以下指南将帮助您启动并运行，使用三端点配置连接到 OpenAI 服务，以实现对话、补全和嵌入功能。

步骤 1：安装 MLflow AI Gateway

首先，您需要在您的机器上安装 MLflow AI Gateway。您可以从 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 Gateway 应暴露的端点。让我们创建一个包含三个端点的文件，使用 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 Gateway 的系统上的某个位置。

步骤 4：启动网关服务器

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

使用 MLflow AI Gateway 的 start 命令并指定配置文件的路径

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

如果您未指定主机，将使用本地主机地址。

如果您未指定端口，将使用端口 5000。

gunicorn 的工作进程数默认为 2 个。

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

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

文档端点允许与端点直接交互，并通过点击端点定义条目中的“立即尝试”选项来提交实际请求给提供商服务。

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

有关详细信息，请参阅客户端 API 部分。

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

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

步骤 8：比较提供商模型

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

首先，更新 MLflow AI Gateway 配置 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 模型来查询这些端点，并使用提示工程等技术构建应用程序特定的逻辑。有关更多信息，请参阅网关服务器和 MLflow 模型。

概念

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

提供商

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

支持的提供商

下表显示了 MLflow AI Gateway 中每个 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	✅	✅	✅
Gemini	✅*	✅*	✅**

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

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

注意

* MLflow 模型服务仅在输出返回格式与端点兼容时才适用于对话或补全。响应必须符合 {"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 Gateway 支持以下提供商

• 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 提供的模型。
• gemini：用于 Gemini 提供的模型。

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

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

端点

端点是 MLflow AI Gateway 运行方式的核心。每个端点都充当用户的代理端点，将请求转发到配置文件中指定的底层模型和提供商。

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

• name：这是端点的唯一标识符。通过 MLflow AI Gateway 进行 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

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

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

在配置文件中定义端点时，请确保每个名称都是唯一的，以防止冲突。重复的端点名称将引发 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 Gateway 在尝试将请求路由到该模型时将返回 HTTP 4xx 错误。

重要

请务必查阅指定提供商的最新文档，以确保您想要使用的模型支持您正在配置的端点类型。

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

配置网关服务器

MLflow AI Gateway 依赖于用户提供的 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 Gateway 配置非常容易更新。只需编辑配置文件并保存更改，MLflow AI Gateway 将自动更新端点，无需中断或停机。这使您可以在保持应用程序稳定可靠的同时，尝试新的提供商或模型类型。

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

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

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

警告

MLflow AI Gateway 提供对计费的外部 LLM 服务的直接访问。强烈建议限制对此服务器的访问。有关指南，请参阅安全部分。

如果您倾向于使用环境变量（推荐），可以在您的 shell 环境中定义它。例如

export OPENAI_API_KEY="your_openai_api_key"

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

AI 网关服务器配置详情

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

配置文件以 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 Gateway 将尝试使用标准凭证提供商链扮演此角色，如果角色凭证过期，将对其进行续订。

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

MLflow 模型服务

配置参数	必需	默认值	描述
model_server_url	是	不适用	这是 MLflow 模型服务器的 URL。

请注意，在使用 MLflow 模型服务时，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 密钥。

Gemini

配置参数	必需	默认值	描述
gemini_api_key	是	不适用	这是 Gemini 服务的 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 Gateway 配置并启动后，即可接收来自用户的流量。

标准查询参数

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

补全

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

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

对话

类型为 llm/v1/chat 的对话端点的标准参数是

查询参数	类型	必需	默认值	描述
messages	消息数组	是	不适用	对话中的消息列表，从中生成新消息（对话补全）。有关消息结构的信息，请参阅消息。
n	整数	否	1	为指定提示生成的对话补全数量，介于 1 到 5 之间。
temperature	浮点数	否	0.0	要使用的采样温度，介于 0 和 1 之间。值越高，输出越随机；值越低，输出越确定。
max_tokens	整数	否	无	最大补全长度，介于 1 到无穷大（无限制）之间。
stop	字符串数组	否	无	模型应停止生成令牌并返回对话补全的序列。

消息

每个对话消息是一个字符串字典，包含以下字段

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

嵌入

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

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

附加查询参数

除了标准查询参数外，您还可以将端点提供商支持的任何附加参数作为查询的一部分传递。例如

• 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 Gateway 端点提交查询请求的示例

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 模式。

对话

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 Gateway 的框架，提供了一个自动化的交互式 API 文档界面，可通过 "/docs" 端点（例如，http://my.deployments:9000/docs）访问。这个交互式界面对于探索和测试可用的 API 端点非常方便。

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

MLflow Python 客户端 API

MlflowDeploymentClient 是面向用户的客户端 API，用于与 MLflow AI Gateway 交互。它通过简单易用的 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 部署集成。此集成使用户能够在网关服务器中使用提示工程、检索增强生成和其他 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, name="model")

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

MLflow 模型

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

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

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

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

您还可以构建和部署调用 MLflow AI Gateway 的 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",
        name="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.genai.evaluate() 一起使用，或者像其他模型一样使用部署 API。

REST API

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

以下是一些您可能如何使用 curl 与 MLflow AI Gateway 交互的示例

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

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

   • 补全

     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 Gateway 的 URL。

插件式 LLM 提供商

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

要创建自定义插件，您需要实现一个继承自 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 Gateway 相同的环境中。

重要

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

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

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 Gateway API 文档

API 文档

OpenAI 兼容性

MLflow AI Gateway 与 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="https://:7000/v1")
completion = client.chat.completions.create(
    model="my-chat",
    messages=[{"role": "user", "content": "Hello"}],
)
print(completion.choices[0].message.content)

Unity Catalog 集成

请参阅Unity Catalog 集成，了解如何将 MLflow AI Gateway 与 Unity Catalog 集成。

网关服务器安全注意事项

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

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

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

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

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

例如，以下是带有基本身份验证的 Nginx 简单配置

http {
    server {
        listen 80;

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

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

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

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