mlflow.gateway

class mlflow.gateway.MlflowGatewayClient(gateway_uri: Optional[str] = None)[source]

Note

Experimental: This class may change or be removed in a future release without warning.

Client for interacting with the MLflow Gateway API.

Parameters: gateway_uri – Optional URI of the gateway. If not provided, attempts to resolve from first the stored result of set_gateway_uri(), then the environment variable MLFLOW_GATEWAY_URI.

create_route(name: str, route_type: str, model: Dict[str, Any]) → mlflow.gateway.config.Route [source]

Note

Experimental: This function may change or be removed in a future release without warning.

Create a new route in the Gateway.

Warning

This API is only available when running within Databricks. When running elsewhere, route configuration is handled via updates to the route configuration YAML file that is specified during Gateway server start.

Parameters

name – The name of the route.
route_type – The type of the route (e.g., ‘llm/v1/chat’, ‘llm/v1/completions’, ‘llm/v1/embeddings’).
model –
A dictionary representing the model details to be associated with the route. This dictionary should define:
- The model name (e.g., “gpt-3.5-turbo”)
- The provider (e.g., “openai”, “anthropic”)
- The configuration for the model used in the route

Returns

A serialized representation of the Route data structure, providing information about the name, type, and model details for the newly created route endpoint.

Raises

mlflow.MlflowException – If the function is not running within Databricks.

Note

See the official Databricks documentation for MLflow Gateway for examples of supported model configurations and how to dynamically create new routes within Databricks.

Example usage from within Databricks:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("databricks")

openai_api_key = ...

new_route = gateway_client.create_route(
    "my-new-route",
    "llm/v1/completions",
    {
        "name": "question-answering-bot-1",
        "provider": "openai",
        "config": {
            "openai_api_key": openai_api_key,
            "openai_api_version": "2023-05-10",
            "openai_api_type": "openai/v1/chat/completions",
        },
    },
)

