MLflow AI Gateway 接入 Google PaLM：completions、embeddings 与 chat 三端点配置实战-迪斯科星球

MLflow AI Gateway 接入 Google PaLM：completions、embeddings 与 chat 三端点配置实战

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow

本篇技术指南以 MLflow 仓库中 examples/gateway/palm 示例为骨架，讲解如何通过 MLflow AI Gateway 统一接入 Google PaLM（Pathways Language Model）的文本补全、文本向量化和对话生成三大能力。读完本文，你将掌握 PaLM 端点的 YAML 配置写法、环境变量密钥注入方式、网关启动流程，以及通过mlflow.deployments客户端发起三种请求的完整调用链，并了解 Gateway 底层参数转换与请求路由的实现细节。

PaLM 示例在 MLflow AI Gateway 中的定位

MLflow AI Gateway 允许在单个服务实例上为多个 LLM 提供商配置多个查询端点（endpoint），每个端点可以绑定不同的提供商、模型版本和端点类型。examples/gateway/palm/目录提供了一个完整可运行的 PaLM 接入示例，包含三个文件：

README.md：端点配置说明与 API Key 设置方法；
config.yaml：声明 PaLM 三个端点的完整 YAML 配置；
example.py：使用 Python 客户端实际调用三个端点的可执行脚本。

该示例在 examples/gateway/README.md 中与 OpenAI、Anthropic、Cohere、MosaicML 等提供商并列，作为配置与验证某个 provider 的标准样板。示例中的端点名称（completions、chat、embeddings）仅用于说明，官方建议在真实业务中改为语义明确的自定义名称，避免端点名冲突。

三端点配置解析：completions / embeddings / chat

config.yaml 是示例的核心配置文件，它在同一个配置文件中声明了三个端点，分别对应 PaLM 的三类模型：

endpoints: - name: completions endpoint_type: llm/v1/completions model: provider: palm name: text-bison-001 config: palm_api_key: $PALM_API_KEY - name: embeddings endpoint_type: llm/v1/embeddings model: provider: palm name: embedding-gecko-001 config: palm_api_key: $PALM_API_KEY - name: chat endpoint_type: llm/v1/chat model: provider: palm name: chat-bison-001 config: palm_api_key: $PALM_API_KEY

各字段含义如下：

字段	值	说明
`name`	`completions`/`embeddings`/`chat`	端点在网关内的唯一标识，后续调用时通过该名称路由请求
`endpoint_type`	`llm/v1/completions`、`llm/v1/embeddings`、`llm/v1/chat`	MLflow 网关标准化的端点类型，统一屏蔽各厂商原生 API 差异
`model.provider`	`palm`	提供商标识，映射到源码中的`Provider.PALM = "palm"`（见 mlflow/gateway/config.py）
`model.name`	`text-bison-001`等	PaLM 具体模型名，会被拼接进请求 URL 调用对应的生成接口
`model.config.palm_api_key`	`$PALM_API_KEY`	环境变量引用，网关启动时解析为真实密钥

三个端点的分工是：text-bison-001负责文本续写，embedding-gecko-001负责把文本转为向量，chat-bison-001负责多轮对话。

密钥字段的底层校验

palm_api_key的写法并非随意命名。在 mlflow/gateway/config.py 中定义了PaLMConfig配置模型，它强制要求该字段存在，并通过_resolve_api_key_from_input校验器支持两种取值方式：直接写明文密钥，或写$ENV_VAR_NAME形式的环境变量引用。因此$PALM_API_KEY会在配置加载阶段被解析为环境变量PALM_API_KEY的值，密钥本身不会明文落盘在配置文件中。

设置 PaLM API Key

调用 PaLM 需要 Google 提供的 PaLM API 密钥。设置方式与示例 README 一致，将密钥写入环境变量：

export PALM_API_KEY=...

请将...替换为在 Google AI 开发者平台申请的密钥。配置中的$PALM_API_KEY引用的正是该环境变量。需要说明的是，由于 PaLM 模型系列已逐步被 Gemini 取代，当前仓库的PaLMProvider在初始化时会抛出FutureWarning，提示“PaLM provider 已弃用并将在未来的 MLflow 版本中移除”（见 mlflow/gateway/providers/palm.py）。如需长期使用，建议评估仓库中同目录下的 Gemini provider 作为迁移方向。

启动网关并验证端点

按照 examples/gateway/README.md 的流程，PaLM 示例同样适用：

安装带 genai 依赖的 MLflow：

pip install 'mlflow[genai]'

启动网关，指定本示例的配置文件：

mlflow gateway start --config-path examples/gateway/palm/config.yaml --port 7000

网关启动后，可访问http://127.0.0.1:7000/docs查看交互式 API 文档。

通过 Python 客户端调用三个端点

example.py 展示了完整的调用代码，核心是mlflow.deployments.get_deploy_client：

