mlflow.gateway

mlflow.gateway.get_gateway_uri() → str

Warning

MLflow AI gateway is deprecated and has been replaced by the deployments API for generative AI. See https://mlflow.org/docs/latest/llms/gateway/migration.html for migration.

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.set_gateway_uri(gateway_uri: str)

Warning

MLflow AI gateway is deprecated and has been replaced by the deployments API for generative AI. See https://mlflow.org/docs/latest/llms/gateway/migration.html for migration.

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.

Args:

gateway_uri: The full uri of a running MLflow AI Gateway server or, if running on

Databricks, “databricks”.

class mlflow.gateway.base_models.ConfigModel

A pydantic model representing Gateway configuration data, such as an OpenAI completions route definition including route name, model name, API keys, etc.

class mlflow.gateway.config.AI21LabsConfig(*, ai21labs_api_key: str)

ai21labs_api_key: str

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

classmethod validate_ai21labs_api_key(value)

class mlflow.gateway.config.AWSBaseConfig(*, aws_region: Optional[str] = None)

aws_region: Optional[str]

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.AWSIdAndKey(*, aws_region: Optional[str] = None, aws_access_key_id: str, aws_secret_access_key: str, aws_session_token: Optional[str] = None)

aws_access_key_id: str

aws_secret_access_key: str

aws_session_token: Optional[str]

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.AWSRole(*, aws_region: Optional[str] = None, aws_role_arn: str, session_length_seconds: int = 900)

aws_role_arn: str

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

session_length_seconds: int

class mlflow.gateway.config.AliasedConfigModel

Enables use of field aliases in a configuration model for backwards compatibility

