Skip to main content
Open on GitHub

MLflow AI Gateway for LLMs

The MLflow AI Gateway for LLMs is a powerful tool designed to streamline the usage and management of various large language model (LLM) providers, such as OpenAI and Anthropic, within an organization. It offers a high-level interface that simplifies the interaction with these services by providing a unified endpoint to handle specific LLM related requests.

安装与设置​

安装 mlflow 并包含 MLflow GenAI 依赖项：

pip install 'mlflow[genai]'

将OpenAI API密钥设置为环境变量：

export OPENAI_API_KEY=...

创建一个配置文件：

endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: openai
      name: text-davinci-003
      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 gateway start --config-path /path/to/config.yaml

示例由 MLflow​ 提供

The mlflow.langchain module provides an API for logging and loading LangChain models. This module exports multivariate LangChain models in the langchain flavor and univariate LangChain models in the pyfunc flavor.

查看更多信息，请参阅API文档和示例。

补全示例​

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

llm = Mlflow(
    target_uri="http://127.0.0.1:5000",
    endpoint="completions",
)

llm_chain = LLMChain(
    llm=Mlflow,
    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(chain, "model")

model = mlflow.pyfunc.load_model(model_info.model_uri)
print(model.predict([{"adjective": "funny"}]))
API 参考：LLMChain | MLflow

嵌入示例​

from langchain_community.embeddings import MlflowEmbeddings

embeddings = MlflowEmbeddings(
    target_uri="http://127.0.0.1:5000",
    endpoint="embeddings",
)

print(embeddings.embed_query("hello"))
print(embeddings.embed_documents(["hello"]))
API 参考：MlflowEmbeddings

聊天示例​

from langchain_community.chat_models import ChatMlflow
from langchain_core.messages import HumanMessage, SystemMessage

chat = ChatMlflow(
    target_uri="http://127.0.0.1:5000",
    endpoint="chat",
)

messages = [
    SystemMessage(
        content="You are a helpful assistant that translates English to French."
    ),
    HumanMessage(
        content="Translate this sentence from English to French: I love programming."
    ),
]
print(chat(messages))