技术背景介绍

大语言模型（LLM）的迅速发展引发了企业对使用和管理这些模型的需求。MLflow AI Gateway 是为了解决这一需求而设计的工具，它简化了与多个大语言模型提供商（如 OpenAI）的交互。尽管现在已经建议使用 MLflow Deployments 代替 MLflow AI Gateway，这个服务曾提供了一个统一的接口来处理 LLM 相关的请求。

核心原理解析

MLflow AI Gateway 通过配置文件定义了不同的路由，每个路由对应一个特定的 LLM 提供商和模型。这样，用户就能通过一个统一的端点发送请求，而不需要直接与每个提供商的 API 接口交互。MLflow 的这个模块尤其适用于在开发环境中需要频繁调整或测试多个模型的情境。

代码实现演示

接下来，我们通过代码示例来演示如何使用 MLflow AI Gateway 来完成文本生成和向量嵌入任务。

安装与配置

首先，安装 MLflow 并进行必要的配置：

pip install 'mlflow[gateway]'

设置 OpenAI API 密钥：

export OPENAI_API_KEY=your-openai-api-key

创建配置文件 config.yaml：

routes:
  - name: completions
    route_type: llm/v1/completions
    model:
      provider: openai
      name: text-davinci-003
      config:
        openai_api_key: $OPENAI_API_KEY

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

启动 Gateway 服务器：

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

执行文本生成任务

以下是一个通过 LLMChain 执行文本生成的示例：

import mlflow
from langchain.chains import LLMChain, PromptTemplate
from langchain_community.llms import MlflowAIGateway

# 设置 Gateway
gateway = MlflowAIGateway(
    gateway_uri="http://127.0.0.1:5000",  # 本地启动的 Gateway
    route="completions",
    params={
        "temperature": 0.0,
        "top_p": 0.1,
    },
)

# 定义文本生成任务
llm_chain = LLMChain(
    llm=gateway,
    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"}]))

执行向量嵌入任务

from langchain_community.embeddings import MlflowAIGatewayEmbeddings

# 向量嵌入配置
embeddings = MlflowAIGatewayEmbeddings(
    gateway_uri="http://127.0.0.1:5000",
    route="embeddings",
)

# 执行嵌入查询
print(embeddings.embed_query("hello"))
print(embeddings.embed_documents(["hello"]))

应用场景分析

MLflow AI Gateway 非常适合用于企业内部需要经常切换或比较多种 LLM 模型的场景，特别是当这些需求频繁变化时。通过统一的管理和调用接口，开发者可以更加专注于业务逻辑的开发，而不是模型的具体实现细节。

实践建议

• 定期更新和管理 LLM 的配置文件，以适应新的模型版本和API变更。
• 利用 MLflow 提供的模型记录功能，确保训练和运行模型的可追溯性。
• 尽量在开发环境中测试模型性能，以便根据实际需求调整参数。

如果遇到问题欢迎在评论区交流。
—END—