MLflow AI Gateway

danger

MLflow AI Gateway has been deprecated. Please use MLflow Deployments for LLMs instead.

The MLflow AI Gateway service 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. See the MLflow AI Gateway documentation for more details.

Installation and Setup

Install mlflow with MLflow AI Gateway dependencies:

pip install 'mlflow[gateway]'

Set the OpenAI API key as an environment variable:

export OPENAI_API_KEY=...

Create a configuration file:

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

Start the Gateway server:

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

Example provided by 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.

See the API documentation and examples.

Completions Example

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

gateway = MlflowAIGateway(
    gateway_uri="http://127.0.0.1:5000",
    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(chain, "model")

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

Embeddings Example

from langchain.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"]))

Chat Example

from langchain.chat_models import ChatMLflowAIGateway
from langchain.schema import HumanMessage, SystemMessage

chat = ChatMLflowAIGateway(
    gateway_uri="http://127.0.0.1:5000",
    route="chat",
    params={
        "temperature": 0.1
    }
)

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))

Databricks MLflow AI Gateway

Databricks MLflow AI Gateway is in private preview. Please contact a Databricks representative to enroll in the preview.

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.llms import MlflowAIGateway

gateway = MlflowAIGateway(
    gateway_uri="databricks",
    route="completions",
)

llm_chain = LLMChain(
    llm=gateway,
    prompt=PromptTemplate(
        input_variables=["adjective"],
        template="Tell me a {adjective} joke",
    ),
)
result = llm_chain.run(adjective="funny")
print(result)