delete_route(name: str) → None[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Delete an existing route in the Gateway.

Warning

This API is only available when running within Databricks. When running elsewhere, route deletion is handled by removing the corresponding entry from the route configuration YAML file that is specified during Gateway server start.

Parameters: name – The name of the route to delete.
Raises: mlflow.MlflowException – If the function is not running within Databricks.

Example usage from within Databricks:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("databricks")
gateway_client.delete_route("my-existing-route")

property gateway_uri

Get the current value for the URI of the MLflow Gateway.

Returns: The gateway URI.

get_route(name: str)[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Get a specific query route from the gateway. The routes that are available to retrieve are only those that have been configured through the MLflow Gateway Server configuration file (set during server start or through server update commands).

Parameters: name – The name of the route.
Returns: The returned data structure is a serialized representation of the Route data structure, giving information about the name, type, and model details (model name and provider) for the requested route endpoint.

query(route: str, data: Dict[str, Any])[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Submit a query to a configured provider route.

Parameters

route – The name of the route to submit the query to.
data – The data to send in the query. A dictionary representing the per-route specific structure required for a given provider.

Returns

The route’s response as a dictionary, standardized to the route type.

For chat, the structure should be:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("http://my.gateway:8888")

response = gateway_client.query(
    "my-chat-route",
    {"messages": [{"role": "user", "content": "Tell me a joke about rabbits"}, ...]},
)

For completions, the structure should be:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("http://my.gateway:8888")

response = gateway_client.query(
    "my-completions-route", {"prompt": "It's one small step for"}
)

For embeddings, the structure should be:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("http://my.gateway:8888")

response = gateway_client.query(
    "my-embeddings-route",
    {"text": ["It was the best of times", "It was the worst of times"]},
)

search_routes(page_token: Optional[str] = None) → PagedList[mlflow.gateway.config.Route][source]

Note

Experimental: This function may change or be removed in a future release without warning.

Search for routes in the Gateway.

Parameters: page_token – Token specifying the next page of results. It should be obtained from a prior search_routes() call.
Returns: Returns a list of all configured and initialized Route data for the MLflow Gateway Server. The return will be a list of dictionaries that detail the name, type, and model details of each active route endpoint.

mlflow.gateway.create_route(name: str, route_type: str, model: Dict[str, Any]) → mlflow.gateway.config.Route [source]

Note

Experimental: This function may change or be removed in a future release without warning.

Create a new route in the Gateway.

Warning

This API is only available when running within Databricks. When running elsewhere, route configuration is handled via updates to the route configuration YAML file that is specified during Gateway server start.

Parameters

name – The name of the route.
route_type – The type of the route (e.g., ‘llm/v1/chat’, ‘llm/v1/completions’, ‘llm/v1/embeddings’).
model –
A dictionary representing the model details to be associated with the route. This dictionary should define:
- The model name (e.g., “gpt-3.5-turbo”)
- The provider (e.g., “openai”, “anthropic”)
- The configuration for the model used in the route

Returns

A serialized representation of the Route data structure, providing information about the name, type, and model details for the newly created route endpoint.

Note

See the official Databricks documentation for MLflow Gateway for examples of supported model configurations and how to dynamically create new routes within Databricks.

Example usage from within Databricks:

from mlflow.gateway import set_gateway_uri, create_route

set_gateway_uri(gateway_uri="databricks")

openai_api_key = ...

create_route(
    "my-new-route",
    "llm/v1/completions",
    {
        "name": "question-answering-bot-1",
        "provider": "openai",
        "config": {
            "openai_api_key": openai_api_key,
            "openai_api_version": "2023-05-10",
            "openai_api_type": "openai/v1/chat/completions",
        },
    },
)

mlflow.gateway.delete_route(name: str) → None[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Delete an existing route in the Gateway.

Warning

This API is only available when running within Databricks. When running elsewhere, route deletion is handled by removing the corresponding entry from the route configuration YAML file that is specified during Gateway server start.

Parameters: name – The name of the route to delete.

Example usage from within Databricks:

from mlflow.gateway import set_gateway_uri, delete_route

set_gateway_uri(gateway_uri="databricks")

delete_route("my-new-route")

mlflow.gateway.get_gateway_uri() → str[source]: Note

Experimental: This function may change or be removed in a future release without warning.

Returns the currently set MLflow AI Gateway server uri iff set. If the Gateway uri has not been set by using set_gateway_uri, an MlflowException is raised.

mlflow.gateway.get_route(name: str) → mlflow.gateway.config.Route [source]

Note

Experimental: This function may change or be removed in a future release without warning.

Retrieves a specific route from the MLflow Gateway service.

This function creates an instance of MlflowGatewayClient and uses it to fetch a route by its name from the Gateway service.

Parameters: name – The name of the route to fetch.
Returns: An instance of the Route class representing the fetched route.

mlflow.gateway.query(route: str, data)[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Issues a query request to a configured service through a named route on the Gateway Server. This function will interface with a configured route name (examples below) and return the response from the provider in a standardized format.

Parameters

route – The name of the configured route. Route names can be obtained by running mlflow.gateway.search_routes()
data – The request payload to be submitted to the route. The exact configuration of the expected structure varies based on the route configuration.

Returns

The response from the configured route endpoint provider in a standardized format.

Chat example:

from mlflow.gateway import query, set_gateway_uri

set_gateway_uri(gateway_uri="http://my.gateway:9000")
response = query(
    "my_chat_route",
    {"messages": [{"role": "user", "content": "What is the best day of the week?"}]},
)

Completions example:

from mlflow.gateway import query, set_gateway_uri

set_gateway_uri(gateway_uri="http://my.gateway:9000")
response = query("a_completions_route", {"prompt": "Where do we go from"})

Embeddings example:

from mlflow.gateway import query, set_gateway_uri

set_gateway_uri(gateway_uri="http://my.gateway:9000")
response = query(
    "embeddings_route", {"text": ["I like spaghetti", "and sushi", "but not together"]}
)

mlflow.gateway.search_routes() → List[mlflow.gateway.config.Route][source]

Note

Experimental: This function may change or be removed in a future release without warning.

Searches for routes in the MLflow Gateway service.

This function creates an instance of MlflowGatewayClient and uses it to fetch a list of routes from the Gateway service.

Returns: A list of Route instances.

mlflow.gateway.set_gateway_uri(gateway_uri: str)[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Sets the uri of a configured and running MLflow AI Gateway server in a global context. Providing a valid uri and calling this function is required in order to use the MLflow AI Gateway fluent APIs.

Parameters: gateway_uri – The full uri of a running MLflow AI Gateway server or, if running on Databricks, “databricks”.

class mlflow.gateway.client.MlflowGatewayClient(gateway_uri: Optional[str] = None)[source]

Note

Experimental: This class may change or be removed in a future release without warning.

Client for interacting with the MLflow Gateway API.

Parameters: gateway_uri – Optional URI of the gateway. If not provided, attempts to resolve from first the stored result of set_gateway_uri(), then the environment variable MLFLOW_GATEWAY_URI.

create_route(name: str, route_type: str, model: Dict[str, Any]) → mlflow.gateway.config.Route [source]

Note

Experimental: This function may change or be removed in a future release without warning.

Create a new route in the Gateway.

Warning

This API is only available when running within Databricks. When running elsewhere, route configuration is handled via updates to the route configuration YAML file that is specified during Gateway server start.

Parameters

name – The name of the route.
route_type – The type of the route (e.g., ‘llm/v1/chat’, ‘llm/v1/completions’, ‘llm/v1/embeddings’).
model –
A dictionary representing the model details to be associated with the route. This dictionary should define:
- The model name (e.g., “gpt-3.5-turbo”)
- The provider (e.g., “openai”, “anthropic”)
- The configuration for the model used in the route

Returns

A serialized representation of the Route data structure, providing information about the name, type, and model details for the newly created route endpoint.

Raises

mlflow.MlflowException – If the function is not running within Databricks.

Note

See the official Databricks documentation for MLflow Gateway for examples of supported model configurations and how to dynamically create new routes within Databricks.

Example usage from within Databricks:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("databricks")

openai_api_key = ...

new_route = gateway_client.create_route(
    "my-new-route",
    "llm/v1/completions",
    {
        "name": "question-answering-bot-1",
        "provider": "openai",
        "config": {
            "openai_api_key": openai_api_key,
            "openai_api_version": "2023-05-10",
            "openai_api_type": "openai/v1/chat/completions",
        },
    },
)

delete_route(name: str) → None[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Delete an existing route in the Gateway.

Warning

This API is only available when running within Databricks. When running elsewhere, route deletion is handled by removing the corresponding entry from the route configuration YAML file that is specified during Gateway server start.

Parameters: name – The name of the route to delete.
Raises: mlflow.MlflowException – If the function is not running within Databricks.

Example usage from within Databricks:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("databricks")
gateway_client.delete_route("my-existing-route")

property gateway_uri

Get the current value for the URI of the MLflow Gateway.

Returns: The gateway URI.

get_route(name: str)[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Get a specific query route from the gateway. The routes that are available to retrieve are only those that have been configured through the MLflow Gateway Server configuration file (set during server start or through server update commands).

Parameters: name – The name of the route.
Returns: The returned data structure is a serialized representation of the Route data structure, giving information about the name, type, and model details (model name and provider) for the requested route endpoint.

query(route: str, data: Dict[str, Any])[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Submit a query to a configured provider route.

Parameters

route – The name of the route to submit the query to.
data – The data to send in the query. A dictionary representing the per-route specific structure required for a given provider.

Returns

The route’s response as a dictionary, standardized to the route type.

For chat, the structure should be:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("http://my.gateway:8888")

response = gateway_client.query(
    "my-chat-route",
    {"messages": [{"role": "user", "content": "Tell me a joke about rabbits"}, ...]},
)

For completions, the structure should be:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("http://my.gateway:8888")

response = gateway_client.query(
    "my-completions-route", {"prompt": "It's one small step for"}
)

For embeddings, the structure should be:

from mlflow.gateway import MlflowGatewayClient

gateway_client = MlflowGatewayClient("http://my.gateway:8888")

response = gateway_client.query(
    "my-embeddings-route",
    {"text": ["It was the best of times", "It was the worst of times"]},
)

search_routes(page_token: Optional[str] = None) → PagedList[mlflow.gateway.config.Route][source]

Note

Experimental: This function may change or be removed in a future release without warning.

Search for routes in the Gateway.

Parameters: page_token – Token specifying the next page of results. It should be obtained from a prior search_routes() call.
Returns: Returns a list of all configured and initialized Route data for the MLflow Gateway Server. The return will be a list of dictionaries that detail the name, type, and model details of each active route endpoint.

class mlflow.gateway.config.AnthropicConfig(*, anthropic_api_key: str)[source]

anthropic_api_key: str

classmethod validate_anthropic_api_key(value)[source]

class mlflow.gateway.config.CohereConfig(*, cohere_api_key: str)[source]

cohere_api_key: str

classmethod validate_cohere_api_key(value)[source]

class mlflow.gateway.config.GatewayConfig(*, routes: List[mlflow.gateway.config.RouteConfig])[source]

routes: List[mlflow.gateway.config.RouteConfig]

class mlflow.gateway.config.Model(*, name: Optional[str] = None, provider: Union[str, mlflow.gateway.config.Provider], config: Optional[Union[mlflow.gateway.config.CohereConfig, mlflow.gateway.config.OpenAIConfig, mlflow.gateway.config.AnthropicConfig]] = None)[source]

config: Optional[Union[mlflow.gateway.config.CohereConfig, mlflow.gateway.config.OpenAIConfig, mlflow.gateway.config.AnthropicConfig]]

name: Optional[str]

provider: Union[str, mlflow.gateway.config.Provider]

classmethod validate_config(config, values)[source]

classmethod validate_provider(value)[source]

class mlflow.gateway.config.ModelInfo(*, name: Optional[str] = None, provider: mlflow.gateway.config.Provider)[source]

name: Optional[str]

provider: mlflow.gateway.config.Provider

class mlflow.gateway.config.OpenAIAPIType(value)[source]

An enumeration.

AZURE = 'azure'

AZUREAD = 'azuread'

OPENAI = 'openai'

class mlflow.gateway.config.OpenAIConfig(*, openai_api_key: str, openai_api_type: mlflow.gateway.config.OpenAIAPIType = <OpenAIAPIType.OPENAI: 'openai'>, openai_api_base: Optional[str] = None, openai_api_version: Optional[str] = None, openai_deployment_name: Optional[str] = None, openai_organization: Optional[str] = None)[source]

openai_api_base: Optional[str]

openai_api_key: str

openai_api_type: mlflow.gateway.config.OpenAIAPIType

openai_api_version: Optional[str]

openai_deployment_name: Optional[str]

openai_organization: Optional[str]

classmethod validate_field_compatibility(config: Dict[str, Any])[source]

classmethod validate_openai_api_key(value)[source]

class mlflow.gateway.config.Provider(value)[source]

An enumeration.

ANTHROPIC = 'anthropic'

COHERE = 'cohere'

DATABRICKS_MODEL_SERVING = 'databricks-model-serving'

OPENAI = 'openai'

classmethod values()[source]

class mlflow.gateway.config.Route(*, name: str, route_type: mlflow.gateway.config.RouteType, model: mlflow.gateway.config.ModelInfo, route_url: str)[source]

class Config[source]

schema_extra = {'example': {'model': {'name': 'gpt-3.5-turbo', 'provider': 'openai'}, 'name': 'openai-completions', 'route_type': 'llm/v1/completions', 'route_url': '/gateway/routes/completions/invocations'}}

model: mlflow.gateway.config.ModelInfo

name: str

route_type: mlflow.gateway.config.RouteType

route_url: str

class mlflow.gateway.config.RouteConfig(*, name: str, route_type: mlflow.gateway.config.RouteType, model: mlflow.gateway.config.Model)[source]

model: mlflow.gateway.config.Model

name: str

route_type: mlflow.gateway.config.RouteType

to_route() → mlflow.gateway.config.Route [source]

classmethod validate_endpoint_name(route_name)[source]

classmethod validate_model(model)[source]

classmethod validate_route_type(value)[source]

class mlflow.gateway.config.RouteType(value)[source]

An enumeration.

LLM_V1_CHAT = 'llm/v1/chat'

LLM_V1_COMPLETIONS = 'llm/v1/completions'

LLM_V1_EMBEDDINGS = 'llm/v1/embeddings'