model_config: ClassVar[ConfigDict] = {'extra': 'ignore', 'populate_by_name': True, 'validate_by_alias': True, 'validate_by_name': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.AmazonBedrockConfig(*, aws_config: Union[mlflow.gateway.config.AWSRole, mlflow.gateway.config.AWSIdAndKey, mlflow.gateway.config.AWSBaseConfig])

aws_config: Union[mlflow.gateway.config.AWSRole, mlflow.gateway.config.AWSIdAndKey, mlflow.gateway.config.AWSBaseConfig]

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.AnthropicConfig(*, anthropic_api_key: str, anthropic_version: str = '2023-06-01')

anthropic_api_key: str

anthropic_version: str

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

classmethod validate_anthropic_api_key(value)

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

cohere_api_key: str

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

classmethod validate_cohere_api_key(value)

class mlflow.gateway.config.GatewayConfig(*, endpoints: list[mlflow.gateway.config.RouteConfig])

endpoints: list[mlflow.gateway.config.RouteConfig]

model_config: ClassVar[ConfigDict] = {'extra': 'ignore', 'populate_by_name': True, 'validate_by_alias': True, 'validate_by_name': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.GeminiConfig(*, gemini_api_key: str)

gemini_api_key: str

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

classmethod validate_gemini_api_key(value)

class mlflow.gateway.config.HuggingFaceTextGenerationInferenceConfig(*, hf_server_url: str)

hf_server_url: str

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.Limit(*, calls: int, key: Optional[str] = None, renewal_period: str)

calls: int

key: Optional[str]

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

renewal_period: str

class mlflow.gateway.config.LimitsConfig(*, limits: Optional[list[mlflow.gateway.config.Limit]] = [])

limits: Optional[list[mlflow.gateway.config.Limit]]

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class mlflow.gateway.config.MistralConfig(*, mistral_api_key: str)

mistral_api_key: str

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

classmethod validate_mistral_api_key(value)

class mlflow.gateway.config.MlflowModelServingConfig(*, model_server_url: str)

model_config: ClassVar[ConfigDict] = {'extra': 'ignore', 'protected_namespaces': ()}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_server_url: str

class mlflow.gateway.config.Model(*, name: Optional[str] = None, provider: Union[str, mlflow.gateway.config.Provider], config: Optional[mlflow.gateway.base_models.ConfigModel] = None)

config: Optional[mlflow.gateway.base_models.ConfigModel]

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: Optional[str]

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

classmethod validate_config(info, values)

classmethod validate_provider(value)

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

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: Optional[str]

provider: mlflow.gateway.config.Provider

class mlflow.gateway.config.MosaicMLConfig(*, mosaicml_api_key: str, mosaicml_api_base: Optional[str] = None)

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

mosaicml_api_base: Optional[str]

mosaicml_api_key: str

classmethod validate_mosaicml_api_key(value)

class mlflow.gateway.config.OpenAIAPIType(value)

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_api_base: Optional[str] = None, openai_api_version: Optional[str] = None, openai_deployment_name: Optional[str] = None, openai_organization: Optional[str] = None)

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

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(info: dict[str, typing.Any])

classmethod validate_openai_api_key(value)

class mlflow.gateway.config.PaLMConfig(*, palm_api_key: str)

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

palm_api_key: str

classmethod validate_palm_api_key(value)

class mlflow.gateway.config.Provider(value)

An enumeration.

AI21LABS = 'ai21labs'

AMAZON_BEDROCK = 'amazon-bedrock'

ANTHROPIC = 'anthropic'

BEDROCK = 'bedrock'

COHERE = 'cohere'

DATABRICKS = 'databricks'

DATABRICKS_MODEL_SERVING = 'databricks-model-serving'

GEMINI = 'gemini'

HUGGINGFACE_TEXT_GENERATION_INFERENCE = 'huggingface-text-generation-inference'

MISTRAL = 'mistral'

MLFLOW_MODEL_SERVING = 'mlflow-model-serving'

MOSAICML = 'mosaicml'

OPENAI = 'openai'

PALM = 'palm'

TOGETHERAI = 'togetherai'

classmethod values()

class mlflow.gateway.config.Route(*, name: str, route_type: str, model: mlflow.gateway.config.RouteModelInfo, route_url: str, limit: Optional[mlflow.gateway.config.Limit] = None)

class Config

json_schema_extra = {'example': {'model': {'name': 'gpt-4o-mini', 'provider': 'openai'}, 'name': 'openai-completions', 'route_type': 'llm/v1/completions', 'route_url': '/gateway/routes/completions/invocations'}}

limit: Optional[mlflow.gateway.config.Limit]

model: mlflow.gateway.config.RouteModelInfo

model_config: ClassVar[ConfigDict] = {'extra': 'ignore', 'json_schema_extra': {'example': {'model': {'name': 'gpt-4o-mini', 'provider': 'openai'}, 'name': 'openai-completions', 'route_type': 'llm/v1/completions', 'route_url': '/gateway/routes/completions/invocations'}}}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: str

route_type: str

route_url: str

to_endpoint()

class mlflow.gateway.config.RouteConfig(*, name: str, endpoint_type: mlflow.gateway.config.RouteType, model: mlflow.gateway.config.Model, limit: Optional[mlflow.gateway.config.Limit] = None)

endpoint_type: mlflow.gateway.config.RouteType

limit: Optional[mlflow.gateway.config.Limit]

model: mlflow.gateway.config.Model

model_config: ClassVar[ConfigDict] = {'extra': 'ignore', 'populate_by_name': True, 'validate_by_alias': True, 'validate_by_name': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: str

to_route() → mlflow.gateway.config.Route

classmethod validate_endpoint_name(route_name)

classmethod validate_limit(value)

classmethod validate_model(model)

classmethod validate_route_type(value)

classmethod validate_route_type_and_model_name(values)

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

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

name: Optional[str]

provider: str

class mlflow.gateway.config.RouteType(value)

An enumeration.

LLM_V1_CHAT = 'llm/v1/chat'

LLM_V1_COMPLETIONS = 'llm/v1/completions'

LLM_V1_EMBEDDINGS = 'llm/v1/embeddings'

class mlflow.gateway.config.TogetherAIConfig(*, togetherai_api_key: str)

model_config: ClassVar[ConfigDict] = {'extra': 'ignore'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

togetherai_api_key: str

classmethod validate_togetherai_api_key(value)