from mlflow.deployments import get_deploy_client def main(): client = get_deploy_client("http://localhost:7000") print(f"PaLM endpoints: {client.list_endpoints()}\n") print(f"PaLM completions endpoint info: {client.get_endpoint(endpoint='completions')}\n") # Completions request response_completions = client.predict( endpoint="completions", inputs={ "prompt": "What is the world record for flapjack consumption in a single sitting?", "temperature": 0.1, }, ) print(f"PaLM response for completions: {response_completions}") # Embeddings request response_embeddings = client.predict( endpoint="embeddings", inputs={"input": ["Do you carry the Storm Trooper costume in size 2T?"]}, ) print(f"PaLM response for embeddings: {response_embeddings}") # Chat example response_chat = client.predict( endpoint="chat", inputs={ "messages": [ { "role": "system", "content": "You are a talented European rapper with a background in US history", }, { "role": "user", "content": "Please recite the preamble to the US Constitution as if it were " "written today by a rapper from Reykjavík", }, ] }, ) print(f"PaLM response for chat: {response_chat}") if __name__ == "__main__": main()

在网关运行期间执行python examples/gateway/palm/example.py，即可依次观察到三类响应。

三类请求的输入约定

completions（补全）：入参为prompt字符串，可附带temperature等采样参数，返回补全文本；
embeddings（向量化）：入参为input，值为字符串列表，一次可批量向量化多条文本；
chat（对话）：入参为messages消息列表，每条消息含role与content，其中role使用 OpenAI 风格的system/user角色约定。

这种“统一网关输入、后端自动适配”的设计正是 MLflow AI Gateway 的价值所在：业务代码只需面向一套标准化 Schema 编程，切换底层模型时无需改动调用端。

源码视角：PaLM Provider 的请求转换链路

PaLMProvider的实现（mlflow/gateway/providers/palm.py）揭示了网关是如何把统一的端到端请求转换为 PaLM 原生 API 的，其核心逻辑可以概括为三个层次。

1. Provider 注册

在 mlflow/gateway/provider_registry.py 中，PaLMProvider被注册到全局ProviderRegistry，key 为Provider.PALM。网关加载配置时，根据model.provider: palm查表找到该类，从而把completions等端点路由到正确的 provider 实现。

2. 密钥注入与请求地址

PaLMProvider._request方法把配置中的palm_api_key放入x-goog-api-key请求头，并把请求发往https://generativelanguage.googleapis.com/v1beta3/models/基地址（mlflow/gateway/providers/palm.py）。结合model.name，最终请求路径形如：

补全：{model_name}:generateText
对话：{model_name}:generateMessage
向量：{model_name}:batchEmbedText

3. 参数名映射与语义适配

网关统一参数与 PaLM 原生参数存在命名差异，PaLMProvider通过rename_payload_keys完成转换（mlflow/gateway/providers/palm.py）：

MLflow 网关统一参数	PaLM 原生参数	适用端点
`stop`	`stopSequences`	completions / chat
`n`	`candidateCount`	completions / chat
`max_tokens`	`maxOutputTokens`	completions
`input`	`texts`	embeddings

此外还有两处关键适配：

温度缩放：MLflow 网关的温度范围是 0–2，而 PaLM 是 0–1，因此 provider 会把网关传入的temperature乘以 0.5 再转发（mlflow/gateway/providers/palm.py）；
角色字段重命名：chat 请求中 OpenAI 风格的role被改写为 PaLM 的author，messages、examples、context则被统一打包进prompt字段（mlflow/gateway/providers/palm.py）。

同时，provider 会拒绝接收 PaLM 原生参数名（如直接传stopSequences或candidateCount会返回 422 错误并提示改用网关参数名），并针对 chat 端点限制max_tokens的使用——这些都是网关保证请求格式统一的具体手段。

响应标准化

PaLM 的原始响应经过 provider 加工后，被转换为网关统一的 Schema：补全与对话响应均含created、model、choices列表和usage统计，向量响应则返回data列表（每项含embedding向量与index）。由于 PaLM API 不返回 token 统计信息，usage中的 token 数统一置为None（mlflow/gateway/providers/palm.py）。这意味着调用方拿到的是与 OpenAI 等提供商一致的响应结构，可以无缝复用现有的解析代码。

测试佐证与验证

仓库的测试目录 tests/gateway/providers/test_palm.py 提供了对上述行为的完整验证：测试用例构造了与示例配置一致的completions端点配置（provider 为palm、模型text-bison），通过 mock 拦截aiohttp.ClientSession.post断言请求参数映射、密钥注入与响应转换是否符合预期（tests/gateway/providers/test_palm.py）。读者在修改或调试 PaLM 接入时，可以直接参考或运行该测试文件来确认行为没有回归。

小结

通过examples/gateway/palm/这份示例，我们完成了 PaLM 从配置、启动到三种端点调用的全链路梳理：config.yaml声明端点与模型、$PALM_API_KEY注入密钥、example.py通过统一客户端发起请求。底层PaLMProvider则负责参数映射、温度缩放与响应标准化，让业务代码始终面向 MLflow 网关的统一 API。结合示例、源码与测试三份证据，读者既可以照抄配置文件快速跑通 PaLM 接入，也能深入理解网关 provider 的扩展机制，为迁移到 Gemini 或接入其他提供商打下基础。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析