star

venv/lib/python3.9/site-packages/mlflow/gateway/__init__.py

@@ -0,0 +1,24 @@
+from mlflow.gateway.client import MlflowGatewayClient
+from mlflow.gateway.fluent import (
+    create_route,
+    delete_route,
+    get_limits,
+    get_route,
+    query,
+    search_routes,
+    set_limits,
+)
+from mlflow.gateway.utils import get_gateway_uri, set_gateway_uri
+
+__all__ = [
+    "create_route",
+    "delete_route",
+    "get_route",
+    "set_limits",
+    "get_limits",
+    "get_gateway_uri",
+    "MlflowGatewayClient",
+    "query",
+    "search_routes",
+    "set_gateway_uri",
+]

venv/lib/python3.9/site-packages/mlflow/gateway/app.py

@@ -0,0 +1,436 @@
+import functools
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.openapi.docs import get_swagger_ui_html
+from fastapi.responses import FileResponse, RedirectResponse
+from pydantic import BaseModel
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.errors import RateLimitExceeded
+from slowapi.util import get_remote_address
+
+from mlflow.deployments.server.config import Endpoint
+from mlflow.deployments.server.constants import (
+    MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE,
+    MLFLOW_DEPLOYMENTS_ENDPOINTS_BASE,
+    MLFLOW_DEPLOYMENTS_HEALTH_ENDPOINT,
+    MLFLOW_DEPLOYMENTS_LIMITS_BASE,
+    MLFLOW_DEPLOYMENTS_LIST_ENDPOINTS_PAGE_SIZE,
+    MLFLOW_DEPLOYMENTS_QUERY_SUFFIX,
+)
+from mlflow.environment_variables import (
+    MLFLOW_GATEWAY_CONFIG,
+    MLFLOW_GATEWAY_RATE_LIMITS_STORAGE_URI,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.base_models import SetLimitsModel
+from mlflow.gateway.config import (
+    GatewayConfig,
+    LimitsConfig,
+    Route,
+    RouteConfig,
+    RouteType,
+    _load_route_config,
+)
+from mlflow.gateway.constants import (
+    MLFLOW_GATEWAY_CRUD_ROUTE_BASE,
+    MLFLOW_GATEWAY_HEALTH_ENDPOINT,
+    MLFLOW_GATEWAY_LIMITS_BASE,
+    MLFLOW_GATEWAY_ROUTE_BASE,
+    MLFLOW_GATEWAY_SEARCH_ROUTES_PAGE_SIZE,
+    MLFLOW_QUERY_SUFFIX,
+)
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers import get_provider
+from mlflow.gateway.schemas import chat, completions, embeddings
+from mlflow.gateway.utils import SearchRoutesToken, make_streaming_response
+from mlflow.version import VERSION
+
+
+class GatewayAPI(FastAPI):
+    def __init__(self, config: GatewayConfig, limiter: Limiter, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.state.limiter = limiter
+        self.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+        self.dynamic_routes: dict[str, RouteConfig] = {}
+        self.set_dynamic_routes(config, limiter)
+
+    def set_dynamic_routes(self, config: GatewayConfig, limiter: Limiter) -> None:
+        self.dynamic_routes.clear()
+        for route in config.routes:
+            # TODO: Remove deployments server URLs after deprecation window elapses
+            self.add_api_route(
+                path=(
+                    MLFLOW_DEPLOYMENTS_ENDPOINTS_BASE + route.name + MLFLOW_DEPLOYMENTS_QUERY_SUFFIX
+                ),
+                endpoint=_route_type_to_endpoint(route, limiter, "deployments"),
+                methods=["POST"],
+            )
+            self.add_api_route(
+                path=f"{MLFLOW_GATEWAY_ROUTE_BASE}{route.name}{MLFLOW_QUERY_SUFFIX}",
+                endpoint=_route_type_to_endpoint(route, limiter, "gateway"),
+                methods=["POST"],
+                include_in_schema=False,
+            )
+            self.dynamic_routes[route.name] = route
+
+    def get_dynamic_route(self, route_name: str) -> Optional[Route]:
+        return r.to_route() if (r := self.dynamic_routes.get(route_name)) else None
+
+
+def _translate_http_exception(func):
+    """
+    Decorator for translating MLflow exceptions to HTTP exceptions
+    """
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except AIGatewayException as e:
+            raise HTTPException(status_code=e.status_code, detail=e.detail)
+
+    return wrapper
+
+
+def _create_chat_endpoint(config: RouteConfig):
+    prov = get_provider(config.model.provider)(config)
+
+    # https://slowapi.readthedocs.io/en/latest/#limitations-and-known-issues
+    @_translate_http_exception
+    async def _chat(
+        request: Request, payload: chat.RequestPayload
+    ) -> Union[chat.ResponsePayload, chat.StreamResponsePayload]:
+        if payload.stream:
+            return await make_streaming_response(prov.chat_stream(payload))
+        else:
+            return await prov.chat(payload)
+
+    return _chat
+
+
+def _create_completions_endpoint(config: RouteConfig):
+    prov = get_provider(config.model.provider)(config)
+
+    @_translate_http_exception
+    async def _completions(
+        request: Request, payload: completions.RequestPayload
+    ) -> Union[completions.ResponsePayload, completions.StreamResponsePayload]:
+        if payload.stream:
+            return await make_streaming_response(prov.completions_stream(payload))
+        else:
+            return await prov.completions(payload)
+
+    return _completions
+
+
+def _create_embeddings_endpoint(config: RouteConfig):
+    prov = get_provider(config.model.provider)(config)
+
+    @_translate_http_exception
+    async def _embeddings(
+        request: Request, payload: embeddings.RequestPayload
+    ) -> embeddings.ResponsePayload:
+        return await prov.embeddings(payload)
+
+    return _embeddings
+
+
+async def _custom(request: Request):
+    return request.json()
+
+
+def _route_type_to_endpoint(config: RouteConfig, limiter: Limiter, key: str):
+    provider_to_factory = {
+        RouteType.LLM_V1_CHAT: _create_chat_endpoint,
+        RouteType.LLM_V1_COMPLETIONS: _create_completions_endpoint,
+        RouteType.LLM_V1_EMBEDDINGS: _create_embeddings_endpoint,
+    }
+    if factory := provider_to_factory.get(config.route_type):
+        handler = factory(config)
+        if limit := config.limit:
+            limit_value = f"{limit.calls}/{limit.renewal_period}"
+            handler.__name__ = f"{handler.__name__}_{config.name}_{key}"
+            return limiter.limit(limit_value)(handler)
+        else:
+            return handler
+
+    raise HTTPException(
+        status_code=404,
+        detail=f"Unexpected route type {config.route_type!r} for route {config.name!r}.",
+    )
+
+
+class HealthResponse(BaseModel):
+    status: str
+
+
+class ListEndpointsResponse(BaseModel):
+    endpoints: list[Endpoint]
+    next_page_token: Optional[str] = None
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "endpoints": [
+                    {
+                        "name": "openai-chat",
+                        "endpoint_type": "llm/v1/chat",
+                        "model": {
+                            "name": "gpt-4o-mini",
+                            "provider": "openai",
+                        },
+                        "limit": {"calls": 1, "key": None, "renewal_period": "minute"},
+                    },
+                    {
+                        "name": "anthropic-completions",
+                        "endpoint_type": "llm/v1/completions",
+                        "model": {
+                            "name": "claude-instant-100k",
+                            "provider": "anthropic",
+                        },
+                    },
+                    {
+                        "name": "cohere-embeddings",
+                        "endpoint_type": "llm/v1/embeddings",
+                        "model": {
+                            "name": "embed-english-v2.0",
+                            "provider": "cohere",
+                        },
+                    },
+                ],
+                "next_page_token": "eyJpbmRleCI6IDExfQ==",
+            }
+        }
+
+
+class SearchRoutesResponse(BaseModel):
+    routes: list[Route]
+    next_page_token: Optional[str] = None
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "routes": [
+                    {
+                        "name": "openai-chat",
+                        "route_type": "llm/v1/chat",
+                        "model": {
+                            "name": "gpt-4o-mini",
+                            "provider": "openai",
+                        },
+                    },
+                    {
+                        "name": "anthropic-completions",
+                        "route_type": "llm/v1/completions",
+                        "model": {
+                            "name": "claude-instant-100k",
+                            "provider": "anthropic",
+                        },
+                    },
+                    {
+                        "name": "cohere-embeddings",
+                        "route_type": "llm/v1/embeddings",
+                        "model": {
+                            "name": "embed-english-v2.0",
+                            "provider": "cohere",
+                        },
+                    },
+                ],
+                "next_page_token": "eyJpbmRleCI6IDExfQ==",
+            }
+        }
+
+
+def create_app_from_config(config: GatewayConfig) -> GatewayAPI:
+    """
+    Create the GatewayAPI app from the gateway configuration.
+    """
+    limiter = Limiter(
+        key_func=get_remote_address, storage_uri=MLFLOW_GATEWAY_RATE_LIMITS_STORAGE_URI.get()
+    )
+    app = GatewayAPI(
+        config=config,
+        limiter=limiter,
+        title="MLflow AI Gateway",
+        description="The core deployments API for reverse proxy interface using remote inference "
+        "endpoints within MLflow",
+        version=VERSION,
+        docs_url=None,
+    )
+
+    @app.get("/", include_in_schema=False)
+    async def index():
+        return RedirectResponse(url="/docs")
+
+    @app.get("/favicon.ico", include_in_schema=False)
+    async def favicon():
+        for directory in ["build", "public"]:
+            favicon_file = Path(__file__).parent.parent.joinpath(
+                "server", "js", directory, "favicon.ico"
+            )
+            if favicon_file.exists():
+                return FileResponse(favicon_file)
+        raise HTTPException(status_code=404, detail="favicon.ico not found")
+
+    @app.get("/docs", include_in_schema=False)
+    async def docs():
+        return get_swagger_ui_html(
+            openapi_url="/openapi.json",
+            title="MLflow AI Gateway",
+            swagger_favicon_url="/favicon.ico",
+        )
+
+    # TODO: Remove deployments server URLs after deprecation window elapses
+    @app.get(MLFLOW_DEPLOYMENTS_HEALTH_ENDPOINT)
+    @app.get(MLFLOW_GATEWAY_HEALTH_ENDPOINT, include_in_schema=False)
+    async def health() -> HealthResponse:
+        return {"status": "OK"}
+
+    # TODO: Remove deployments server URLs after deprecation window elapses
+    @app.get(MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE + "{endpoint_name}")
+    async def get_endpoint(endpoint_name: str) -> Endpoint:
+        if matched := app.get_dynamic_route(endpoint_name):
+            return matched.to_endpoint()
+
+        raise HTTPException(
+            status_code=404,
+            detail=f"The endpoint '{endpoint_name}' is not present or active on the server. Please "
+            "verify the endpoint name.",
+        )
+
+    @app.get(MLFLOW_GATEWAY_CRUD_ROUTE_BASE + "{route_name}", include_in_schema=False)
+    async def get_route(route_name: str) -> Route:
+        if matched := app.get_dynamic_route(route_name):
+            return matched
+
+        raise HTTPException(
+            status_code=404,
+            detail=f"The route '{route_name}' is not present or active on the server. Please "
+            "verify the route name.",
+        )
+
+    # TODO: Remove deployments server URLs after deprecation window elapses
+    @app.get(MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE)
+    async def list_endpoints(page_token: Optional[str] = None) -> ListEndpointsResponse:
+        start_idx = SearchRoutesToken.decode(page_token).index if page_token is not None else 0
+
+        end_idx = start_idx + MLFLOW_DEPLOYMENTS_LIST_ENDPOINTS_PAGE_SIZE
+        routes = list(app.dynamic_routes.values())
+        result = {
+            "endpoints": [route.to_route().to_endpoint() for route in routes[start_idx:end_idx]]
+        }
+        if len(routes[end_idx:]) > 0:
+            next_page_token = SearchRoutesToken(index=end_idx)
+            result["next_page_token"] = next_page_token.encode()
+
+        return result
+
+    @app.get(MLFLOW_GATEWAY_CRUD_ROUTE_BASE, include_in_schema=False)
+    async def search_routes(page_token: Optional[str] = None) -> SearchRoutesResponse:
+        start_idx = SearchRoutesToken.decode(page_token).index if page_token is not None else 0
+
+        end_idx = start_idx + MLFLOW_GATEWAY_SEARCH_ROUTES_PAGE_SIZE
+        routes = list(app.dynamic_routes.values())
+        result = {"routes": [r.to_route() for r in routes[start_idx:end_idx]]}
+        if len(routes[end_idx:]) > 0:
+            next_page_token = SearchRoutesToken(index=end_idx)
+            result["next_page_token"] = next_page_token.encode()
+
+        return result
+
+    # TODO: Remove deployments server URLs after deprecation window elapses
+    @app.get(MLFLOW_DEPLOYMENTS_LIMITS_BASE + "{endpoint}")
+    @app.get(MLFLOW_GATEWAY_LIMITS_BASE + "{endpoint}", include_in_schema=False)
+    async def get_limits(endpoint: str) -> LimitsConfig:
+        raise HTTPException(status_code=501, detail="The get_limits API is not available yet.")
+
+    # TODO: Remove deployments server URLs after deprecation window elapses
+    @app.post(MLFLOW_DEPLOYMENTS_LIMITS_BASE)
+    @app.post(MLFLOW_GATEWAY_LIMITS_BASE, include_in_schema=False)
+    async def set_limits(payload: SetLimitsModel) -> LimitsConfig:
+        raise HTTPException(status_code=501, detail="The set_limits API is not available yet.")
+
+    def _look_up_route(name: str) -> Optional[Route]:
+        if r := app.dynamic_routes.get(name):
+            return r
+
+        raise HTTPException(
+            status_code=400,
+            detail=f"Route {name} not found in the configuration.",
+        )
+
+    @app.post("/v1/chat/completions")
+    async def openai_chat_handler(
+        request: Request, payload: chat.RequestPayload
+    ) -> chat.ResponsePayload:
+        route = _look_up_route(payload.model)
+        if route.route_type != RouteType.LLM_V1_CHAT:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Endpoint {route.name!r} is not a chat endpoint.",
+            )
+
+        prov = get_provider(route.model.provider)(route)
+        payload.model = None  # provider rejects a request with model field, must be set to None
+        if payload.stream:
+            return await make_streaming_response(prov.chat_stream(payload))
+        else:
+            return await prov.chat(payload)
+
+    @app.post("/v1/completions")
+    async def openai_completions_handler(
+        request: Request, payload: completions.RequestPayload
+    ) -> completions.ResponsePayload:
+        route = _look_up_route(payload.model)
+        if route.route_type != RouteType.LLM_V1_COMPLETIONS:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Endpoint {route.name!r} is not a completions endpoint.",
+            )
+
+        prov = get_provider(route.model.provider)(route)
+        payload.model = None  # provider rejects a request with model field, must be set to None
+        if payload.stream:
+            return await make_streaming_response(prov.completions_stream(payload))
+        else:
+            return await prov.completions(payload)
+
+    @app.post("/v1/embeddings")
+    async def openai_embeddings_handler(
+        request: Request, payload: embeddings.RequestPayload
+    ) -> embeddings.ResponsePayload:
+        route = _look_up_route(payload.model)
+        if route.route_type != RouteType.LLM_V1_EMBEDDINGS:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Endpoint {route.name!r} is not an embeddings endpoint.",
+            )
+
+        prov = get_provider(route.model.provider)(route)
+        payload.model = None  # provider rejects a request with model field, must be set to None
+        return await prov.embeddings(payload)
+
+    return app
+
+
+def create_app_from_path(config_path: Union[str, Path]) -> GatewayAPI:
+    """
+    Load the path and generate the GatewayAPI app instance.
+    """
+    config = _load_route_config(config_path)
+    return create_app_from_config(config)
+
+
+def create_app_from_env() -> GatewayAPI:
+    """
+    Load the path from the environment variable and generate the GatewayAPI app instance.
+    """
+    if config_path := MLFLOW_GATEWAY_CONFIG.get():
+        return create_app_from_path(config_path)
+
+    raise MlflowException(
+        f"Environment variable {MLFLOW_GATEWAY_CONFIG!r} is not set. "
+        "Please set it to the path of the gateway configuration file."
+    )

venv/lib/python3.9/site-packages/mlflow/gateway/base_models.py

@@ -0,0 +1,60 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class RequestModel(
+    BaseModel,
+    # Allow extra fields for pydantic request models, e.g. to support
+    # vendor-specific embeddings parameters
+    extra="allow",
+):
+    """
+    A pydantic model representing Gateway request data, such as a chat or completions request
+    """
+
+
+class ResponseModel(
+    BaseModel,
+    # Ignore extra fields for pydantic response models to ensure a consistent response
+    # experience for clients across different backends
+    extra="ignore",
+):
+    """
+    A pydantic model representing Gateway response data, such as information about a Gateway
+    Route returned in response to a GetRoute request
+    """
+
+
+class ConfigModel(
+    BaseModel,
+    # Ignore extra fields for pydantic config models, since they are unused
+    extra="ignore",
+):
+    """
+    A pydantic model representing Gateway configuration data, such as an OpenAI completions
+    route definition including route name, model name, API keys, etc.
+    """
+
+
+class LimitModel(
+    BaseModel,
+    # Ignore extra fields for pydantic limit models, since they are unused
+    extra="ignore",
+):
+    """
+    A pydantic model representing Gateway Limit data, such as renewal period, limit
+    key, limit value, etc.
+    """
+
+
+class SetLimitsModel(
+    BaseModel,
+    # Ignore extra fields for pydantic limit models, since they are unused
+    extra="ignore",
+):
+    route: str
+    limits: list[dict[str, Any]]
+    """
+    A pydantic model representing Gateway SetLimits request body, containing route and limits.
+    """

venv/lib/python3.9/site-packages/mlflow/gateway/cli.py

@@ -0,0 +1,48 @@
+import click
+
+from mlflow.environment_variables import MLFLOW_GATEWAY_CONFIG
+from mlflow.gateway.config import _validate_config
+from mlflow.gateway.runner import run_app
+from mlflow.utils.os import is_windows
+
+
+def validate_config_path(_ctx, _param, value):
+    try:
+        _validate_config(value)
+        return value
+    except Exception as e:
+        raise click.BadParameter(str(e))
+
+
+@click.group("gateway", help="Manage the MLflow Gateway service")
+def commands():
+    pass
+
+
+@commands.command("start", help="Start the MLflow Gateway service")
+@click.option(
+    "--config-path",
+    envvar=MLFLOW_GATEWAY_CONFIG.name,
+    callback=validate_config_path,
+    required=True,
+    help="The path to the gateway configuration file.",
+)
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    help="The network address to listen on (default: 127.0.0.1).",
+)
+@click.option(
+    "--port",
+    default=5000,
+    help="The port to listen on (default: 5000).",
+)
+@click.option(
+    "--workers",
+    default=2,
+    help="The number of workers.",
+)
+def start(config_path: str, host: str, port: str, workers: int):
+    if is_windows():
+        raise click.ClickException("MLflow AI Gateway does not support Windows.")
+    run_app(config_path=config_path, host=host, port=port, workers=workers)

venv/lib/python3.9/site-packages/mlflow/gateway/client.py

@@ -0,0 +1,456 @@
+import json
+import logging
+from typing import Any, Optional
+
+import requests.exceptions
+
+from mlflow import MlflowException
+from mlflow.gateway.config import LimitsConfig, Route
+from mlflow.gateway.constants import (
+    MLFLOW_GATEWAY_CLIENT_QUERY_RETRY_CODES,
+    MLFLOW_GATEWAY_CLIENT_QUERY_TIMEOUT_SECONDS,
+    MLFLOW_GATEWAY_CRUD_ROUTE_BASE,
+    MLFLOW_GATEWAY_LIMITS_BASE,
+    MLFLOW_GATEWAY_ROUTE_BASE,
+    MLFLOW_QUERY_SUFFIX,
+)
+from mlflow.gateway.utils import (
+    assemble_uri_path,
+    gateway_deprecated,
+    get_gateway_uri,
+    resolve_route_url,
+)
+from mlflow.protos.databricks_pb2 import BAD_REQUEST
+from mlflow.store.entities.paged_list import PagedList
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.rest_utils import augmented_raise_for_status, http_request
+from mlflow.utils.uri import get_uri_scheme
+
+_logger = logging.getLogger(__name__)
+
+
+@gateway_deprecated
+class MlflowGatewayClient:
+    """
+    Client for interacting with the MLflow Gateway API.
+
+    Args:
+        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`.
+    """
+
+    def __init__(self, gateway_uri: Optional[str] = None):
+        self._gateway_uri = gateway_uri or get_gateway_uri()
+
+    def _is_databricks_host(self) -> bool:
+        return (
+            self._gateway_uri == "databricks" or get_uri_scheme(self._gateway_uri) == "databricks"
+        )
+
+    @property
+    def _host_creds(self):
+        """
+        NB: When `MlflowGatewayClient` is used as an instance variable in a custom pyfunc model, it
+        is pickled in the environment where the custom pyfunc model is defined (e.g. a notebook).
+        When the model is moved to a different environment, e.g. model serving, new credentials
+        need to be resolved from within the new environment. Accordingly, we re-resolve host
+        credentials every time a request is made.
+        """
+        if self._is_databricks_host():
+            return get_databricks_host_creds(self._gateway_uri)
+        else:
+            return get_default_host_creds(self._gateway_uri)
+
+    @property
+    def gateway_uri(self):
+        """
+        Get the current value for the URI of the MLflow Gateway.
+
+        Returns:
+            The gateway URI.
+        """
+        return self._gateway_uri
+
+    def _call_endpoint(self, method: str, route: str, json_body: Optional[str] = None):
+        """
+        Call a specific endpoint on the Gateway API.
+
+        Args:
+            method: The HTTP method to use.
+            route: The API route to call.
+            json_body: Optional JSON body to include in the request.
+
+        Returns:
+            The server's response.
+
+        """
+        if json_body:
+            json_body = json.loads(json_body)
+
+        call_kwargs = {}
+        if method.lower() == "get":
+            call_kwargs["params"] = json_body
+        else:
+            call_kwargs["json"] = json_body
+
+        response = http_request(
+            host_creds=self._host_creds,
+            endpoint=route,
+            method=method,
+            timeout=MLFLOW_GATEWAY_CLIENT_QUERY_TIMEOUT_SECONDS,
+            retry_codes=MLFLOW_GATEWAY_CLIENT_QUERY_RETRY_CODES,
+            raise_on_status=False,
+            **call_kwargs,
+        )
+        augmented_raise_for_status(response)
+        return response
+
+    @gateway_deprecated
+    def get_route(self, name: str):
+        """
+        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).
+
+        Args:
+            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.
+
+        """
+        route = assemble_uri_path([MLFLOW_GATEWAY_CRUD_ROUTE_BASE, name])
+        response = self._call_endpoint("GET", route).json()
+        response["route_url"] = resolve_route_url(self._gateway_uri, response["route_url"])
+
+        return Route(**response)
+
+    @gateway_deprecated
+    def search_routes(self, page_token: Optional[str] = None) -> PagedList[Route]:
+        """
+        Search for routes in the Gateway.
+
+        Args:
+            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.
+
+        """
+        request_parameters = {"page_token": page_token} if page_token is not None else None
+        response_json = self._call_endpoint(
+            "GET", MLFLOW_GATEWAY_CRUD_ROUTE_BASE, json_body=json.dumps(request_parameters)
+        ).json()
+        routes = [
+            Route(
+                **{
+                    **resp,
+                    "route_url": resolve_route_url(
+                        self._gateway_uri,
+                        resp["route_url"],
+                    ),
+                }
+            )
+            for resp in response_json.get("routes", [])
+        ]
+        next_page_token = response_json.get("next_page_token")
+        return PagedList(routes, next_page_token)
+
+    @gateway_deprecated
+    def create_route(
+        self, name: str, route_type: Optional[str] = None, model: Optional[dict[str, Any]] = None
+    ) -> Route:
+        """
+        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.
+
+        Args:
+            name: The name of the route. This parameter is required for all routes.
+            route_type: The type of the route (e.g., 'llm/v1/chat', 'llm/v1/completions',
+                'llm/v1/embeddings'). This parameter is required for routes that are not managed by
+                Databricks (the provider isn't 'databricks').
+            model: A dictionary representing the model details to be associated with the route.
+                This parameter is required for all routes. This dictionary should define:
+
+                    - The model name (e.g., "gpt-4o-mini")
+                    - 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:
+
+        .. code-block:: python
+
+            from mlflow.gateway import MlflowGatewayClient
+
+            gateway_client = MlflowGatewayClient("databricks")
+
+            openai_api_key = ...
+
+            new_route = gateway_client.create_route(
+                name="my-route",
+                route_type="llm/v1/completions",
+                model={
+                    "name": "question-answering-bot",
+                    "provider": "openai",
+                    "openai_config": {
+                        "openai_api_key": openai_api_key,
+                    },
+                },
+            )
+        """
+        if not self._is_databricks_host():
+            raise MlflowException(
+                "The create_route API is only available when running within "
+                "Databricks. Route creation is handled through creating a "
+                "configuration YAML file during startup or through updating a "
+                "running Gateway server.",
+                error_code=BAD_REQUEST,
+            )
+        payload = {
+            "name": name,
+            "route_type": route_type,
+            "model": model,
+        }
+        response = self._call_endpoint(
+            "POST", MLFLOW_GATEWAY_CRUD_ROUTE_BASE, json.dumps(payload)
+        ).json()
+        return Route(**response)
+
+    @gateway_deprecated
+    def delete_route(self, name: str) -> None:
+        """
+        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.
+
+        Args:
+            name: The name of the route to delete.
+
+        Raises:
+            mlflow.MlflowException: If the function is not running within Databricks.
+
+        Example usage from within Databricks:
+
+        .. code-block:: python
+
+            from mlflow.gateway import MlflowGatewayClient
+
+            gateway_client = MlflowGatewayClient("databricks")
+            gateway_client.delete_route("my-existing-route")
+        """
+        if not self._is_databricks_host():
+            raise MlflowException(
+                "The delete_route API is only available when running within Databricks. Route "
+                "deletion is handled through uploading a modified configuration YAML file to the "
+                "location specified when starting the Gateway server. To delete a route, remove "
+                "the route entry from the configuration file.",
+                error_code=BAD_REQUEST,
+            )
+        route = assemble_uri_path([MLFLOW_GATEWAY_CRUD_ROUTE_BASE, name])
+        self._call_endpoint("DELETE", route)
+
+    @gateway_deprecated
+    def query(self, route: str, data: dict[str, Any]):
+        """
+        Submit a query to a configured provider route.
+
+        Args:
+            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.
+
+                For chat, the structure should be:
+
+                .. code-block:: python
+
+                    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:
+
+                .. code-block:: python
+
+                    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:
+
+                .. code-block:: python
+
+                    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"]},
+                    )
+
+                Additional parameters that are valid for a given provider and route configuration
+                can be included with the request as shown below, using an openai completions route
+                request as an example:
+
+                .. code-block:: python
+
+                    from mlflow.gateway import MlflowGatewayClient
+
+                    gateway_client = MlflowGatewayClient("http://my.gateway:8888")
+
+                    response = gateway_client.query(
+                        "my-completions-route",
+                        {
+                            "prompt": "Give me an example of a properly formatted pytest unit test",
+                            "temperature": 0.3,
+                            "max_tokens": 500,
+                        },
+                    )
+
+        Returns:
+            The route's response as a dictionary, standardized to the route type.
+
+        """
+
+        data = json.dumps(data)
+
+        query_route = assemble_uri_path([MLFLOW_GATEWAY_ROUTE_BASE, route, MLFLOW_QUERY_SUFFIX])
+
+        try:
+            return self._call_endpoint("POST", query_route, data).json()
+        except MlflowException as e:
+            if isinstance(e.__cause__, requests.exceptions.Timeout):
+                timeout_message = (
+                    "The provider has timed out while generating a response to your "
+                    "query. Please evaluate the available parameters for the query "
+                    "that you are submitting. Some parameter values and inputs can "
+                    "increase the computation time beyond the allowable route "
+                    f"timeout of {MLFLOW_GATEWAY_CLIENT_QUERY_TIMEOUT_SECONDS} "
+                    "seconds."
+                )
+                raise MlflowException(message=timeout_message, error_code=BAD_REQUEST)
+            else:
+                raise e
+
+    @gateway_deprecated
+    def set_limits(self, route: str, limits: list[dict[str, Any]]) -> LimitsConfig:
+        """
+        Set limits on an existing route in the Gateway.
+
+        .. warning::
+
+            This API is **only available** when running within Databricks.
+
+        Args:
+            route: The name of the route to set limits on.
+            limits: Limits (Array of dictionary) to set on the route. Each limit is defined by a
+                dictionary representing the limit details to be associated with the route. This
+                dictionary should define:
+
+                - renewal_period: a string representing the length of the window to enforce limit
+                  on (only supports "minute" for now).
+                - calls: a non-negative integer representing the number of calls allowed per
+                  renewal_period (e.g., 10, 0, 55).
+                - key: an optional string represents per route limit or per user limit ("user" for
+                  per user limit, "route" for per route limit, if not supplied, default to per
+                  route limit).
+
+        Returns:
+            The returned data structure is a serialized representation of the `Limit`
+            data structure, giving information about the renewal_period, key, and calls.
+
+        Example usage:
+
+        .. code-block:: python
+
+            from mlflow.gateway import MlflowGatewayClient
+
+            gateway_client = MlflowGatewayClient("databricks")
+
+            gateway_client.set_limits(
+                "my-new-route", [{"key": "user", "renewal_period": "minute", "calls": 50}]
+            )
+        """
+        payload = {
+            "route": route,
+            "limits": limits,
+        }
+
+        response = self._call_endpoint(
+            "POST", MLFLOW_GATEWAY_LIMITS_BASE, json.dumps(payload)
+        ).json()
+        return LimitsConfig(**response)
+
+    @gateway_deprecated
+    def get_limits(self, route: str) -> LimitsConfig:
+        """
+        Get limits of an existing route in the Gateway.
+
+        .. warning::
+
+            This API is **only available** when connected to a Databricks-hosted AI Gateway.
+
+        Args:
+            route: The name of the route to get limits of.
+
+        Returns:
+            The returned data structure is a serialized representation of the `Limit` data
+            structure, giving information about the renewal_period, key, and calls.
+
+        Example usage:
+
+        .. code-block:: python
+
+            from mlflow.gateway import MlflowGatewayClient
+
+            gateway_client = MlflowGatewayClient("databricks")
+
+            gateway_client.get_limits("my-new-route")
+        """
+        if not route:
+            raise MlflowException("A non-empty string is required for the route.", BAD_REQUEST)
+        route_uri = assemble_uri_path([MLFLOW_GATEWAY_LIMITS_BASE, route])
+        response = self._call_endpoint("GET", route_uri).json()
+        return LimitsConfig(**response)

venv/lib/python3.9/site-packages/mlflow/gateway/config.py

@@ -0,0 +1,519 @@
+import json
+import logging
+import os
+import pathlib
+from enum import Enum
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import pydantic
+import yaml
+from packaging.version import Version
+from pydantic import ConfigDict, Field, ValidationError
+from pydantic.json import pydantic_encoder
+
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.base_models import ConfigModel, LimitModel, ResponseModel
+from mlflow.gateway.constants import (
+    MLFLOW_AI_GATEWAY_MOSAICML_CHAT_SUPPORTED_MODEL_PREFIXES,
+    MLFLOW_GATEWAY_ROUTE_BASE,
+    MLFLOW_QUERY_SUFFIX,
+)
+from mlflow.gateway.utils import (
+    check_configuration_deprecated_fields,
+    check_configuration_route_name_collisions,
+    is_valid_ai21labs_model,
+    is_valid_endpoint_name,
+    is_valid_mosiacml_chat_model,
+)
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER, field_validator, model_validator
+
+_logger = logging.getLogger(__name__)
+
+if IS_PYDANTIC_V2_OR_NEWER:
+    from pydantic import SerializeAsAny
+
+
+class Provider(str, Enum):
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    COHERE = "cohere"
+    AI21LABS = "ai21labs"
+    MLFLOW_MODEL_SERVING = "mlflow-model-serving"
+    MOSAICML = "mosaicml"
+    HUGGINGFACE_TEXT_GENERATION_INFERENCE = "huggingface-text-generation-inference"
+    PALM = "palm"
+    GEMINI = "gemini"
+    BEDROCK = "bedrock"
+    AMAZON_BEDROCK = "amazon-bedrock"  # an alias for bedrock
+    # Note: The following providers are only supported on Databricks
+    DATABRICKS_MODEL_SERVING = "databricks-model-serving"
+    DATABRICKS = "databricks"
+    MISTRAL = "mistral"
+    TOGETHERAI = "togetherai"
+
+    @classmethod
+    def values(cls):
+        return {p.value for p in cls}
+
+
+class TogetherAIConfig(ConfigModel):
+    togetherai_api_key: str
+
+    @field_validator("togetherai_api_key", mode="before")
+    def validate_togetherai_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class RouteType(str, Enum):
+    LLM_V1_COMPLETIONS = "llm/v1/completions"
+    LLM_V1_CHAT = "llm/v1/chat"
+    LLM_V1_EMBEDDINGS = "llm/v1/embeddings"
+
+
+class CohereConfig(ConfigModel):
+    cohere_api_key: str
+
+    @field_validator("cohere_api_key", mode="before")
+    def validate_cohere_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class AI21LabsConfig(ConfigModel):
+    ai21labs_api_key: str
+
+    @field_validator("ai21labs_api_key", mode="before")
+    def validate_ai21labs_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class MosaicMLConfig(ConfigModel):
+    mosaicml_api_key: str
+    mosaicml_api_base: Optional[str] = None
+
+    @field_validator("mosaicml_api_key", mode="before")
+    def validate_mosaicml_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class OpenAIAPIType(str, Enum):
+    OPENAI = "openai"
+    AZURE = "azure"
+    AZUREAD = "azuread"
+
+    @classmethod
+    def _missing_(cls, value):
+        """
+        Implements case-insensitive matching of API type strings
+        """
+        for api_type in cls:
+            if api_type.value == value.lower():
+                return api_type
+
+        raise MlflowException.invalid_parameter_value(f"Invalid OpenAI API type '{value}'")
+
+
+class OpenAIConfig(ConfigModel):
+    openai_api_key: str
+    openai_api_type: 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
+
+    @field_validator("openai_api_key", mode="before")
+    def validate_openai_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+    @classmethod
+    def _validate_field_compatibility(cls, info: dict[str, Any]):
+        if not isinstance(info, dict):
+            return info
+        api_type = (info.get("openai_api_type") or OpenAIAPIType.OPENAI).lower()
+        if api_type == OpenAIAPIType.OPENAI:
+            if info.get("openai_deployment_name") is not None:
+                raise MlflowException.invalid_parameter_value(
+                    f"OpenAI route configuration can only specify a value for "
+                    f"'openai_deployment_name' if 'openai_api_type' is '{OpenAIAPIType.AZURE}' "
+                    f"or '{OpenAIAPIType.AZUREAD}'. Found type: '{api_type}'"
+                )
+            if info.get("openai_api_base") is None:
+                info["openai_api_base"] = "https://api.openai.com/v1"
+        elif api_type in (OpenAIAPIType.AZURE, OpenAIAPIType.AZUREAD):
+            if info.get("openai_organization") is not None:
+                raise MlflowException.invalid_parameter_value(
+                    f"OpenAI route configuration can only specify a value for "
+                    f"'openai_organization' if 'openai_api_type' is '{OpenAIAPIType.OPENAI}'"
+                )
+            base_url = info.get("openai_api_base")
+            deployment_name = info.get("openai_deployment_name")
+            api_version = info.get("openai_api_version")
+            if (base_url, deployment_name, api_version).count(None) > 0:
+                raise MlflowException.invalid_parameter_value(
+                    f"OpenAI route configuration must specify 'openai_api_base', "
+                    f"'openai_deployment_name', and 'openai_api_version' if 'openai_api_type' is "
+                    f"'{OpenAIAPIType.AZURE}' or '{OpenAIAPIType.AZUREAD}'."
+                )
+        else:
+            raise MlflowException.invalid_parameter_value(f"Invalid OpenAI API type '{api_type}'")
+
+        return info
+
+    @model_validator(mode="before")
+    def validate_field_compatibility(cls, info: dict[str, Any]):
+        return cls._validate_field_compatibility(info)
+
+
+class AnthropicConfig(ConfigModel):
+    anthropic_api_key: str
+    anthropic_version: str = "2023-06-01"
+
+    @field_validator("anthropic_api_key", mode="before")
+    def validate_anthropic_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class PaLMConfig(ConfigModel):
+    palm_api_key: str
+
+    @field_validator("palm_api_key", mode="before")
+    def validate_palm_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class GeminiConfig(ConfigModel):
+    gemini_api_key: str
+
+    @field_validator("gemini_api_key", mode="before")
+    def validate_gemini_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class MlflowModelServingConfig(ConfigModel):
+    model_server_url: str
+
+    # Workaround to suppress warning that Pydantic raises when a field name starts with "model_".
+    # https://github.com/mlflow/mlflow/issues/10335
+    model_config = pydantic.ConfigDict(protected_namespaces=())
+
+
+class HuggingFaceTextGenerationInferenceConfig(ConfigModel):
+    hf_server_url: str
+
+
+class AWSBaseConfig(pydantic.BaseModel):
+    aws_region: Optional[str] = None
+
+
+class AWSRole(AWSBaseConfig):
+    aws_role_arn: str
+    session_length_seconds: int = 15 * 60
+
+
+class AWSIdAndKey(AWSBaseConfig):
+    aws_access_key_id: str
+    aws_secret_access_key: str
+    aws_session_token: Optional[str] = None
+
+
+class AmazonBedrockConfig(ConfigModel):
+    # order here is important, at least for pydantic<2
+    aws_config: Union[AWSRole, AWSIdAndKey, AWSBaseConfig]
+
+
+class MistralConfig(ConfigModel):
+    mistral_api_key: str
+
+    @field_validator("mistral_api_key", mode="before")
+    def validate_mistral_api_key(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
+class ModelInfo(ResponseModel):
+    name: Optional[str] = None
+    provider: Provider
+
+
+def _resolve_api_key_from_input(api_key_input):
+    """
+    Resolves the provided API key.
+
+    Input formats accepted:
+
+    - Path to a file as a string which will have the key loaded from it
+    - environment variable name that stores the api key
+    - the api key itself
+    """
+
+    if not isinstance(api_key_input, str):
+        raise MlflowException.invalid_parameter_value(
+            "The api key provided is not a string. Please provide either an environment "
+            "variable key, a path to a file containing the api key, or the api key itself"
+        )
+
+    # try reading as an environment variable
+    if api_key_input.startswith("$"):
+        env_var_name = api_key_input[1:]
+        if env_var := os.getenv(env_var_name):
+            return env_var
+        else:
+            raise MlflowException.invalid_parameter_value(
+                f"Environment variable {env_var_name!r} is not set"
+            )
+
+    # try reading from a local path
+    file = pathlib.Path(api_key_input)
+    try:
+        if file.is_file():
+            return file.read_text()
+    except OSError:
+        # `is_file` throws an OSError if `api_key_input` exceeds the maximum filename length
+        # (e.g., 255 characters on Unix).
+        pass
+
+    # if the key itself is passed, return
+    return api_key_input
+
+
+class Model(ConfigModel):
+    name: Optional[str] = None
+    provider: Union[str, Provider]
+    if IS_PYDANTIC_V2_OR_NEWER:
+        config: Optional[SerializeAsAny[ConfigModel]] = None
+    else:
+        config: Optional[ConfigModel] = None
+
+    @field_validator("provider", mode="before")
+    def validate_provider(cls, value):
+        from mlflow.gateway.provider_registry import provider_registry
+
+        if isinstance(value, Provider):
+            return value
+        formatted_value = value.replace("-", "_").upper()
+        if formatted_value in Provider.__members__:
+            return Provider[formatted_value]
+        if value in provider_registry.keys():
+            return value
+        raise MlflowException.invalid_parameter_value(f"The provider '{value}' is not supported.")
+
+    @classmethod
+    def _validate_config(cls, val, context):
+        from mlflow.gateway.provider_registry import provider_registry
+
+        # For Pydantic v2: 'context' is a ValidationInfo object with a 'data' attribute.
+        # For Pydantic v1: 'context' is dict-like 'values'.
+        if IS_PYDANTIC_V2_OR_NEWER:
+            provider = context.data.get("provider")
+        else:
+            provider = context.get("provider") if context else None
+
+        if provider:
+            config_type = provider_registry.get(provider).CONFIG_TYPE
+            return config_type(**val) if isinstance(val, dict) else val
+        raise MlflowException.invalid_parameter_value(
+            "A provider must be provided for each gateway route."
+        )
+
+    @field_validator("config", mode="before")
+    def validate_config(cls, info, values):
+        return cls._validate_config(info, values)
+
+
+class AliasedConfigModel(ConfigModel):
+    """
+    Enables use of field aliases in a configuration model for backwards compatibility
+    """
+
+    if Version(pydantic.__version__) >= Version("2.0"):
+        model_config = ConfigDict(populate_by_name=True)
+    else:
+
+        class Config:
+            allow_population_by_field_name = True
+
+
+class Limit(LimitModel):
+    calls: int
+    key: Optional[str] = None
+    renewal_period: str
+
+
+class LimitsConfig(ConfigModel):
+    limits: Optional[list[Limit]] = []
+
+
+class RouteConfig(AliasedConfigModel):
+    name: str
+    route_type: RouteType = Field(alias="endpoint_type")
+    model: Model
+    limit: Optional[Limit] = None
+
+    @field_validator("name")
+    def validate_endpoint_name(cls, route_name):
+        if not is_valid_endpoint_name(route_name):
+            raise MlflowException.invalid_parameter_value(
+                "The route name provided contains disallowed characters for a url endpoint. "
+                f"'{route_name}' is invalid. Names cannot contain spaces or any non "
+                "alphanumeric characters other than hyphen and underscore."
+            )
+        return route_name
+
+    @field_validator("model", mode="before")
+    def validate_model(cls, model):
+        if model:
+            model_instance = Model(**model)
+            if model_instance.provider in Provider.values() and model_instance.config is None:
+                raise MlflowException.invalid_parameter_value(
+                    "A config must be supplied when setting a provider. The provider entry for "
+                    f"{model_instance.provider} is incorrect."
+                )
+        return model
+
+    @model_validator(mode="after", skip_on_failure=True)
+    def validate_route_type_and_model_name(cls, values):
+        if IS_PYDANTIC_V2_OR_NEWER:
+            route_type = values.route_type
+            model = values.model
+        else:
+            route_type = values.get("route_type")
+            model = values.get("model")
+        if (
+            model
+            and model.provider == "mosaicml"
+            and route_type == RouteType.LLM_V1_CHAT
+            and not is_valid_mosiacml_chat_model(model.name)
+        ):
+            raise MlflowException.invalid_parameter_value(
+                f"An invalid model has been specified for the chat route. '{model.name}'. "
+                f"Ensure the model selected starts with one of: "
+                f"{MLFLOW_AI_GATEWAY_MOSAICML_CHAT_SUPPORTED_MODEL_PREFIXES}"
+            )
+        if model and model.provider == "ai21labs" and not is_valid_ai21labs_model(model.name):
+            raise MlflowException.invalid_parameter_value(
+                f"An Unsupported AI21Labs model has been specified: '{model.name}'. "
+                f"Please see documentation for supported models."
+            )
+        return values
+
+    @field_validator("route_type", mode="before")
+    def validate_route_type(cls, value):
+        if value in RouteType._value2member_map_:
+            return value
+        raise MlflowException.invalid_parameter_value(f"The route_type '{value}' is not supported.")
+
+    @field_validator("limit", mode="before")
+    def validate_limit(cls, value):
+        from limits import parse
+
+        if value:
+            limit = Limit(**value)
+            try:
+                parse(f"{limit.calls}/{limit.renewal_period}")
+            except ValueError:
+                raise MlflowException.invalid_parameter_value(
+                    "Failed to parse the rate limit configuration."
+                    "Please make sure limit.calls is a positive number and"
+                    "limit.renewal_period is a right granularity"
+                )
+
+        return value
+
+    def to_route(self) -> "Route":
+        return Route(
+            name=self.name,
+            route_type=self.route_type,
+            model=RouteModelInfo(
+                name=self.model.name,
+                provider=self.model.provider,
+            ),
+            route_url=f"{MLFLOW_GATEWAY_ROUTE_BASE}{self.name}{MLFLOW_QUERY_SUFFIX}",
+            limit=self.limit,
+        )
+
+
+class RouteModelInfo(ResponseModel):
+    name: Optional[str] = None
+    # Use `str` instead of `Provider` enum to allow gateway backends such as Databricks to
+    # support new providers without breaking the gateway client.
+    provider: str
+
+
+_ROUTE_EXTRA_SCHEMA = {
+    "example": {
+        "name": "openai-completions",
+        "route_type": "llm/v1/completions",
+        "model": {
+            "name": "gpt-4o-mini",
+            "provider": "openai",
+        },
+        "route_url": "/gateway/routes/completions/invocations",
+    }
+}
+
+
+class Route(ConfigModel):
+    name: str
+    route_type: str
+    model: RouteModelInfo
+    route_url: str
+    limit: Optional[Limit] = None
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _ROUTE_EXTRA_SCHEMA
+        else:
+            schema_extra = _ROUTE_EXTRA_SCHEMA
+
+    def to_endpoint(self):
+        from mlflow.deployments.server.config import Endpoint
+
+        return Endpoint(
+            name=self.name,
+            endpoint_type=self.route_type,
+            model=self.model,
+            endpoint_url=self.route_url,
+            limit=self.limit,
+        )
+
+
+class GatewayConfig(AliasedConfigModel):
+    routes: list[RouteConfig] = Field(alias="endpoints")
+
+
+def _load_route_config(path: Union[str, Path]) -> GatewayConfig:
+    """
+    Reads the gateway configuration yaml file from the storage location and returns an instance
+    of the configuration RouteConfig class
+    """
+    if isinstance(path, str):
+        path = Path(path)
+    try:
+        configuration = yaml.safe_load(path.read_text())
+    except Exception as e:
+        raise MlflowException.invalid_parameter_value(
+            f"The file at {path} is not a valid yaml file"
+        ) from e
+    check_configuration_deprecated_fields(configuration)
+    check_configuration_route_name_collisions(configuration)
+    try:
+        return GatewayConfig(**configuration)
+    except ValidationError as e:
+        raise MlflowException.invalid_parameter_value(
+            f"The gateway configuration is invalid: {e}"
+        ) from e
+
+
+def _save_route_config(config: GatewayConfig, path: Union[str, Path]) -> None:
+    if isinstance(path, str):
+        path = Path(path)
+    path.write_text(yaml.safe_dump(json.loads(json.dumps(config.dict(), default=pydantic_encoder))))
+
+
+def _validate_config(config_path: str) -> GatewayConfig:
+    if not os.path.exists(config_path):
+        raise MlflowException.invalid_parameter_value(f"{config_path} does not exist")
+
+    try:
+        return _load_route_config(config_path)
+    except ValidationError as e:
+        raise MlflowException.invalid_parameter_value(f"Invalid gateway configuration: {e}") from e

venv/lib/python3.9/site-packages/mlflow/gateway/constants.py

@@ -0,0 +1,41 @@
+MLFLOW_GATEWAY_HEALTH_ENDPOINT = "/health"
+MLFLOW_GATEWAY_CRUD_ROUTE_BASE = "/api/2.0/gateway/routes/"
+MLFLOW_GATEWAY_LIMITS_BASE = "/api/2.0/gateway/limits/"
+MLFLOW_GATEWAY_ROUTE_BASE = "/gateway/"
+MLFLOW_QUERY_SUFFIX = "/invocations"
+MLFLOW_GATEWAY_SEARCH_ROUTES_PAGE_SIZE = 3000
+
+# Specifies the timeout for the Gateway server to declare a request submitted to a provider has
+# timed out.
+MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS = 300
+
+# Specifies the timeout for the MLflowGatewayClient APIs to declare a request has timed out
+MLFLOW_GATEWAY_CLIENT_QUERY_TIMEOUT_SECONDS = 300
+
+# Abridged retryable error codes for the interface to the Gateway Server.
+# These are modified from the standard MLflow Tracking server retry codes for the MLflowClient to
+# remove timeouts from the list of the retryable conditions. A long-running timeout with
+# retries for the proxied providers generally indicates an issue with the underlying query or
+# the model being served having issues responding to the query due to parameter configuration.
+MLFLOW_GATEWAY_CLIENT_QUERY_RETRY_CODES = frozenset(
+    [
+        429,  # Too many requests
+        500,  # Server Error
+        502,  # Bad Gateway
+        503,  # Service Unavailable
+    ]
+)
+
+# Provider constants
+MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS = 1_000_000
+MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS = 200_000
+
+# MLflow model serving constants
+MLFLOW_SERVING_RESPONSE_KEY = "predictions"
+
+# MosaicML constants
+# MosaicML supported chat model names
+# These validated names are used for the MosaicML provider due to the need to perform prompt
+# translations prior to sending a request payload to their chat endpoints.
+# to reduce the need to case-match, supported model prefixes are lowercase.
+MLFLOW_AI_GATEWAY_MOSAICML_CHAT_SUPPORTED_MODEL_PREFIXES = ["llama2"]

venv/lib/python3.9/site-packages/mlflow/gateway/exceptions.py

@@ -0,0 +1,14 @@
+class AIGatewayConfigException(Exception):
+    pass
+
+
+class AIGatewayException(Exception):
+    """
+    A custom exception class for handling exceptions raised by the AI Gateway.
+    This will be transformed into an HTTPException before being returned to the client.
+    """
+
+    def __init__(self, status_code: int, detail: str):
+        self.status_code = status_code
+        self.detail = detail
+        super().__init__(detail)

venv/lib/python3.9/site-packages/mlflow/gateway/fluent.py

@@ -0,0 +1,254 @@
+from typing import Any, Optional
+
+from mlflow.gateway.client import MlflowGatewayClient
+from mlflow.gateway.config import LimitsConfig, Route
+from mlflow.gateway.constants import MLFLOW_GATEWAY_SEARCH_ROUTES_PAGE_SIZE
+from mlflow.gateway.utils import gateway_deprecated
+from mlflow.utils import get_results_from_paginated_fn
+
+
+@gateway_deprecated
+def get_route(name: str) -> Route:
+    """
+    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.
+
+    Args:
+        name: The name of the route to fetch.
+
+    Returns:
+        An instance of the Route class representing the fetched route.
+
+    """
+    return MlflowGatewayClient().get_route(name)
+
+
+@gateway_deprecated
+def search_routes() -> list[Route]:
+    """
+    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.
+
+    """
+
+    def pagination_wrapper_func(_, next_page_token):
+        return MlflowGatewayClient().search_routes(page_token=next_page_token)
+
+    return get_results_from_paginated_fn(
+        paginated_fn=pagination_wrapper_func,
+        max_results_per_page=MLFLOW_GATEWAY_SEARCH_ROUTES_PAGE_SIZE,
+        max_results=None,
+    )
+
+
+@gateway_deprecated
+def create_route(
+    name: str, route_type: Optional[str] = None, model: Optional[dict[str, Any]] = None
+) -> Route:
+    """
+    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.
+
+    Args:
+        name: The name of the route. This parameter is required for all routes.
+        route_type: The type of the route (e.g., 'llm/v1/chat', 'llm/v1/completions',
+            'llm/v1/embeddings'). This parameter is required for routes that are
+            not managed by Databricks (the provider isn't 'databricks').
+        model: A dictionary representing the model details to be associated with the route.
+            This parameter is required for all routes. This dictionary should define:
+            - The model name (e.g., "gpt-4o-mini")
+            - 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:
+
+    .. code-block:: python
+
+        from mlflow.gateway import set_gateway_uri, create_route
+
+        set_gateway_uri(gateway_uri="databricks")
+
+        openai_api_key = ...
+
+        create_route(
+            name="my-route",
+            route_type="llm/v1/completions",
+            model={
+                "name": "question-answering-bot",
+                "provider": "openai",
+                "openai_config": {
+                    "openai_api_key": openai_api_key,
+                },
+            },
+        )
+    """
+    return MlflowGatewayClient().create_route(name, route_type, model)
+
+
+@gateway_deprecated
+def delete_route(name: str) -> None:
+    """
+    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.
+
+    Args:
+        name: The name of the route to delete.
+
+    Example usage from within Databricks:
+
+    .. code-block:: python
+
+        from mlflow.gateway import set_gateway_uri, delete_route
+
+        set_gateway_uri(gateway_uri="databricks")
+
+        delete_route("my-new-route")
+    """
+    MlflowGatewayClient().delete_route(name)
+
+
+@gateway_deprecated
+def set_limits(route: str, limits: list[dict[str, Any]]) -> LimitsConfig:
+    """
+    Set limits on an existing route in the Gateway.
+
+    .. warning::
+
+        This API is **only available** when running within Databricks.
+
+    Args:
+        route: The name of the route to set limits on.
+        limits: Limits to set on the route.
+
+    Example usage from within Databricks:
+
+    .. code-block:: python
+
+        from mlflow.gateway import set_gateway_uri, set_limits
+
+        set_gateway_uri(gateway_uri="databricks")
+
+        set_limits("my-new-route", [{"key": "user", "renewal_period": "minute", "calls": 50}])
+    """
+    return MlflowGatewayClient().set_limits(route=route, limits=limits)
+
+
+@gateway_deprecated
+def get_limits(route: str) -> LimitsConfig:
+    """
+    Get limits of an existing route in the Gateway.
+
+    .. warning::
+
+        This API is **only available** when connected to a Databricks-hosted AI Gateway.
+
+    Args:
+        route: The name of the route to get limits of.
+
+    Example usage from within Databricks:
+
+    .. code-block:: python
+
+        from mlflow.gateway import set_gateway_uri, get_limits
+
+        set_gateway_uri(gateway_uri="databricks")
+
+        get_limits("my-new-route")
+    """
+    return MlflowGatewayClient().get_limits(route=route)
+
+
+@gateway_deprecated
+def query(route: str, data):
+    """
+    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.
+
+    Args:
+        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:
+
+    .. code-block:: python
+
+        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:
+
+    .. code-block:: python
+
+        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:
+
+    .. code-block:: python
+
+        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"]}
+        )
+
+    Additional parameters that are valid for a given provider and route configuration can be
+    included with the request as shown below, using an openai completions route request as
+    an example:
+
+    .. code-block:: python
+
+        from mlflow.gateway import query, set_gateway_uri
+
+        set_gateway_uri(gateway_uri="http://my.gateway:9000")
+        response = query(
+            "a_completions_route",
+            {
+                "prompt": "Give me an example of a properly formatted pytest unit test",
+                "temperature": 0.6,
+                "max_tokens": 1000,
+            },
+        )
+    """
+    return MlflowGatewayClient().query(route, data)

venv/lib/python3.9/site-packages/mlflow/gateway/provider_registry.py

@@ -0,0 +1,73 @@
+from typing import Union
+
+from mlflow import MlflowException
+from mlflow.gateway.config import Provider
+from mlflow.gateway.providers import BaseProvider
+from mlflow.utils.plugins import get_entry_points
+
+
+class ProviderRegistry:
+    def __init__(self):
+        self._providers: dict[Union[str, Provider], type[BaseProvider]] = {}
+
+    def register(self, name: str, provider: type[BaseProvider]):
+        if name in self._providers:
+            raise MlflowException.invalid_parameter_value(
+                f"Provider {name} is already registered: {self._providers[name]}"
+            )
+        self._providers[name] = provider
+
+    def get(self, name: str) -> type[BaseProvider]:
+        if name not in self._providers:
+            raise MlflowException.invalid_parameter_value(f"Provider {name} not found")
+        return self._providers[name]
+
+    def keys(self):
+        return list(self._providers.keys())
+
+
+def _register_default_providers(registry: ProviderRegistry):
+    from mlflow.gateway.providers.ai21labs import AI21LabsProvider
+    from mlflow.gateway.providers.anthropic import AnthropicProvider
+    from mlflow.gateway.providers.bedrock import AmazonBedrockProvider
+    from mlflow.gateway.providers.cohere import CohereProvider
+    from mlflow.gateway.providers.gemini import GeminiProvider
+    from mlflow.gateway.providers.huggingface import HFTextGenerationInferenceServerProvider
+    from mlflow.gateway.providers.mistral import MistralProvider
+    from mlflow.gateway.providers.mlflow import MlflowModelServingProvider
+    from mlflow.gateway.providers.mosaicml import MosaicMLProvider
+    from mlflow.gateway.providers.openai import OpenAIProvider
+    from mlflow.gateway.providers.palm import PaLMProvider
+    from mlflow.gateway.providers.togetherai import TogetherAIProvider
+
+    registry.register(Provider.OPENAI, OpenAIProvider)
+    registry.register(Provider.ANTHROPIC, AnthropicProvider)
+    registry.register(Provider.COHERE, CohereProvider)
+    registry.register(Provider.AI21LABS, AI21LabsProvider)
+    registry.register(Provider.MOSAICML, MosaicMLProvider)
+    registry.register(Provider.PALM, PaLMProvider)
+    registry.register(Provider.GEMINI, GeminiProvider)
+    registry.register(Provider.MLFLOW_MODEL_SERVING, MlflowModelServingProvider)
+    registry.register(Provider.BEDROCK, AmazonBedrockProvider)
+    registry.register(Provider.AMAZON_BEDROCK, AmazonBedrockProvider)
+    registry.register(
+        Provider.HUGGINGFACE_TEXT_GENERATION_INFERENCE, HFTextGenerationInferenceServerProvider
+    )
+    registry.register(Provider.MISTRAL, MistralProvider)
+    registry.register(Provider.TOGETHERAI, TogetherAIProvider)
+
+
+def _register_plugin_providers(registry: ProviderRegistry):
+    providers = get_entry_points("mlflow.gateway.providers")
+    for p in providers:
+        cls = p.load()
+        registry.register(p.name, cls)
+
+
+def is_supported_provider(name: str) -> bool:
+    return name in provider_registry.keys()
+
+
+provider_registry = ProviderRegistry()
+_register_default_providers(provider_registry)
+_register_plugin_providers(provider_registry)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/__init__.py

@@ -0,0 +1,8 @@
+from mlflow.gateway.config import Provider
+from mlflow.gateway.providers.base import BaseProvider
+
+
+def get_provider(provider: Provider) -> type[BaseProvider]:
+    from mlflow.gateway.provider_registry import provider_registry
+
+    return provider_registry.get(provider)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/ai21labs.py

@@ -0,0 +1,87 @@
+import time
+
+from mlflow.gateway.config import AI21LabsConfig, RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider
+from mlflow.gateway.providers.utils import rename_payload_keys, send_request
+from mlflow.gateway.schemas import completions
+
+
+class AI21LabsProvider(BaseProvider):
+    NAME = "AI21Labs"
+    CONFIG_TYPE = AI21LabsConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, AI21LabsConfig):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.ai21labs_config: AI21LabsConfig = config.model.config
+        self.headers = {"Authorization": f"Bearer {self.ai21labs_config.ai21labs_api_key}"}
+        self.base_url = f"https://api.ai21.com/studio/v1/{self.config.model.name}/"
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        key_mapping = {
+            "stop": "stopSequences",
+            "n": "numResults",
+            "max_tokens": "maxTokens",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        if payload.get("stream", False):
+            raise AIGatewayException(
+                status_code=422,
+                detail="Setting the 'stream' parameter to 'true' is not supported with the MLflow "
+                "Gateway.",
+            )
+        payload = rename_payload_keys(payload, key_mapping)
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="complete",
+            payload=payload,
+        )
+        # Response example (https://docs.ai21.com/reference/j2-complete-ref)
+        # ```
+        # {
+        #   "id": "7921a78e-d905-c9df-27e3-88e4831e3c3b",
+        #   "prompt": {
+        #     "text": "I will"
+        #   },
+        #   "completions": [
+        #     {
+        #       "data": {
+        #         "text": " complete this"
+        #       },
+        #       "finishReason": {
+        #         "reason": "length",
+        #         "length": 2
+        #       }
+        #     }
+        #   ]
+        # }
+        # ```
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=self.config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=c["data"]["text"],
+                    finish_reason=c["finishReason"]["reason"],
+                )
+                for idx, c in enumerate(resp["completions"])
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )

venv/lib/python3.9/site-packages/mlflow/gateway/providers/anthropic.py

@@ -0,0 +1,357 @@
+import json
+import time
+from typing import AsyncIterable
+
+from mlflow.gateway.config import AnthropicConfig, RouteConfig
+from mlflow.gateway.constants import (
+    MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS,
+    MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS,
+)
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.utils import rename_payload_keys, send_request, send_stream_request
+from mlflow.gateway.schemas import chat, completions
+
+
+class AnthropicAdapter(ProviderAdapter):
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        key_mapping = {"stop": "stop_sequences"}
+        payload["model"] = config.model.name
+        payload = rename_payload_keys(payload, key_mapping)
+
+        if "top_p" in payload and "temperature" in payload:
+            raise AIGatewayException(
+                status_code=422, detail="Cannot set both 'temperature' and 'top_p' parameters."
+            )
+
+        max_tokens = payload.get("max_tokens", MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS)
+        if max_tokens > MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS:
+            raise AIGatewayException(
+                status_code=422,
+                detail="Invalid value for max_tokens: cannot exceed "
+                f"{MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS}.",
+            )
+        payload["max_tokens"] = max_tokens
+
+        if payload.pop("n", 1) != 1:
+            raise AIGatewayException(
+                status_code=422,
+                detail="'n' must be '1' for the Anthropic provider. Received value: '{n}'.",
+            )
+
+        # Cohere uses `system` to set the system message
+        # we concatenate all system messages from the user with a newline
+        system_messages = [m for m in payload["messages"] if m["role"] == "system"]
+        if system_messages:
+            payload["system"] = "\n".join(m["content"] for m in system_messages)
+
+        # remaining messages are chat history
+        # we want to include only user and assistant messages
+        payload["messages"] = [m for m in payload["messages"] if m["role"] in ("user", "assistant")]
+
+        # The range of Anthropic's temperature is 0-1, but ours is 0-2, so we halve it
+        if "temperature" in payload:
+            payload["temperature"] = 0.5 * payload["temperature"]
+
+        return payload
+
+    @classmethod
+    def model_to_chat(cls, resp, config):
+        # API reference: https://docs.anthropic.com/en/api/messages#body-messages
+        #
+        # Example response:
+        # ```
+        # {
+        #   "content": [
+        #     {
+        #       "text": "Blue is often seen as a calming and soothing color.",
+        #       "type": "text"
+        #     },
+        #     {
+        #       "source": {
+        #       "type": "base64",
+        #       "media_type": "image/jpeg",
+        #       "data": "/9j/4AAQSkZJRg...",
+        #       "type": "image",
+        #       }
+        #     }
+        #   ],
+        #   "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF",
+        #   "model": "claude-2.1",
+        #   "role": "assistant",
+        #   "stop_reason": "end_turn",
+        #   "stop_sequence": null,
+        #   "type": "message",
+        #   "usage": {
+        #     "input_tokens": 10,
+        #     "output_tokens": 25
+        #   }
+        # }
+        # ```
+        from mlflow.anthropic.chat import convert_message_to_mlflow_chat
+
+        stop_reason = "length" if resp["stop_reason"] == "max_tokens" else "stop"
+
+        return chat.ResponsePayload(
+            id=resp["id"],
+            created=int(time.time()),
+            object="chat.completion",
+            model=resp["model"],
+            choices=[
+                chat.Choice(
+                    index=0,
+                    # TODO: Remove this casting once
+                    # https://github.com/mlflow/mlflow/pull/14160 is merged
+                    message=chat.ResponseMessage(
+                        **convert_message_to_mlflow_chat(resp).model_dump_compat()
+                    ),
+                    finish_reason=stop_reason,
+                )
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=resp["usage"]["input_tokens"],
+                completion_tokens=resp["usage"]["output_tokens"],
+                total_tokens=resp["usage"]["input_tokens"] + resp["usage"]["output_tokens"],
+            ),
+        )
+
+    @classmethod
+    def chat_streaming_to_model(cls, payload, config):
+        return cls.chat_to_model(payload, config)
+
+    @classmethod
+    def model_to_chat_streaming(cls, resp, config):
+        content = resp.get("delta") or resp.get("content_block") or {}
+        if (stop_reason := content.get("stop_reason")) is not None:
+            stop_reason = "length" if stop_reason == "max_tokens" else "stop"
+        return chat.StreamResponsePayload(
+            id=resp["id"],
+            created=int(time.time()),
+            model=resp["model"],
+            choices=[
+                chat.StreamChoice(
+                    index=resp["index"],
+                    finish_reason=stop_reason,
+                    delta=chat.StreamDelta(
+                        role=None,
+                        content=content.get("text"),
+                    ),
+                )
+            ],
+        )
+
+    @classmethod
+    def model_to_completions(cls, resp, config):
+        stop_reason = "stop" if resp["stop_reason"] == "stop_sequence" else "length"
+
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=resp["model"],
+            choices=[
+                completions.Choice(
+                    index=0,
+                    text=resp["completion"],
+                    finish_reason=stop_reason,
+                )
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        key_mapping = {"max_tokens": "max_tokens_to_sample", "stop": "stop_sequences"}
+
+        payload["model"] = config.model.name
+
+        if "top_p" in payload:
+            raise AIGatewayException(
+                status_code=422,
+                detail="Cannot set both 'temperature' and 'top_p' parameters. "
+                "Please use only the temperature parameter for your query.",
+            )
+        max_tokens = payload.get("max_tokens", MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS)
+
+        if max_tokens > MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS:
+            raise AIGatewayException(
+                status_code=422,
+                detail="Invalid value for max_tokens: cannot exceed "
+                f"{MLFLOW_AI_GATEWAY_ANTHROPIC_MAXIMUM_MAX_TOKENS}.",
+            )
+
+        payload["max_tokens"] = max_tokens
+
+        if payload.get("stream", False):
+            raise AIGatewayException(
+                status_code=422,
+                detail="Setting the 'stream' parameter to 'true' is not supported with the MLflow "
+                "Gateway.",
+            )
+        n = payload.pop("n", 1)
+        if n != 1:
+            raise AIGatewayException(
+                status_code=422,
+                detail=f"'n' must be '1' for the Anthropic provider. Received value: '{n}'.",
+            )
+
+        payload = rename_payload_keys(payload, key_mapping)
+
+        if payload["prompt"].startswith("Human: "):
+            payload["prompt"] = "\n\n" + payload["prompt"]
+
+        if not payload["prompt"].startswith("\n\nHuman: "):
+            payload["prompt"] = "\n\nHuman: " + payload["prompt"]
+
+        if not payload["prompt"].endswith("\n\nAssistant:"):
+            payload["prompt"] = payload["prompt"] + "\n\nAssistant:"
+
+        # The range of Anthropic's temperature is 0-1, but ours is 0-2, so we halve it
+        if "temperature" in payload:
+            payload["temperature"] = 0.5 * payload["temperature"]
+
+        return payload
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        raise NotImplementedError
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        raise NotImplementedError
+
+
+class AnthropicProvider(BaseProvider, AnthropicAdapter):
+    NAME = "Anthropic"
+    CONFIG_TYPE = AnthropicConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, AnthropicConfig):
+            raise TypeError(f"Invalid config type {config.model.config}")
+        self.anthropic_config: AnthropicConfig = config.model.config
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return {
+            "x-api-key": self.anthropic_config.anthropic_api_key,
+            "anthropic-version": self.anthropic_config.anthropic_version,
+        }
+
+    @property
+    def base_url(self) -> str:
+        return "https://api.anthropic.com/v1"
+
+    @property
+    def adapter_class(self) -> type[ProviderAdapter]:
+        return AnthropicAdapter
+
+    def get_endpoint_url(self, route_type: str) -> str:
+        if route_type == "llm/v1/chat":
+            return f"{self.base_url}/messages"
+        elif route_type == "llm/v1/completions":
+            return f"{self.base_url}/complete"
+        else:
+            raise ValueError(f"Invalid route type {route_type}")
+
+    async def chat_stream(
+        self, payload: chat.RequestPayload
+    ) -> AsyncIterable[chat.StreamResponsePayload]:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        stream = send_stream_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="messages",
+            payload=AnthropicAdapter.chat_streaming_to_model(payload, self.config),
+        )
+
+        indices = []
+        metadata = {}
+        async for chunk in stream:
+            chunk = chunk.strip()
+            if not chunk:
+                continue
+
+            # No handling on "event" lines
+            prefix, content = chunk.split(b":", 1)
+            if prefix != b"data":
+                continue
+
+            # See https://docs.anthropic.com/claude/reference/messages-streaming
+            resp = json.loads(content.decode("utf-8"))
+
+            # response id and model are only present in `message_start`
+            if resp["type"] == "message_start":
+                metadata["id"] = resp["message"]["id"]
+                metadata["model"] = resp["message"]["model"]
+                continue
+
+            if resp["type"] not in (
+                "message_delta",
+                "content_block_start",
+                "content_block_delta",
+            ):
+                continue
+
+            index = resp.get("index")
+            if index is not None and index not in indices:
+                indices.append(index)
+
+            resp.update(metadata)
+            if resp["type"] == "message_delta":
+                for index in indices:
+                    yield AnthropicAdapter.model_to_chat_streaming(
+                        {**resp, "index": index},
+                        self.config,
+                    )
+            else:
+                yield AnthropicAdapter.model_to_chat_streaming(resp, self.config)
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="messages",
+            payload=AnthropicAdapter.chat_to_model(payload, self.config),
+        )
+        return AnthropicAdapter.model_to_chat(resp, self.config)
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="complete",
+            payload=AnthropicAdapter.completions_to_model(payload, self.config),
+        )
+
+        # Example response:
+        # Documentation: https://docs.anthropic.com/claude/reference/complete_post
+        # ```
+        # {
+        #     "completion": " Hello! My name is Claude."
+        #     "stop_reason": "stop_sequence",
+        #     "model": "claude-instant-1.1",
+        #     "truncated": False,
+        #     "stop": None,
+        #     "log_id": "dee173f87ddf1357da639dee3c38d833",
+        #     "exception": None,
+        # }
+        # ```
+
+        return AnthropicAdapter.model_to_completions(resp, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/base.py

@@ -0,0 +1,127 @@
+from abc import ABC, abstractmethod
+from typing import AsyncIterable
+
+from mlflow.gateway.base_models import ConfigModel
+from mlflow.gateway.config import RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.schemas import chat, completions, embeddings
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class BaseProvider(ABC):
+    """
+    Base class for MLflow Gateway providers.
+    """
+
+    NAME: str = ""
+    SUPPORTED_ROUTE_TYPES: tuple[str, ...]
+    CONFIG_TYPE: type[ConfigModel]
+
+    def __init__(self, config: RouteConfig):
+        if self.NAME == "":
+            raise ValueError(
+                f"{self.__class__.__name__} is a subclass of BaseProvider and must "
+                f"override 'NAME' attribute as a non-empty string."
+            )
+
+        if not hasattr(self, "CONFIG_TYPE") or not issubclass(self.CONFIG_TYPE, ConfigModel):
+            raise ValueError(
+                f"{self.__class__.__name__} is a subclass of BaseProvider and must "
+                f"override 'CONFIG_TYPE' attribute as a subclass of ConfigModel."
+            )
+
+        self.config = config
+
+    async def chat_stream(
+        self, payload: chat.RequestPayload
+    ) -> AsyncIterable[chat.StreamResponsePayload]:
+        raise AIGatewayException(
+            status_code=501,
+            detail=f"The chat streaming route is not implemented for {self.NAME} models.",
+        )
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        raise AIGatewayException(
+            status_code=501,
+            detail=f"The chat route is not implemented for {self.NAME} models.",
+        )
+
+    async def completions_stream(
+        self, payload: completions.RequestPayload
+    ) -> AsyncIterable[completions.StreamResponsePayload]:
+        raise AIGatewayException(
+            status_code=501,
+            detail=f"The completions streaming route is not implemented for {self.NAME} models.",
+        )
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        raise AIGatewayException(
+            status_code=501,
+            detail=f"The completions route is not implemented for {self.NAME} models.",
+        )
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        raise AIGatewayException(
+            status_code=501,
+            detail=f"The embeddings route is not implemented for {self.NAME} models.",
+        )
+
+    @staticmethod
+    def check_for_model_field(payload):
+        if "model" in payload:
+            raise AIGatewayException(
+                status_code=422,
+                detail="The parameter 'model' is not permitted to be passed. The route being "
+                "queried already defines a model instance.",
+            )
+
+
+class ProviderAdapter(ABC):
+    @classmethod
+    @abstractmethod
+    def model_to_embeddings(cls, resp, config): ...
+
+    @classmethod
+    @abstractmethod
+    def model_to_completions(cls, resp, config): ...
+
+    @classmethod
+    def model_to_completions_streaming(cls, resp, config):
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def completions_to_model(cls, payload, config): ...
+
+    @classmethod
+    def completions_streaming_to_model(cls, payload, config):
+        raise NotImplementedError
+
+    @classmethod
+    def model_to_chat(cls, resp, config):
+        raise NotImplementedError
+
+    @classmethod
+    def model_to_chat_streaming(cls, resp, config):
+        raise NotImplementedError
+
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        raise NotImplementedError
+
+    @classmethod
+    def chat_streaming_to_model(cls, payload, config):
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def embeddings_to_model(cls, payload, config): ...
+
+    @classmethod
+    def check_keys_against_mapping(cls, mapping, payload):
+        for k1, k2 in mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=400, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )

venv/lib/python3.9/site-packages/mlflow/gateway/providers/bedrock.py

@@ -0,0 +1,299 @@
+import json
+import time
+from enum import Enum
+
+from mlflow.gateway.config import AmazonBedrockConfig, AWSIdAndKey, AWSRole, RouteConfig
+from mlflow.gateway.constants import (
+    MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS,
+)
+from mlflow.gateway.exceptions import AIGatewayConfigException, AIGatewayException
+from mlflow.gateway.providers.anthropic import AnthropicAdapter
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.cohere import CohereAdapter
+from mlflow.gateway.providers.utils import rename_payload_keys
+from mlflow.gateway.schemas import completions
+
+AWS_BEDROCK_ANTHROPIC_MAXIMUM_MAX_TOKENS = 8191
+
+
+class AmazonBedrockAnthropicAdapter(AnthropicAdapter):
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        payload = super().chat_to_model(payload, config)
+        # "model" keys are not supported in Bedrock"
+        payload.pop("model", None)
+        return payload
+
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        payload = super().completions_to_model(payload, config)
+
+        if "\n\nHuman:" not in payload.get("stop_sequences", []):
+            payload.setdefault("stop_sequences", []).append("\n\nHuman:")
+
+        payload["max_tokens_to_sample"] = min(
+            payload.get("max_tokens_to_sample", MLFLOW_AI_GATEWAY_ANTHROPIC_DEFAULT_MAX_TOKENS),
+            AWS_BEDROCK_ANTHROPIC_MAXIMUM_MAX_TOKENS,
+        )
+
+        # "model" keys are not supported in Bedrock"
+        payload.pop("model", None)
+        return payload
+
+    @classmethod
+    def model_to_completions(cls, payload, config):
+        payload["model"] = config.model.name
+        return super().model_to_completions(payload, config)
+
+
+class AWSTitanAdapter(ProviderAdapter):
+    # TODO handle top_p, top_k, etc.
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        n = payload.pop("n", 1)
+        if n != 1:
+            raise AIGatewayException(
+                status_code=422,
+                detail=f"'n' must be '1' for AWS Titan models. Received value: '{n}'.",
+            )
+
+        # The range of Titan's temperature is 0-1, but ours is 0-2, so we halve it
+        if "temperature" in payload:
+            payload["temperature"] = 0.5 * payload["temperature"]
+        return {
+            "inputText": payload.pop("prompt"),
+            "textGenerationConfig": rename_payload_keys(
+                payload, {"max_tokens": "maxTokenCount", "stop": "stopSequences"}
+            ),
+        }
+
+    @classmethod
+    def model_to_completions(cls, resp, config):
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=candidate.get("outputText"),
+                    finish_reason=None,
+                )
+                for idx, candidate in enumerate(resp.get("results", []))
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        raise NotImplementedError
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        raise NotImplementedError
+
+
+class AI21Adapter(ProviderAdapter):
+    # TODO handle top_p, top_k, etc.
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        return rename_payload_keys(
+            payload,
+            {
+                "stop": "stopSequences",
+                "n": "numResults",
+                "max_tokens": "maxTokens",
+            },
+        )
+
+    @classmethod
+    def model_to_completions(cls, resp, config):
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=candidate.get("data", {}).get("text"),
+                    finish_reason=None,
+                )
+                for idx, candidate in enumerate(resp.get("completions", []))
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        raise NotImplementedError
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        raise NotImplementedError
+
+
+class AmazonBedrockModelProvider(Enum):
+    AMAZON = "amazon"
+    COHERE = "cohere"
+    AI21 = "ai21"
+    ANTHROPIC = "anthropic"
+
+    @property
+    def adapter_class(self) -> type[ProviderAdapter]:
+        return AWS_MODEL_PROVIDER_TO_ADAPTER.get(self)
+
+    @classmethod
+    def of_str(cls, name: str):
+        name = name.lower()
+
+        for opt in cls:
+            if opt.name.lower() == name or opt.value.lower() == name:
+                return opt
+
+
+AWS_MODEL_PROVIDER_TO_ADAPTER = {
+    AmazonBedrockModelProvider.COHERE: CohereAdapter,
+    AmazonBedrockModelProvider.ANTHROPIC: AmazonBedrockAnthropicAdapter,
+    AmazonBedrockModelProvider.AMAZON: AWSTitanAdapter,
+    AmazonBedrockModelProvider.AI21: AI21Adapter,
+}
+
+
+class AmazonBedrockProvider(BaseProvider):
+    NAME = "Amazon Bedrock"
+    CONFIG_TYPE = AmazonBedrockConfig
+
+    def __init__(self, config: RouteConfig):
+        super().__init__(config)
+
+        if config.model.config is None or not isinstance(config.model.config, AmazonBedrockConfig):
+            raise TypeError(f"Invalid config type {config.model.config}")
+        self.bedrock_config: AmazonBedrockConfig = config.model.config
+        self._client = None
+        self._client_created = 0
+
+    def _client_expired(self):
+        if not isinstance(self.bedrock_config.aws_config, AWSRole):
+            return False
+
+        return (
+            (time.monotonic_ns() - self._client_created)
+            >= (self.bedrock_config.aws_config.session_length_seconds) * 1_000_000_000,
+        )
+
+    def get_bedrock_client(self):
+        import boto3
+        import botocore.exceptions
+
+        if self._client is not None and not self._client_expired():
+            return self._client
+
+        session = boto3.Session(**self._construct_session_args())
+
+        try:
+            self._client, self._client_created = (
+                session.client(
+                    service_name="bedrock-runtime",
+                    **self._construct_client_args(session),
+                ),
+                time.monotonic_ns(),
+            )
+            return self._client
+        except botocore.exceptions.UnknownServiceError as e:
+            raise AIGatewayConfigException(
+                "Cannot create Amazon Bedrock client; ensure boto3/botocore "
+                "linked from the Amazon Bedrock user guide are installed. "
+                "Otherwise likely missing credentials or accessing account without to "
+                "Amazon Bedrock Private Preview"
+            ) from e
+
+    def _construct_session_args(self):
+        session_args = {
+            "region_name": self.bedrock_config.aws_config.aws_region,
+        }
+
+        return {k: v for k, v in session_args.items() if v}
+
+    def _construct_client_args(self, session):
+        aws_config = self.bedrock_config.aws_config
+
+        if isinstance(aws_config, AWSRole):
+            role = session.client(service_name="sts").assume_role(
+                RoleArn=aws_config.aws_role_arn,
+                RoleSessionName="ai-gateway-bedrock",
+                DurationSeconds=aws_config.session_length_seconds,
+            )
+            return {
+                "aws_access_key_id": role["Credentials"]["AccessKeyId"],
+                "aws_secret_access_key": role["Credentials"]["SecretAccessKey"],
+                "aws_session_token": role["Credentials"]["SessionToken"],
+            }
+        elif isinstance(aws_config, AWSIdAndKey):
+            return {
+                "aws_access_key_id": aws_config.aws_access_key_id,
+                "aws_secret_access_key": aws_config.aws_secret_access_key,
+                "aws_session_token": aws_config.aws_session_token,
+            }
+        else:
+            return {}
+
+    @property
+    def _underlying_provider(self):
+        if (not self.config.model.name) or "." not in self.config.model.name:
+            return None
+        provider = self.config.model.name.split(".")[0]
+        return AmazonBedrockModelProvider.of_str(provider)
+
+    @property
+    def adapter_class(self) -> type[ProviderAdapter]:
+        provider = self._underlying_provider
+        if not provider:
+            raise AIGatewayException(
+                status_code=422,
+                detail=f"Unknown Amazon Bedrock model type {self._underlying_provider}",
+            )
+        adapter = provider.adapter_class
+        if not adapter:
+            raise AIGatewayException(
+                status_code=422,
+                detail=f"Don't know how to handle {self._underlying_provider} for Amazon Bedrock",
+            )
+        return adapter
+
+    def _request(self, body):
+        import botocore.exceptions
+
+        try:
+            response = self.get_bedrock_client().invoke_model(
+                body=json.dumps(body).encode(),
+                modelId=self.config.model.name,
+                # defaults
+                # save=False,
+                accept="application/json",
+                contentType="application/json",
+            )
+            return json.loads(response.get("body").read())
+
+        # TODO work though botocore.exceptions to make this catchable.
+        # except botocore.exceptions.ValidationException as e:
+        #     raise HTTPException(status_code=422, detail=str(e)) from e
+
+        except botocore.exceptions.ReadTimeoutError as e:
+            raise AIGatewayException(status_code=408) from e
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        self.check_for_model_field(payload)
+        payload = jsonable_encoder(payload, exclude_none=True, exclude_defaults=True)
+        payload = self.adapter_class.completions_to_model(payload, self.config)
+        response = self._request(payload)
+        return self.adapter_class.model_to_completions(response, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/cohere.py

@@ -0,0 +1,457 @@
+import json
+import time
+from typing import Any, AsyncGenerator, AsyncIterable
+
+from mlflow.gateway.config import CohereConfig, RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.utils import rename_payload_keys, send_request, send_stream_request
+from mlflow.gateway.schemas import chat, completions, embeddings
+
+
+class CohereAdapter(ProviderAdapter):
+    @staticmethod
+    def _scale_temperature(payload):
+        # The range of Cohere's temperature is 0-5, but ours is 0-2, so we scale it.
+        if temperature := payload.get("temperature"):
+            payload["temperature"] = 2.5 * temperature
+        return payload
+
+    @classmethod
+    def model_to_completions(cls, resp, config):
+        # Response example (https://docs.cohere.com/reference/generate)
+        # ```
+        # {
+        #   "id": "string",
+        #   "generations": [
+        #     {
+        #       "id": "string",
+        #       "text": "string"
+        #     }
+        #   ],
+        #   "prompt": "string"
+        # }
+        # ```
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=c["text"],
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["generations"])
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    @classmethod
+    def model_to_completions_streaming(cls, resp, config):
+        # Response example (https://docs.cohere.com/reference/generate)
+        #
+        # Streaming chunks:
+        # ```
+        # {"index":0,"text":" Hi","is_finished":false,"event_type":"text-generation"}
+        # ```
+        # ```
+        # {"index":1,"text":" Hi","is_finished":false,"event_type":"text-generation"}
+        # ```
+        # notes: "index" is only present if "num_generations" > 1
+        #
+        # Final chunk:
+        # ```
+        # {"is_finished":true,"event_type":"stream-end","finish_reason":"COMPLETE",
+        #   "response":{"id":"b32a70c5-8c91-4f96-958f-d942801ed22f",
+        #       "generations":[
+        #           {
+        #               "id":"5d5d0851-35ac-4c25-a9a9-2fbb391bd415",
+        #               "index":0,
+        #               "text":" Hi there! How can I assist you today? ",
+        #               "finish_reason":"COMPLETE"
+        #           },
+        #           {
+        #               "id":"0a24787f-504e-470e-a088-0bf801a2c72d",
+        #               "index":1,
+        #               "text":" Hi there, how can I assist you today? ",
+        #               "finish_reason":"COMPLETE"
+        #           }
+        #       ],
+        #       "prompt":"Hello"
+        #   }}
+        # ```
+        response = resp.get("response")
+        return completions.StreamResponsePayload(
+            id=response["id"] if response else None,
+            created=int(time.time()),
+            model=config.model.name,
+            choices=[
+                completions.StreamChoice(
+                    index=resp.get("index", 0),
+                    finish_reason=resp.get("finish_reason"),
+                    text=resp.get("text"),
+                )
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        # Response example (https://docs.cohere.com/reference/embed):
+        # ```
+        # {
+        #   "id": "bc57846a-3e56-4327-8acc-588ca1a37b8a",
+        #   "texts": [
+        #     "hello world"
+        #   ],
+        #   "embeddings": [
+        #     [
+        #       3.25,
+        #       0.7685547,
+        #       2.65625,
+        #       ...
+        #       -0.30126953,
+        #       -2.3554688,
+        #       1.2597656
+        #     ]
+        #   ],
+        #   "meta": [
+        #     {
+        #       "api_version": [
+        #         {
+        #           "version": "1"
+        #         }
+        #       ]
+        #     }
+        #   ]
+        # }
+        # ```
+        return embeddings.ResponsePayload(
+            data=[
+                embeddings.EmbeddingObject(
+                    embedding=output,
+                    index=idx,
+                )
+                for idx, output in enumerate(resp["embeddings"])
+            ],
+            model=config.model.name,
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        key_mapping = {
+            "stop": "stop_sequences",
+            "n": "num_generations",
+        }
+        cls.check_keys_against_mapping(key_mapping, payload)
+        payload = cls._scale_temperature(payload)
+        return rename_payload_keys(payload, key_mapping)
+
+    @classmethod
+    def completions_streaming_to_model(cls, payload, config):
+        return cls.completions_to_model(payload, config)
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        key_mapping = {"input": "texts"}
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        return rename_payload_keys(payload, key_mapping)
+
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        if payload["n"] != 1:
+            raise AIGatewayException(
+                status_code=422,
+                detail=f"Parameter n must be 1 for Cohere chat, got {payload['n']}.",
+            )
+        del payload["n"]
+
+        if "stop" in payload:
+            raise AIGatewayException(
+                status_code=422,
+                detail="Parameter stop is not supported for Cohere chat.",
+            )
+        payload = cls._scale_temperature(payload)
+
+        messages = payload.pop("messages")
+        last_message = messages.pop()  # pydantic enforces min_items=1
+        if last_message["role"] != "user":
+            raise AIGatewayException(
+                status_code=422,
+                detail=f"Last message must be from user, got {last_message['role']}.",
+            )
+        payload["message"] = last_message["content"]
+
+        # Cohere uses `preamble_override` to set the system message
+        # we concatenate all system messages from the user with a newline
+        system_messages = [m for m in messages if m["role"] == "system"]
+        if len(system_messages) > 0:
+            payload["preamble_override"] = "\n".join(m["content"] for m in system_messages)
+
+        # remaining messages are chat history
+        # we want to include only user and assistant messages
+        messages = [m for m in messages if m["role"] in ("user", "assistant")]
+        if messages:
+            payload["chat_history"] = [
+                {
+                    "role": "USER" if m["role"] == "user" else "CHATBOT",
+                    "message": m["content"],
+                }
+                for m in messages
+            ]
+        return payload
+
+    @classmethod
+    def chat_streaming_to_model(cls, payload, config):
+        return cls.chat_to_model(payload, config)
+
+    @classmethod
+    def model_to_chat(cls, resp, config):
+        # Response example (https://docs.cohere.com/reference/chat)
+        # ```
+        # {
+        #   "response_id": "string",
+        #   "text": "string",
+        #   "generation_id": "string",
+        #   "token_count": {
+        #     "prompt_tokens": 0,
+        #     "response_tokens": 0,
+        #     "total_tokens": 0,
+        #     "billed_tokens": 0
+        #   },
+        #   "meta": {
+        #     "api_version": {
+        #       "version": "1"
+        #     },
+        #     "billed_units": {
+        #       "input_tokens": 0,
+        #       "output_tokens": 0
+        #     }
+        #   },
+        #   "tool_inputs": null
+        # }
+        # ```
+        return chat.ResponsePayload(
+            id=resp["response_id"],
+            object="chat.completion",
+            created=int(time.time()),
+            model=config.model.name,
+            choices=[
+                chat.Choice(
+                    index=0,
+                    message=chat.ResponseMessage(
+                        role="assistant",
+                        content=resp["text"],
+                    ),
+                    finish_reason=None,
+                ),
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=resp["token_count"]["prompt_tokens"],
+                completion_tokens=resp["token_count"]["response_tokens"],
+                total_tokens=resp["token_count"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_chat_streaming(cls, resp, config):
+        # Response example (https://docs.cohere.com/reference/chat)
+        # Streaming chunks:
+        # ```
+        # {
+        #   "is_finished":false,
+        #   "event_type":"stream-start",
+        #   "generation_id":"string"
+        # }
+        # {"is_finished":false,"event_type":"text-generation","text":"How"}
+        # {"is_finished":false,"event_type":"text-generation","text":" are"}
+        # {"is_finished":false,"event_type":"text-generation","text":" you"}
+        # {
+        #   "is_finished":true,
+        #   "event_type":"stream-end",
+        #   "response":{
+        #     "response_id":"string",
+        #     "text":"How are you",
+        #     "generation_id":"string",
+        #     "token_count":{
+        #       "prompt_tokens":83,"response_tokens":63,"total_tokens":146,"billed_tokens":128
+        #     },
+        #     "tool_inputs":null
+        #   },
+        #   "finish_reason":"COMPLETE"
+        # }
+        # ```
+        response = resp.get("response")
+        return chat.StreamResponsePayload(
+            # first chunk has "generation_id" but not "response_id"
+            id=response["response_id"] if response else None,
+            created=int(time.time()),
+            model=config.model.name,
+            choices=[
+                chat.StreamChoice(
+                    index=0,
+                    finish_reason=resp.get("finish_reason"),
+                    delta=chat.StreamDelta(
+                        role=None,
+                        content=resp.get("text"),
+                    ),
+                )
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=response["token_count"]["prompt_tokens"] if response else None,
+                completion_tokens=response["token_count"]["response_tokens"] if response else None,
+                total_tokens=response["token_count"]["total_tokens"] if response else None,
+            ),
+        )
+
+
+class CohereProvider(BaseProvider):
+    NAME = "Cohere"
+    CONFIG_TYPE = CohereConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, CohereConfig):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.cohere_config: CohereConfig = config.model.config
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return {"Authorization": f"Bearer {self.cohere_config.cohere_api_key}"}
+
+    @property
+    def base_url(self) -> str:
+        return "https://api.cohere.ai/v1"
+
+    @property
+    def adapter_class(self) -> type[ProviderAdapter]:
+        return CohereAdapter
+
+    def get_endpoint_url(self, route_type: str) -> str:
+        if route_type == "llm/v1/chat":
+            return f"{self.base_url}/chat"
+        elif route_type == "llm/v1/completions":
+            return f"{self.base_url}/generate"
+        elif route_type == "llm/v1/embeddings":
+            return f"{self.base_url}/embed"
+        else:
+            raise ValueError(f"Invalid route type {route_type}")
+
+    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        return await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path=path,
+            payload=payload,
+        )
+
+    def _stream_request(self, path: str, payload: dict[str, Any]) -> AsyncGenerator[bytes, None]:
+        return send_stream_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path=path,
+            payload=payload,
+        )
+
+    async def chat_stream(
+        self, payload: chat.RequestPayload
+    ) -> AsyncIterable[chat.StreamResponsePayload]:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        stream = self._stream_request(
+            "chat",
+            {
+                "model": self.config.model.name,
+                **CohereAdapter.chat_streaming_to_model(payload, self.config),
+            },
+        )
+        async for chunk in stream:
+            if not chunk:
+                continue
+
+            resp = json.loads(chunk)
+            if resp["event_type"] == "stream-start":
+                continue
+            yield CohereAdapter.model_to_chat_streaming(resp, self.config)
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await self._request(
+            "chat",
+            {
+                "model": self.config.model.name,
+                **CohereAdapter.chat_to_model(payload, self.config),
+            },
+        )
+        return CohereAdapter.model_to_chat(resp, self.config)
+
+    async def completions_stream(
+        self, payload: completions.RequestPayload
+    ) -> AsyncIterable[completions.StreamResponsePayload]:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        stream = self._stream_request(
+            "generate",
+            {
+                "model": self.config.model.name,
+                **CohereAdapter.completions_streaming_to_model(payload, self.config),
+            },
+        )
+        async for chunk in stream:
+            if not chunk:
+                continue
+
+            resp = json.loads(chunk)
+            yield CohereAdapter.model_to_completions_streaming(resp, self.config)
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await self._request(
+            "generate",
+            {
+                "model": self.config.model.name,
+                **CohereAdapter.completions_to_model(payload, self.config),
+            },
+        )
+        return CohereAdapter.model_to_completions(resp, self.config)
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await self._request(
+            "embed",
+            {
+                "model": self.config.model.name,
+                **CohereAdapter.embeddings_to_model(payload, self.config),
+            },
+        )
+        return CohereAdapter.model_to_embeddings(resp, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/gemini.py

@@ -0,0 +1,150 @@
+from typing import Any
+
+from mlflow.gateway.config import GeminiConfig, RouteConfig
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.utils import send_request
+from mlflow.gateway.schemas import embeddings
+
+
+class GeminiAdapter(ProviderAdapter):
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        # Example payload for the embedding API.
+        # Documentation: https://ai.google.dev/api/embeddings#v1beta.ContentEmbedding
+        #
+        # {
+        #     "requests": [
+        #         {
+        #             "model": "models/text-embedding-004",
+        #             "content": {
+        #                 "parts": [
+        #                     {
+        #                         "text": "What is the meaning of life?"
+        #                     }
+        #                 ]
+        #             }
+        #         },
+        #         {
+        #             "model": "models/text-embedding-004",
+        #             "content": {
+        #                 "parts": [
+        #                     {
+        #                         "text": "How much wood would a woodchuck chuck?"
+        #                     }
+        #                 ]
+        #             }
+        #         },
+        #         {
+        #             "model": "models/text-embedding-004",
+        #             "content": {
+        #                 "parts": [
+        #                     {
+        #                         "text": "How does the brain work?"
+        #                     }
+        #                 ]
+        #             }
+        #         }
+        #     ]
+        # }
+
+        texts = payload["input"]
+        if isinstance(texts, str):
+            texts = [texts]
+        return (
+            {"content": {"parts": [{"text": texts[0]}]}}
+            if len(texts) == 1
+            else {
+                "requests": [
+                    {"model": f"models/{config.model.name}", "content": {"parts": [{"text": text}]}}
+                    for text in texts
+                ]
+            }
+        )
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        # Documentation: https://ai.google.dev/api/embeddings#v1beta.ContentEmbedding
+        #
+        # Example Response:
+        # {
+        #   "embeddings": [
+        #     {
+        #       "values": [
+        #         3.25,
+        #         0.7685547,
+        #         2.65625,
+        #         ...,
+        #         -0.30126953,
+        #         -2.3554688,
+        #         1.2597656
+        #       ]
+        #     }
+        #   ]
+        # }
+
+        data = [
+            embeddings.EmbeddingObject(embedding=item.get("values", []), index=i)
+            for i, item in enumerate(resp.get("embeddings") or [resp.get("embedding", {})])
+        ]
+
+        # Create and return response payload directly
+        return embeddings.ResponsePayload(
+            data=data,
+            model=config.model.name,
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+
+class GeminiProvider(BaseProvider):
+    NAME = "Gemini"
+    CONFIG_TYPE = GeminiConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, GeminiConfig):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.gemini_config: GeminiConfig = config.model.config
+
+    @property
+    def headers(self):
+        return {"x-goog-api-key": self.gemini_config.gemini_api_key}
+
+    @property
+    def base_url(self):
+        return "https://generativelanguage.googleapis.com/v1beta/models"
+
+    @property
+    def adapter_class(self):
+        return GeminiAdapter
+
+    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        return await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path=path,
+            payload=payload,
+        )
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+
+        embedding_payload = self.adapter_class.embeddings_to_model(payload, self.config)
+        # Documentation: https://ai.google.dev/api/embeddings
+
+        # Use the batch endpoint if payload contains "requests"
+        if "requests" in embedding_payload:
+            endpoint_suffix = ":batchEmbedContents"
+        else:
+            endpoint_suffix = ":embedContent"
+
+        resp = await self._request(
+            f"{self.config.model.name}{endpoint_suffix}",
+            embedding_payload,
+        )
+        return self.adapter_class.model_to_embeddings(resp, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/huggingface.py

@@ -0,0 +1,115 @@
+import time
+from typing import Any
+
+from mlflow.gateway.config import HuggingFaceTextGenerationInferenceConfig, RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider
+from mlflow.gateway.providers.utils import (
+    rename_payload_keys,
+    send_request,
+)
+from mlflow.gateway.schemas import completions
+
+
+class HFTextGenerationInferenceServerProvider(BaseProvider):
+    NAME = "Hugging Face Text Generation Inference"
+    CONFIG_TYPE = HuggingFaceTextGenerationInferenceConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(
+            config.model.config, HuggingFaceTextGenerationInferenceConfig
+        ):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.huggingface_config: HuggingFaceTextGenerationInferenceConfig = config.model.config
+        self.headers = {"Content-Type": "application/json"}
+
+    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        return await send_request(
+            headers=self.headers,
+            base_url=self.huggingface_config.hf_server_url,
+            path=path,
+            payload=payload,
+        )
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        key_mapping = {
+            "max_tokens": "max_new_tokens",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+
+        # HF TGI does not support generating multiple candidates.
+        n = payload.pop("n", 1)
+        if n != 1:
+            raise AIGatewayException(
+                status_code=422,
+                detail="'n' must be '1' for the Text Generation Inference provider."
+                f"Received value: '{n}'.",
+            )
+        prompt = payload.pop("prompt")
+        parameters = rename_payload_keys(payload, key_mapping)
+
+        # The range of HF TGI's temperature is 0-100, but ours is 0-2, so we multiply
+        # by 50
+        payload["temperature"] = 50 * payload["temperature"]
+        # HF TGI does not support 0 temperature
+        parameters["temperature"] = max(payload["temperature"], 1e-3)
+        parameters["details"] = True
+        parameters["decoder_input_details"] = True
+
+        final_payload = {"inputs": prompt, "parameters": parameters}
+
+        resp = await self._request(
+            "generate",
+            final_payload,
+        )
+
+        # Example Response:
+        # Documentation: https://huggingface.github.io/text-generation-inference/#/Text%20Generation%20Inference/compat_generate
+        # {'details': {'best_of_sequences': [{'finish_reason': 'length',
+        #     'generated_text': 'test',
+        #     'generated_tokens': 1,
+        #     'prefill': [{'id': 0, 'logprob': -0.34, 'text': 'test'}],
+        #     'seed': 42,
+        #     'tokens': [{'id': 0, 'logprob': -0.34, 'special': False, 'text': 'test'}],
+        #     'top_tokens': [[{'id': 0,
+        #     'logprob': -0.34,
+        #     'special': False,
+        #     'text': 'test'}]]}],
+        # 'finish_reason': 'length',
+        # 'generated_tokens': 1,
+        # 'prefill': [{'id': 0, 'logprob': -0.34, 'text': 'test'}],
+        # 'seed': 42,
+        # 'tokens': [{'id': 0, 'logprob': -0.34, 'special': False, 'text': 'test'}],
+        # 'top_tokens': [[{'id': 0,
+        #     'logprob': -0.34,
+        #     'special': False,
+        #     'text': 'test'}]]},
+        # 'generated_text': 'test'}
+        output_tokens = resp["details"]["generated_tokens"]
+        input_tokens = len(resp["details"]["prefill"])
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=self.config.model.name,
+            choices=[
+                completions.Choice(
+                    index=0,
+                    text=resp["generated_text"],
+                    finish_reason=resp["details"]["finish_reason"],
+                )
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=input_tokens,
+                completion_tokens=output_tokens,
+                total_tokens=input_tokens + output_tokens,
+            ),
+        )

venv/lib/python3.9/site-packages/mlflow/gateway/providers/mistral.py

@@ -0,0 +1,206 @@
+import time
+from typing import Any
+
+from mlflow.gateway.config import MistralConfig, RouteConfig
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.utils import send_request
+from mlflow.gateway.schemas import chat, completions, embeddings
+
+
+class MistralAdapter(ProviderAdapter):
+    @classmethod
+    def model_to_completions(cls, resp, config):
+        # Response example (https://docs.mistral.ai/api/#operation/createChatCompletion)
+        # ```
+        # {
+        #   "id": "string",
+        #   "object": "string",
+        #   "created": "integer",
+        #   "model": "string",
+        #   "choices": [
+        #     {
+        #       "index": "integer",
+        #       "message": {
+        #           "role": "string",
+        #           "content": "string"
+        #       },
+        #       "finish_reason": "string",
+        #     }
+        #   ],
+        #   "usage":
+        #   {
+        #       "prompt_tokens": "integer",
+        #       "completion_tokens": "integer",
+        #       "total_tokens": "integer",
+        #   }
+        # }
+        # ```
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=c["message"]["content"],
+                    finish_reason=c["finish_reason"],
+                )
+                for idx, c in enumerate(resp["choices"])
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                completion_tokens=resp["usage"]["completion_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_chat(cls, resp, config):
+        # Response example (https://docs.mistral.ai/api/#operation/createChatCompletion)
+        return chat.ResponsePayload(
+            id=resp["id"],
+            object=resp["object"],
+            created=resp["created"],
+            model=resp["model"],
+            choices=[
+                chat.Choice(
+                    index=idx,
+                    message=chat.ResponseMessage(
+                        role=c["message"]["role"],
+                        content=c["message"].get("content"),
+                        tool_calls=(
+                            (calls := c["message"].get("tool_calls"))
+                            and [chat.ToolCall(**c) for c in calls]
+                        ),
+                    ),
+                    finish_reason=c.get("finish_reason"),
+                )
+                for idx, c in enumerate(resp["choices"])
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                completion_tokens=resp["usage"]["completion_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        # Response example (https://docs.mistral.ai/api/#operation/createEmbedding):
+        # ```
+        # {
+        #   "id": "string",
+        #   "object": "string",
+        #   "data": [
+        #     {
+        #       "object": "string",
+        #       "embedding":
+        #       [
+        #           float,
+        #           float
+        #       ]
+        #       "index": "integer",
+        #     }
+        #   ],
+        #   "model": "string",
+        #   "usage":
+        #   {
+        #       "prompt_tokens": "integer",
+        #       "total_tokens": "integer",
+        #   }
+        # }
+        # ```
+        return embeddings.ResponsePayload(
+            data=[
+                embeddings.EmbeddingObject(
+                    embedding=data["embedding"],
+                    index=data["index"],
+                )
+                for data in resp["data"]
+            ],
+            model=config.model.name,
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        payload["model"] = config.model.name
+        payload.pop("stop", None)
+        payload.pop("n", None)
+        payload["messages"] = [{"role": "user", "content": payload.pop("prompt")}]
+
+        # The range of Mistral's temperature is 0-1, but ours is 0-2, so we scale it.
+        if "temperature" in payload:
+            payload["temperature"] = 0.5 * payload["temperature"]
+
+        return payload
+
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        return {"model": config.model.name, **payload}
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        return {"model": config.model.name, **payload}
+
+
+class MistralProvider(BaseProvider):
+    NAME = "Mistral"
+    CONFIG_TYPE = MistralConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, MistralConfig):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.mistral_config: MistralConfig = config.model.config
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return {"Authorization": f"Bearer {self.mistral_config.mistral_api_key}"}
+
+    @property
+    def base_url(self) -> str:
+        return "https://api.mistral.ai/v1"
+
+    @property
+    def adapter_class(self) -> type[ProviderAdapter]:
+        return MistralAdapter
+
+    def get_endpoint_url(self, route_type: str) -> str:
+        if route_type == "llm/v1/chat":
+            return f"{self.base_url}/chat/completions"
+        else:
+            raise ValueError(f"Invalid route type {route_type}")
+
+    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        return await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path=path,
+            payload=payload,
+        )
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await self._request(
+            "chat/completions",
+            MistralAdapter.completions_to_model(payload, self.config),
+        )
+        return MistralAdapter.model_to_completions(resp, self.config)
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await self._request(
+            "embeddings",
+            MistralAdapter.embeddings_to_model(payload, self.config),
+        )
+        return MistralAdapter.model_to_embeddings(resp, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/mlflow.py

@@ -0,0 +1,238 @@
+import time
+
+from pydantic import BaseModel, StrictFloat, StrictStr, ValidationError
+
+from mlflow.gateway.config import MlflowModelServingConfig, RouteConfig
+from mlflow.gateway.constants import MLFLOW_SERVING_RESPONSE_KEY
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider
+from mlflow.gateway.providers.utils import send_request
+from mlflow.gateway.schemas import chat, completions, embeddings
+from mlflow.utils.pydantic_utils import field_validator
+
+
+class ServingTextResponse(BaseModel):
+    predictions: list[StrictStr]
+
+    @field_validator("predictions", mode="before")
+    def extract_choices(cls, predictions):
+        if isinstance(predictions, list) and not predictions:
+            raise ValueError("The input list is empty")
+        if isinstance(predictions, dict):
+            if "choices" not in predictions and len(predictions) > 1:
+                raise ValueError(
+                    "The dict format is invalid for this route type. Ensure the served model "
+                    "returns a dict key containing 'choices'"
+                )
+            if len(predictions) == 1:
+                predictions = next(iter(predictions.values()))
+            else:
+                predictions = predictions.get("choices", predictions)
+            if not predictions:
+                raise ValueError("The input list is empty")
+        return predictions
+
+
+class EmbeddingsResponse(BaseModel):
+    predictions: list[list[StrictFloat]]
+
+    @field_validator("predictions", mode="before")
+    def validate_predictions(cls, predictions):
+        if isinstance(predictions, list) and not predictions:
+            raise ValueError("The input list is empty")
+        if isinstance(predictions, list) and all(
+            isinstance(item, list) and not item for item in predictions
+        ):
+            raise ValueError("One or more lists in the returned prediction response are empty")
+        elif all(isinstance(item, float) for item in predictions):
+            return [predictions]
+        else:
+            return predictions
+
+
+class MlflowModelServingProvider(BaseProvider):
+    NAME = "MLflow Model Serving"
+    CONFIG_TYPE = MlflowModelServingConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(
+            config.model.config, MlflowModelServingConfig
+        ):
+            raise TypeError(f"Invalid config type {config.model.config}")
+        self.mlflow_config: MlflowModelServingConfig = config.model.config
+        self.headers = {"Content-Type": "application/json"}
+
+    @staticmethod
+    def _extract_mlflow_response_key(response):
+        if MLFLOW_SERVING_RESPONSE_KEY not in response:
+            raise AIGatewayException(
+                status_code=502,
+                detail=f"The response is missing the required key: {MLFLOW_SERVING_RESPONSE_KEY}.",
+            )
+        return response[MLFLOW_SERVING_RESPONSE_KEY]
+
+    @staticmethod
+    def _process_payload(payload, key):
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+
+        input_data = payload.pop(key, None)
+        request_payload = {"inputs": input_data if isinstance(input_data, list) else [input_data]}
+
+        if payload:
+            request_payload["params"] = payload
+
+        return request_payload
+
+    @staticmethod
+    def _process_completions_response_for_mlflow_serving(response):
+        try:
+            validated_response = ServingTextResponse(**response)
+            inference_data = validated_response.predictions
+        except ValidationError as e:
+            raise AIGatewayException(status_code=502, detail=str(e))
+
+        return [
+            completions.Choice(index=idx, text=entry, finish_reason=None)
+            for idx, entry in enumerate(inference_data)
+        ]
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        # Example request to MLflow REST API server for completions:
+        # {
+        #     "inputs": ["hi", "hello", "bye"],
+        #     "params": {
+        #         "temperature": 0.5,
+        #         "top_k": 3,
+        #     }
+        # }
+
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.mlflow_config.model_server_url,
+            path="invocations",
+            payload=self._process_payload(payload, "prompt"),
+        )
+
+        # Example response:
+        # {"predictions": ["hello", "hi", "goodbye"]}
+
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=self.config.model.name,
+            choices=self._process_completions_response_for_mlflow_serving(resp),
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    def _process_chat_response_for_mlflow_serving(self, response):
+        try:
+            validated_response = ServingTextResponse(**response)
+            inference_data = validated_response.predictions
+        except ValidationError as e:
+            raise AIGatewayException(status_code=502, detail=str(e))
+
+        return [
+            {"message": {"role": "assistant", "content": entry}, "metadata": {}}
+            for entry in inference_data
+        ]
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        # Example request to MLflow REST API for chat:
+        # {
+        #     "inputs": ["question"],
+        #     "params": ["temperature": 0.2],
+        # }
+
+        payload = self._process_payload(payload, "messages")
+
+        query_count = len(payload["inputs"])
+        if query_count > 1:
+            raise AIGatewayException(
+                status_code=422,
+                detail="MLflow chat models are only capable of processing a single query at a "
+                f"time. The request submitted consists of {query_count} queries.",
+            )
+
+        payload["inputs"] = [payload["inputs"][0]["content"]]
+
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.mlflow_config.model_server_url,
+            path="invocations",
+            payload=payload,
+        )
+
+        # Example response:
+        # {"predictions": ["answer"]}
+
+        return chat.ResponsePayload(
+            created=int(time.time()),
+            model=self.config.model.name,
+            choices=[
+                chat.Choice(
+                    index=idx,
+                    message=chat.ResponseMessage(
+                        role=c["message"]["role"], content=c["message"]["content"]
+                    ),
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(self._process_chat_response_for_mlflow_serving(resp))
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    def _process_embeddings_response_for_mlflow_serving(self, response):
+        try:
+            validated_response = EmbeddingsResponse(**response)
+            inference_data = validated_response.predictions
+        except ValidationError as e:
+            raise AIGatewayException(status_code=502, detail=str(e))
+
+        return inference_data
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        # Example request to MLflow REST API server for embeddings:
+        # {
+        #     "inputs": ["a sentence", "another sentence"],
+        #     "params": {
+        #         "output_value": "token_embeddings",
+        #     }
+        # }
+
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.mlflow_config.model_server_url,
+            path="invocations",
+            payload=self._process_payload(payload, "input"),
+        )
+
+        # Example response:
+        # {"predictions": [[0.100, -0.234, 0.002, ...], [0.222, -0.111, 0.134, ...]]}
+
+        return embeddings.ResponsePayload(
+            data=[
+                embeddings.EmbeddingObject(
+                    embedding=embedding,
+                    index=idx,
+                )
+                for idx, embedding in enumerate(
+                    self._process_embeddings_response_for_mlflow_serving(resp)
+                )
+            ],
+            model=self.config.model.name,
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=None,
+                total_tokens=None,
+            ),
+        )

venv/lib/python3.9/site-packages/mlflow/gateway/providers/mosaicml.py

@@ -0,0 +1,301 @@
+import time
+from contextlib import contextmanager
+from typing import Any
+
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.config import MosaicMLConfig, RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider
+from mlflow.gateway.providers.utils import rename_payload_keys, send_request
+from mlflow.gateway.schemas import chat, completions, embeddings
+
+
+class MosaicMLProvider(BaseProvider):
+    NAME = "MosaicML"
+    CONFIG_TYPE = MosaicMLConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, MosaicMLConfig):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.mosaicml_config: MosaicMLConfig = config.model.config
+
+    async def _request(self, model: str, payload: dict[str, Any]) -> dict[str, Any]:
+        headers = {"Authorization": f"{self.mosaicml_config.mosaicml_api_key}"}
+        return await send_request(
+            headers=headers,
+            base_url=self.mosaicml_config.mosaicml_api_base
+            or "https://models.hosted-on.mosaicml.hosting",
+            path=model + "/v1/predict",
+            payload=payload,
+        )
+
+    # NB: as this parser performs no blocking operations, we are intentionally not defining it
+    # as async due to the overhead of spawning an additional thread if we did.
+    @staticmethod
+    def _parse_chat_messages_to_prompt(messages: list[chat.RequestMessage]) -> str:
+        """
+        This parser is based on the format described in
+        https://huggingface.co/blog/llama2#how-to-prompt-llama-2 .
+        The expected format is:
+            "<s>[INST] <<SYS>>
+            {{ system_prompt }}
+            <</SYS>>
+
+            {{ user_msg_1 }} [/INST] {{ model_answer_1 }} </s>
+            <s>[INST] {{ user_msg_2 }} [/INST]"
+        """
+        prompt = "<s>"  # Always start with an opening <s> tag
+        for m in messages:
+            if m.role == "system" or m.role == "user":
+                inst = m.content
+
+                # Wrap system messages in <<SYS>> tags
+                if m.role == "system":
+                    inst = f"<<SYS>> {inst} <</SYS>>"
+
+                # Close the [INST] tag
+                inst += " [/INST]"
+
+                # If the previous message was a system/user message,
+                # remove previous closing [/INST] tag
+                if prompt.endswith("[/INST]"):
+                    prompt = prompt[:-7]
+                # Otherwise, add an opening [INST] tag
+                else:
+                    inst = f"[INST] {inst}"
+                prompt += inst
+            elif m.role == "assistant":
+                # Add statement closing/opening tags by default
+                prompt += f" {m.content} </s><s>"
+            else:
+                raise MlflowException.invalid_parameter_value(
+                    f"Invalid role {m.role} inputted. Must be one of 'system', "
+                    "'user', or 'assistant'.",
+                )
+
+        # Remove the last </s><s> tags if they exist to allow for
+        # assistant completion prompts.
+        if prompt.endswith("</s><s>"):
+            prompt = prompt[:-7]
+        return prompt
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        # Extract the List[RequestMessage] from the RequestPayload
+        messages = payload.messages
+        payload = jsonable_encoder(payload, exclude_none=True)
+        # remove the messages from the remaining configuration items
+        payload.pop("messages", None)
+        self.check_for_model_field(payload)
+        key_mapping = {
+            "max_tokens": "max_new_tokens",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        payload = rename_payload_keys(payload, key_mapping)
+
+        # Handle 'prompt' field in payload
+        try:
+            prompt = [self._parse_chat_messages_to_prompt(messages)]
+        except MlflowException as e:
+            raise AIGatewayException(
+                status_code=422, detail=f"An invalid request structure was submitted. {e.message}"
+            )
+        # Construct final payload structure
+        final_payload = {"inputs": prompt, "parameters": payload}
+
+        # Input data structure for Mosaic Text Completion endpoint
+        #
+        # {"inputs": [prompt],
+        #  {
+        #    "parameters": {
+        #      "temperature": 0.2
+        #    }
+        #   }
+        # }
+        with custom_token_allowance_exceeded_handling():
+            resp = await self._request(
+                self.config.model.name,
+                final_payload,
+            )
+        # Response example
+        # (https://docs.mosaicml.com/en/latest/inference.html#text-completion-models)
+        # ```
+        # {
+        #   "outputs": [
+        #     "string",
+        #   ],
+        # }
+        # ```
+        return chat.ResponsePayload(
+            created=int(time.time()),
+            model=self.config.model.name,
+            choices=[
+                chat.Choice(
+                    index=idx,
+                    message=chat.ResponseMessage(role="assistant", content=c),
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["outputs"])
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        key_mapping = {
+            "max_tokens": "max_new_tokens",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        payload = rename_payload_keys(payload, key_mapping)
+
+        # Handle 'prompt' field in payload
+        prompt = payload.pop("prompt")
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        # Construct final payload structure
+        final_payload = {"inputs": prompt, "parameters": payload}
+
+        # Input data structure for Mosaic Text Completion endpoint
+        #
+        # {"inputs": [prompt],
+        #  {
+        #    "parameters": {
+        #      "temperature": 0.2
+        #    }
+        #   }
+        # }
+
+        with custom_token_allowance_exceeded_handling():
+            resp = await self._request(
+                self.config.model.name,
+                final_payload,
+            )
+
+        # Response example
+        # (https://docs.mosaicml.com/en/latest/inference.html#text-completion-models)
+        # ```
+        # {
+        #   "outputs": [
+        #     "string",
+        #   ],
+        # }
+        # ```
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=self.config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=c,
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["outputs"])
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        key_mapping = {"input": "inputs"}
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        payload = rename_payload_keys(payload, key_mapping)
+
+        # Ensure 'inputs' is a list of strings
+        if isinstance(payload["inputs"], str):
+            payload["inputs"] = [payload["inputs"]]
+
+        resp = await self._request(
+            self.config.model.name,
+            payload,
+        )
+        # Response example
+        # (https://docs.mosaicml.com/en/latest/inference.html#text-embedding-models):
+        # ```
+        # {
+        #   "outputs": [
+        #     [
+        #       3.25,
+        #       0.7685547,
+        #       2.65625,
+        #       ...
+        #       -0.30126953,
+        #       -2.3554688,
+        #       1.2597656
+        #     ]
+        #   ]
+        # }
+        # ```
+        return embeddings.ResponsePayload(
+            data=[
+                embeddings.EmbeddingObject(
+                    embedding=output,
+                    index=idx,
+                )
+                for idx, output in enumerate(resp["outputs"])
+            ],
+            model=self.config.model.name,
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+
+@contextmanager
+def custom_token_allowance_exceeded_handling():
+    """
+    Context manager handler for specific error messages that are incorrectly set as server-side
+    errors, but are in actuality an issue with the request sent to the external provider.
+    """
+    from fastapi import HTTPException
+
+    try:
+        yield
+    except HTTPException as e:
+        status_code = e.status_code
+        detail = e.detail or {}
+
+        if (
+            status_code == 500
+            and detail
+            and any(
+                detail.get("message", "").startswith(x)
+                for x in (
+                    "Error: max output tokens is limited to",
+                    "Error: prompt token count",
+                )
+            )
+        ):
+            raise HTTPException(status_code=422, detail=detail)
+        else:
+            raise

venv/lib/python3.9/site-packages/mlflow/gateway/providers/openai.py

@@ -0,0 +1,600 @@
+import json
+import os
+from typing import TYPE_CHECKING, AsyncIterable
+from urllib.parse import urlparse, urlunparse
+
+from mlflow.environment_variables import MLFLOW_ENABLE_UC_FUNCTIONS
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.config import OpenAIAPIType, OpenAIConfig, RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.utils import send_request, send_stream_request
+from mlflow.gateway.schemas import chat, completions, embeddings
+from mlflow.gateway.uc_function_utils import (
+    _UC_FUNCTION,
+    TokenUsageAccumulator,
+    execute_function,
+    get_func_schema,
+    join_uc_functions,
+    parse_uc_functions,
+    prepend_uc_functions,
+)
+from mlflow.gateway.utils import handle_incomplete_chunks, strip_sse_prefix
+from mlflow.utils.uri import append_to_uri_path, append_to_uri_query_params
+
+if TYPE_CHECKING:
+    from databricks.sdk import FunctionInfo
+
+
+# To mock the WorkspaceClient in tests
+def _get_workspace_client():
+    try:
+        from databricks.sdk import WorkspaceClient
+
+        return WorkspaceClient()
+    except ImportError:
+        raise AIGatewayException(
+            message="Databricks SDK is required to use Unity Catalog integration",
+            error_code=404,
+        )
+
+
+class OpenAIAdapter(ProviderAdapter):
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        return cls._add_model_to_payload_if_necessary(payload, config)
+
+    @classmethod
+    def completion_to_model(cls, payload, config):
+        return cls._add_model_to_payload_if_necessary(payload, config)
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        return cls._add_model_to_payload_if_necessary(payload, config)
+
+    @classmethod
+    def _add_model_to_payload_if_necessary(cls, payload, config):
+        # NB: For Azure OpenAI, the deployment name (which is included in the URL) specifies
+        # the model; it is not specified in the payload. For OpenAI outside of Azure, the
+        # model is always specified in the payload
+        if config.model.config.openai_api_type not in (OpenAIAPIType.AZURE, OpenAIAPIType.AZUREAD):
+            return {"model": config.model.name, **payload}
+        else:
+            return payload
+
+    @classmethod
+    def model_to_chat(cls, resp, config):
+        # Response example (https://platform.openai.com/docs/api-reference/chat/create)
+        # ```
+        # {
+        #    "id":"chatcmpl-abc123",
+        #    "object":"chat.completion",
+        #    "created":1677858242,
+        #    "model":"gpt-4o-mini",
+        #    "usage":{
+        #       "prompt_tokens":13,
+        #       "completion_tokens":7,
+        #       "total_tokens":20
+        #    },
+        #    "choices":[
+        #       {
+        #          "message":{
+        #             "role":"assistant",
+        #             "content":"\n\nThis is a test!"
+        #          },
+        #          "finish_reason":"stop",
+        #          "index":0
+        #       }
+        #    ]
+        # }
+        # ```
+        return chat.ResponsePayload(
+            id=resp["id"],
+            object=resp["object"],
+            created=resp["created"],
+            model=resp["model"],
+            choices=[
+                chat.Choice(
+                    index=idx,
+                    message=chat.ResponseMessage(
+                        role=c["message"]["role"],
+                        content=c["message"].get("content"),
+                        tool_calls=(
+                            (calls := c["message"].get("tool_calls"))
+                            and [chat.ToolCall(**c) for c in calls]
+                        ),
+                    ),
+                    finish_reason=c.get("finish_reason"),
+                )
+                for idx, c in enumerate(resp["choices"])
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                completion_tokens=resp["usage"]["completion_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_chat_streaming(cls, resp, config):
+        return chat.StreamResponsePayload(
+            id=resp["id"],
+            object=resp["object"],
+            created=resp["created"],
+            model=resp["model"],
+            choices=[
+                chat.StreamChoice(
+                    index=c["index"],
+                    finish_reason=c["finish_reason"],
+                    delta=chat.StreamDelta(
+                        role=c["delta"].get("role"), content=c["delta"].get("content")
+                    ),
+                )
+                for c in resp["choices"]
+            ],
+        )
+
+    @classmethod
+    def model_to_completions(self, resp, config):
+        # Response example (https://platform.openai.com/docs/api-reference/completions/create)
+        # ```
+        # {
+        #   "id": "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7",
+        #   "object": "text_completion",
+        #   "created": 1589478378,
+        #   "model": "text-davinci-003",
+        #   "choices": [
+        #     {
+        #       "text": "\n\nThis is indeed a test",
+        #       "index": 0,
+        #       "logprobs": null,
+        #       "finish_reason": "length"
+        #     }
+        #   ],
+        #   "usage": {
+        #     "prompt_tokens": 5,
+        #     "completion_tokens": 7,
+        #     "total_tokens": 12
+        #   }
+        # }
+        # ```
+        return completions.ResponsePayload(
+            id=resp["id"],
+            # The chat models response from OpenAI is of object type "chat.completion". Since
+            # we're using the completions response format here, we hardcode the "text_completion"
+            # object type in the response instead
+            object="text_completion",
+            created=resp["created"],
+            model=resp["model"],
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=c["message"]["content"],
+                    finish_reason=c["finish_reason"],
+                )
+                for idx, c in enumerate(resp["choices"])
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                completion_tokens=resp["usage"]["completion_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_completions_streaming(cls, resp, config):
+        return completions.StreamResponsePayload(
+            id=resp["id"],
+            # The chat models response from OpenAI is of object type "chat.completion.chunk".
+            # Since we're using the completions response format here, we hardcode the
+            # "text_completion_chunk" object type in the response instead
+            object="text_completion_chunk",
+            created=resp["created"],
+            model=resp["model"],
+            choices=[
+                completions.StreamChoice(
+                    index=c["index"],
+                    finish_reason=c["finish_reason"],
+                    text=c["delta"].get("content"),
+                )
+                for c in resp["choices"]
+            ],
+        )
+
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        # Response example (https://platform.openai.com/docs/api-reference/embeddings/create):
+        # ```
+        # {
+        #   "object": "list",
+        #   "data": [
+        #     {
+        #       "object": "embedding",
+        #       "embedding": [
+        #         0.0023064255,
+        #         -0.009327292,
+        #         .... (1536 floats total for ada-002)
+        #         -0.0028842222,
+        #       ],
+        #       "index": 0
+        #     }
+        #   ],
+        #   "model": "text-embedding-ada-002",
+        #   "usage": {
+        #     "prompt_tokens": 8,
+        #     "total_tokens": 8
+        #   }
+        # }
+        # ```
+        return embeddings.ResponsePayload(
+            data=[
+                embeddings.EmbeddingObject(
+                    embedding=d["embedding"],
+                    index=idx,
+                )
+                for idx, d in enumerate(resp["data"])
+            ],
+            model=resp["model"],
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+
+class OpenAIProvider(BaseProvider):
+    NAME = "OpenAI"
+    CONFIG_TYPE = OpenAIConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, OpenAIConfig):
+            # Should be unreachable
+            raise MlflowException.invalid_parameter_value(
+                "Invalid config type {config.model.config}"
+            )
+        self.openai_config: OpenAIConfig = config.model.config
+
+    @property
+    def base_url(self):
+        api_type = self.openai_config.openai_api_type
+        if api_type == OpenAIAPIType.OPENAI:
+            base_url = self.openai_config.openai_api_base or "https://api.openai.com/v1"
+            if (api_version := self.openai_config.openai_api_version) is not None:
+                return append_to_uri_query_params(base_url, ("api-version", api_version))
+            else:
+                return base_url
+        elif api_type in (OpenAIAPIType.AZURE, OpenAIAPIType.AZUREAD):
+            openai_url = append_to_uri_path(
+                self.openai_config.openai_api_base,
+                "openai",
+                "deployments",
+                self.openai_config.openai_deployment_name,
+            )
+            return append_to_uri_query_params(
+                openai_url,
+                ("api-version", self.openai_config.openai_api_version),
+            )
+        else:
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid OpenAI API type '{self.openai_config.openai_api_type}'"
+            )
+
+    @property
+    def headers(self):
+        api_type = self.openai_config.openai_api_type
+        if api_type == OpenAIAPIType.OPENAI:
+            headers = {
+                "Authorization": f"Bearer {self.openai_config.openai_api_key}",
+            }
+            if org := self.openai_config.openai_organization:
+                headers["OpenAI-Organization"] = org
+            return headers
+        elif api_type == OpenAIAPIType.AZUREAD:
+            return {
+                "Authorization": f"Bearer {self.openai_config.openai_api_key}",
+            }
+        elif api_type == OpenAIAPIType.AZURE:
+            return {
+                "api-key": self.openai_config.openai_api_key,
+            }
+        else:
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid OpenAI API type '{self.openai_config.openai_api_type}'"
+            )
+
+    @property
+    def adapter_class(self):
+        return OpenAIAdapter
+
+    def get_endpoint_url(self, route_type: str) -> str:
+        if route_type == "llm/v1/chat":
+            route_path = "chat/completions"
+        elif route_type == "llm/v1/completions":
+            route_path = "completions"
+        elif route_type == "llm/v1/embeddings":
+            route_path = "embeddings"
+        else:
+            raise ValueError(f"Invalid route type {route_type}")
+
+        # Append the route path to the base URL. Note that we cannot simply append the route path
+        # at the end of the base URL because it has query parameters for the Azure OpenAI case.
+        parsed_base_url = urlparse(self.base_url)
+        return urlunparse(parsed_base_url._replace(path=f"{parsed_base_url.path}/{route_path}"))
+
+    async def chat_stream(
+        self, payload: chat.RequestPayload
+    ) -> AsyncIterable[chat.StreamResponsePayload]:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+
+        stream = send_stream_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="chat/completions",
+            payload=self.adapter_class.chat_to_model(payload, self.config),
+        )
+
+        async for chunk in handle_incomplete_chunks(stream):
+            chunk = chunk.strip()
+            if not chunk:
+                continue
+
+            data = strip_sse_prefix(chunk.decode("utf-8"))
+            if data == "[DONE]":
+                return
+
+            resp = json.loads(data)
+            yield OpenAIAdapter.model_to_chat_streaming(resp, self.config)
+
+    async def _chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+
+        return await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="chat/completions",
+            payload=self.adapter_class.chat_to_model(payload, self.config),
+        )
+
+    async def _chat_uc_function(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        workspace_client = _get_workspace_client()
+        warehouse_id = os.environ.get("DATABRICKS_WAREHOUSE_ID")
+        if warehouse_id is None:
+            raise AIGatewayException(
+                status_code=400,
+                detail="DATABRICKS_WAREHOUSE_ID environment variable is not set",
+            )
+
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+
+        token_usage_accumulator = TokenUsageAccumulator()
+        user_tool_messages = [m for m in payload["messages"] if m["role"] == "tool"]
+        user_tool_calls = next(
+            (m["tool_calls"] for m in payload["messages"] if "tool_calls" in m), None
+        )
+        if (
+            user_tool_messages
+            and user_tool_calls
+            and (result := parse_uc_functions(payload["messages"][0]["content"]))
+        ):
+            uc_func_calls, uc_func_messages = result
+            messages = [
+                *[m for m in payload["messages"] if m["role"] == "tool" or "tool_calls" in m],
+                # Join UC function calls and user tool calls
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": uc_func_calls + user_tool_calls,
+                },
+                *uc_func_messages,
+                *user_tool_messages,
+            ]
+            resp = await send_request(
+                headers=self.headers,
+                base_url=self.base_url,
+                path="chat/completions",
+                payload=self.adapter_class.chat_to_model(
+                    {
+                        **payload,
+                        "messages": messages,
+                    },
+                    self.config,
+                ),
+            )
+            token_usage_accumulator.update(resp.get("usage", {}))
+        elif any(t["type"] == _UC_FUNCTION for t in payload.get("tools", [])):
+            updated_tools = []
+            uc_func_mapping: dict[str, "FunctionInfo"] = {}
+            for tool in payload.get("tools", []):
+                if tool["type"] == _UC_FUNCTION:
+                    function_name = tool[_UC_FUNCTION]["name"]
+                    function = workspace_client.functions.get(function_name)
+                    param_metadata = get_func_schema(function)
+                    t = {
+                        "type": "function",
+                        "function": param_metadata,
+                    }
+                    uc_func_mapping[t["function"]["name"]] = function
+                    updated_tools.append(t)
+                else:
+                    updated_tools.append(tool)
+
+            payload["tools"] = updated_tools
+
+            messages = payload.pop("messages", [])
+            uc_func_calls = []
+            user_tool_calls = []
+            resp = None
+            for _ in range(20):  # loop until we get a response without tool_calls
+                resp = await send_request(
+                    headers=self.headers,
+                    base_url=self.base_url,
+                    path="chat/completions",
+                    payload=self.adapter_class.chat_to_model(
+                        {
+                            **payload,
+                            "messages": messages,
+                        },
+                        self.config,
+                    ),
+                )
+                token_usage_accumulator.update(resp.get("usage", {}))
+                # TODO to support n > 1.
+                assistant_msg = resp["choices"][0]["message"]
+                tool_calls = assistant_msg.get("tool_calls")
+                if tool_calls is None:
+                    if uc_func_calls:
+                        original_content = resp["choices"][0]["message"]["content"]
+                        resp["choices"][0]["message"]["content"] = prepend_uc_functions(
+                            original_content, uc_func_calls
+                        )
+
+                    if user_tool_calls:
+                        # Is this line unreachable?
+                        resp["choices"][0]["message"]["tool_calls"] = user_tool_calls
+
+                    break
+
+                tool_messages = []
+                for tool_call in tool_calls:
+                    func = tool_call["function"]
+                    parameters = json.loads(func["arguments"])
+                    if func_info := uc_func_mapping.get(func["name"]):
+                        result = execute_function(
+                            ws=workspace_client,
+                            warehouse_id=warehouse_id,
+                            function=function,
+                            parameters=parameters,
+                        )
+                        tool_messages.append(
+                            {
+                                "role": "tool",
+                                "tool_call_id": tool_call["id"],
+                                "content": result.to_json(),
+                            }
+                        )
+
+                        uc_func_calls.append(
+                            (
+                                {
+                                    "id": tool_call["id"],
+                                    "name": func_info.full_name,
+                                    "arguments": func["arguments"],
+                                },
+                                {
+                                    "tool_call_id": tool_call["id"],
+                                    "content": result.to_json(),
+                                },
+                            )
+                        )
+                    else:
+                        user_tool_calls.append(
+                            {
+                                "id": tool_call["id"],
+                                "type": "function",
+                                "function": {
+                                    "name": func["name"],
+                                    "arguments": func["arguments"],
+                                },
+                            }
+                        )
+
+                if message_content := assistant_msg.pop("content", None):
+                    messages.append({"role": "assistant", "content": message_content})
+                messages += [assistant_msg, *tool_messages]
+
+                if user_tool_calls:
+                    # We can't go on without a response from the user, so we break here
+                    if uc_func_calls:
+                        resp["choices"][0]["message"]["content"] = join_uc_functions(uc_func_calls)
+
+                    resp["choices"][0]["message"]["tool_calls"] = user_tool_calls
+                    break
+            else:
+                raise AIGatewayException(
+                    status_code=500,
+                    detail="Max iterations reached",
+                )
+        else:
+            # No UC functions to execute
+            resp = await send_request(
+                headers=self.headers,
+                base_url=self.base_url,
+                path="chat/completions",
+                payload=self.adapter_class.chat_to_model(payload, self.config),
+            )
+            token_usage_accumulator.update(resp.get("usage", {}))
+
+        # Update the token usage
+        resp["usage"].update(token_usage_accumulator.dict())
+
+        return resp
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        if MLFLOW_ENABLE_UC_FUNCTIONS.get():
+            resp = await self._chat_uc_function(payload)
+        else:
+            resp = await self._chat(payload)
+
+        return OpenAIAdapter.model_to_chat(resp, self.config)
+
+    async def completions_stream(
+        self, payload: completions.RequestPayload
+    ) -> AsyncIterable[completions.StreamResponsePayload]:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        stream = send_stream_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="completions",
+            payload=OpenAIAdapter.completion_to_model(payload, self.config),
+        )
+
+        async for chunk in handle_incomplete_chunks(stream):
+            chunk = chunk.strip()
+            if not chunk:
+                continue
+
+            data = strip_sse_prefix(chunk.decode("utf-8"))
+            if data == "[DONE]":
+                return
+
+            resp = json.loads(data)
+            yield OpenAIAdapter.model_to_completions_streaming(resp, self.config)
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="completions",
+            payload=OpenAIAdapter.completion_to_model(payload, self.config),
+        )
+        return OpenAIAdapter.model_to_completions(resp, self.config)
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        resp = await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path="embeddings",
+            payload=OpenAIAdapter.embeddings_to_model(payload, self.config),
+        )
+        return OpenAIAdapter.model_to_embeddings(resp, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/palm.py

@@ -0,0 +1,219 @@
+import time
+from typing import Any
+
+from mlflow.gateway.config import PaLMConfig, RouteConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider
+from mlflow.gateway.providers.utils import rename_payload_keys, send_request
+from mlflow.gateway.schemas import chat, completions, embeddings
+
+
+class PaLMProvider(BaseProvider):
+    NAME = "PaLM"
+    CONFIG_TYPE = PaLMConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, PaLMConfig):
+            raise TypeError(f"Unexpected config type {config.model.config}")
+        self.palm_config: PaLMConfig = config.model.config
+
+    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        headers = {"x-goog-api-key": self.palm_config.palm_api_key}
+        return await send_request(
+            headers=headers,
+            base_url="https://generativelanguage.googleapis.com/v1beta3/models/",
+            path=path,
+            payload=payload,
+        )
+
+    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        if "max_tokens" in payload or "maxOutputTokens" in payload:
+            raise AIGatewayException(
+                status_code=422, detail="Max tokens is not supported for PaLM chat."
+            )
+        key_mapping = {
+            "stop": "stopSequences",
+            "n": "candidateCount",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        payload = rename_payload_keys(payload, key_mapping)
+        # The range of PaLM's temperature is 0-1, but ours is 0-2, so we halve it
+        payload["temperature"] = 0.5 * payload["temperature"]
+
+        # Replace 'role' with 'author' in payload
+        for m in payload["messages"]:
+            m["author"] = m.pop("role")
+
+        # Map 'messages', 'examples, and 'context' to 'prompt'
+        prompt = {"messages": payload.pop("messages")}
+        if "examples" in payload:
+            prompt["examples"] = payload.pop("examples")
+        if "context" in payload:
+            prompt["context"] = payload.pop("context")
+        payload["prompt"] = prompt
+
+        resp = await self._request(
+            f"{self.config.model.name}:generateMessage",
+            payload,
+        )
+
+        # Response example
+        # (https://developers.generativeai.google/api/rest/generativelanguage/models/generateMessage)
+        # ```
+        # {
+        #   "candidates": [
+        #     {
+        #       "author": "1",
+        #       "content": "Hi there! How can I help you today?"
+        #     }
+        #   ],
+        #   "messages": [
+        #     {
+        #       "author": "0",
+        #       "content": "hi"
+        #     }
+        #   ]
+        # }
+        # ```
+        return chat.ResponsePayload(
+            created=int(time.time()),
+            model=self.config.model.name,
+            choices=[
+                chat.Choice(
+                    index=idx,
+                    message=chat.ResponseMessage(role=c["author"], content=c["content"]),
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["candidates"])
+            ],
+            usage=chat.ChatUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        key_mapping = {
+            "stop": "stopSequences",
+            "n": "candidateCount",
+            "max_tokens": "maxOutputTokens",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        payload = rename_payload_keys(payload, key_mapping)
+        # The range of PaLM's temperature is 0-1, but ours is 0-2, so we halve it
+        payload["temperature"] = 0.5 * payload["temperature"]
+        payload["prompt"] = {"text": payload["prompt"]}
+        resp = await self._request(
+            f"{self.config.model.name}:generateText",
+            payload,
+        )
+        # Response example (https://developers.generativeai.google/api/rest/generativelanguage/models/generateText)
+        # ```
+        # {
+        #   "candidates": [
+        #     {
+        #       "output": "Once upon a time, there was a young girl named Lily...",
+        #       "safetyRatings": [
+        #         {
+        #           "category": "HARM_CATEGORY_DEROGATORY",
+        #           "probability": "NEGLIGIBLE"
+        #         }, ...
+        #       ]
+        #     {
+        #       "output": "Once upon a time, there was a young boy named Billy...",
+        #       "safetyRatings": [
+        #           ...
+        #       ]
+        #     }
+        #   ]
+        # }
+        # ```
+        return completions.ResponsePayload(
+            created=int(time.time()),
+            object="text_completion",
+            model=self.config.model.name,
+            choices=[
+                completions.Choice(
+                    index=idx,
+                    text=c["output"],
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["candidates"])
+            ],
+            usage=completions.CompletionsUsage(
+                prompt_tokens=None,
+                completion_tokens=None,
+                total_tokens=None,
+            ),
+        )
+
+    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+        self.check_for_model_field(payload)
+        key_mapping = {
+            "input": "texts",
+        }
+        for k1, k2 in key_mapping.items():
+            if k2 in payload:
+                raise AIGatewayException(
+                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
+                )
+        payload = rename_payload_keys(payload, key_mapping)
+
+        resp = await self._request(
+            f"{self.config.model.name}:batchEmbedText",
+            payload,
+        )
+
+        # Batch-text response example (https://developers.generativeai.google/api/rest/generativelanguage/models/batchEmbedText):
+        # ```
+        # {
+        #   "embeddings": [
+        #     {
+        #       "value": [
+        #         3.25,
+        #         0.7685547,
+        #         2.65625,
+        #         ...
+        #         -0.30126953,
+        #         -2.3554688,
+        #         1.2597656
+        #       ]
+        #     }
+        #   ]
+        # }
+        # ```
+        return embeddings.ResponsePayload(
+            data=[
+                embeddings.EmbeddingObject(
+                    embedding=embedding["value"],
+                    index=idx,
+                )
+                for idx, embedding in enumerate(resp["embeddings"])
+            ],
+            model=self.config.model.name,
+            usage=embeddings.EmbeddingsUsage(
+                prompt_tokens=None,
+                total_tokens=None,
+            ),
+        )

venv/lib/python3.9/site-packages/mlflow/gateway/providers/togetherai.py

@@ -0,0 +1,448 @@
+import json
+from typing import Any, AsyncGenerator, AsyncIterable
+
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.config import RouteConfig, TogetherAIConfig
+from mlflow.gateway.exceptions import AIGatewayException
+from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
+from mlflow.gateway.providers.utils import rename_payload_keys, send_request, send_stream_request
+from mlflow.gateway.schemas import chat as chat_schema
+from mlflow.gateway.schemas import completions as completions_schema
+from mlflow.gateway.schemas import embeddings as embeddings_schema
+from mlflow.gateway.utils import strip_sse_prefix
+
+
+class TogetherAIAdapter(ProviderAdapter):
+    @classmethod
+    def model_to_embeddings(cls, resp, config):
+        # Response example: (https://docs.together.ai/docs/embeddings-rest)
+        # ```
+        # {
+        #  "object": "list",
+        #  "data": [
+        #    {
+        #      "object": "embedding",
+        #      "embedding": [
+        #        0.44990748,
+        #        -0.2521129,
+        #        ...
+        #        -0.43091708,
+        #        0.214978
+        #      ],
+        #      "index": 0
+        #    }
+        #  ],
+        #  "model": "togethercomputer/m2-bert-80M-8k-retrieval",
+        #  "request_id": "840fc1b5bb2830cb-SEA"
+        # }
+        # ```
+        return embeddings_schema.ResponsePayload(
+            data=[
+                embeddings_schema.EmbeddingObject(
+                    embedding=item["embedding"],
+                    index=item["index"],
+                )
+                for item in resp["data"]
+            ],
+            model=config.model.name,
+            usage=embeddings_schema.EmbeddingsUsage(prompt_tokens=None, total_tokens=None),
+        )
+
+    @classmethod
+    def model_to_completions(cls, resp, config):
+        # Example response (https://docs.together.ai/reference/completions):
+        # {
+        #  "id": "8447f286bbdb67b3-SJC",
+        #  "choices": [
+        #    {
+        #      "text": "Example text."
+        #    }
+        #  ],
+        #  "usage": {
+        #    "prompt_tokens": 16,
+        #    "completion_tokens": 78,
+        #    "total_tokens": 94
+        #  },
+        #  "created": 1705089226,
+        #  "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        #  "object": "text_completion"
+        # }
+
+        return completions_schema.ResponsePayload(
+            id=resp["id"],
+            created=resp["created"],
+            model=config.model.name,
+            choices=[
+                completions_schema.Choice(
+                    index=idx,
+                    text=c["text"],
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["choices"])
+            ],
+            usage=completions_schema.CompletionsUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                completion_tokens=resp["usage"]["completion_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_completions_streaming(cls, resp, config):
+        # Response example (after manually calling API):
+        #
+        # {'id': '86d8d6e06df86f61-ATH', 'object': 'completion.chunk',
+        # 'created': 1711977238, 'choices': [{'index': 0, 'text': '   ',
+        # 'logprobs': None, 'finish_reason': None, 'delta': {'token_id': 2287, 'content': '   '}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1', 'usage': None}
+        #
+        # {'id': '86d8d6e06df86f61-ATH', 'object': 'completion.chunk',
+        # 'created': 1711977238, 'choices': [{'index': 0, 'text': ' "', 'logprobs': None,
+        # 'finish_reason': None, 'delta': {'token_id': 345, 'content': ' "'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1', 'usage': None}
+        #
+        # "{'id': '86d8d6e06df86f61-ATH', 'object': 'completion.chunk',
+        # 'created': 1711977238, 'choices': [{'index': 0, 'text': 'name', 'logprobs': None,
+        # 'finish_reason': None, 'delta': {'token_id': 861, 'content': 'name'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1', 'usage': None}
+        #
+        # LAST CHUNK
+        # {'id': '86d8d6e06df86f61-ATH', 'object': 'completion.chunk',
+        # 'created': 1711977238, 'choices': [{'index': 0, 'text': '":', 'logprobs': None,
+        # 'finish_reason': 'length', 'delta': {'token_id': 1264, 'content': '":'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1',
+        # 'usage': {'prompt_tokens': 17, 'completion_tokens': 200, 'total_tokens': 217}}
+        # ":[DONE]
+
+        return completions_schema.StreamResponsePayload(
+            id=resp.get("id"),
+            created=resp.get("created"),
+            model=config.model.name,
+            choices=[
+                completions_schema.StreamChoice(
+                    index=idx,
+                    # TODO this is questionable since the finish reason comes from togetherai api
+                    finish_reason=choice.get("finish_reason"),
+                    text=choice.get("text"),
+                )
+                for idx, choice in enumerate(resp.get("choices", []))
+            ],
+            # usage is not included in OpenAI StreamResponsePayload
+        )
+
+    @classmethod
+    def completions_to_model(cls, payload, config):
+        key_mapping = {
+            # TogetherAI uses logprobs
+            # OpenAI uses top_logprobs
+            "top_logprobs": "logprobs"
+        }
+
+        # in openAI API the logprobs parameter
+        # is a boolean flag.
+        # Insert this here to prevent the user from mixing up the APIs
+        logprobs_in_payload_condition = "logprobs" in payload and not isinstance(
+            payload["logprobs"], int
+        )
+
+        if logprobs_in_payload_condition:
+            raise AIGatewayException(
+                status_code=422,
+                detail="Wrong type for logprobs. It should be an 32bit integer.",
+            )
+
+        openai_top_logprobs_in_payload_condition = "top_logprobs" in payload and not isinstance(
+            payload["top_logprobs"], int
+        )
+
+        if openai_top_logprobs_in_payload_condition:
+            raise AIGatewayException(
+                status_code=422,
+                detail="Wrong type for top_logprobs. It should a 32bit integer.",
+            )
+
+        payload = rename_payload_keys(payload, key_mapping)
+        return {"model": config.model.name, **payload}
+
+    @classmethod
+    def completions_streaming_to_model(cls, payload, config):
+        # parameters for streaming completions are the same as the standard completions
+        return TogetherAIAdapter.completions_to_model(payload, config)
+
+    @classmethod
+    def model_to_chat(cls, resp, config):
+        # Example response (https://docs.together.ai/reference/chat-completions):
+        # {
+        #   "id": "8448080b880415ea-SJC",
+        #   "choices": [
+        #    {
+        #        "message": {
+        #           "role": "assistant",
+        #           "content": "example"
+        #         }
+        #     }
+        #   ],
+        #   "usage": {
+        #     "prompt_tokens": 31,
+        #     "completion_tokens": 455,
+        #     "total_tokens": 486
+        #   },
+        #   "created": 1705090115,
+        #   "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        #   "object": "chat.completion"
+        # }
+        return chat_schema.ResponsePayload(
+            id=resp["id"],
+            object="chat.completion",
+            created=resp["created"],
+            model=config.model.name,
+            choices=[
+                chat_schema.Choice(
+                    index=idx,
+                    message=chat_schema.ResponseMessage(
+                        role="assistant",
+                        content=c["message"]["content"],
+                    ),
+                    finish_reason=None,
+                )
+                for idx, c in enumerate(resp["choices"])
+            ],
+            usage=chat_schema.ChatUsage(
+                prompt_tokens=resp["usage"]["prompt_tokens"],
+                completion_tokens=resp["usage"]["completion_tokens"],
+                total_tokens=resp["usage"]["total_tokens"],
+            ),
+        )
+
+    @classmethod
+    def model_to_chat_streaming(cls, resp, config):
+        # Response example (after running API manually):
+        #
+        # {'id': '86f2cfd18f6b38ca-ATH', 'object': 'chat.completion.chunk',
+        # 'created': 1712249578, 'choices': [{'index': 0, 'text': ' The', 'logprobs': None,
+        # 'finish_reason': None, 'delta': {'token_id': 415, 'content': ' The'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1', 'usage': None}
+        #
+        # {'id': '86f2cfd18f6b38ca-ATH', 'object': 'chat.completion.chunk',
+        # 'created': 1712249578, 'choices': [{'index': 0, 'text': ' City', 'logprobs': None,
+        # 'finish_reason': None, 'delta': {'token_id': 3805, 'content': ' City'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1', 'usage': None}
+        #
+        # {'id': '86f2cfd18f6b38ca-ATH', 'object': 'chat.completion.chunk',
+        # 'created': 1712249578, 'choices': [{'index': 0, 'text': ' of', 'logprobs': None,
+        # 'finish_reason': None, 'delta': {'token_id': 302, 'content': ' of'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1', 'usage': None}
+        #
+        # LAST CHUNK
+        # {'id': '86f2cfd18f6b38ca-ATH', 'object': 'chat.completion.chunk',
+        # 'created': 1712249578, 'choices': [{'index': 0, 'text': ' Paris', 'logprobs': None,
+        # 'finish_reason': 'length', 'delta': {'token_id': 5465, 'content': ' Paris'}}],
+        # 'model': 'mistralai/Mixtral-8x7B-v0.1',
+        # 'usage': {'prompt_tokens': 93, 'completion_tokens': 100, 'total_tokens': 193}}
+
+        return chat_schema.StreamResponsePayload(
+            id=resp["id"],
+            model=config.model.name,
+            object="chat.completion.chunk",
+            created=resp["created"],
+            choices=[
+                chat_schema.StreamChoice(
+                    index=idx,
+                    finish_reason=choice.get("finish_reason"),
+                    delta=chat_schema.StreamDelta(
+                        role=None,
+                        content=choice.get("text"),
+                    ),
+                )
+                # Added enumerate and a default empty list
+                for idx, choice in enumerate(resp.get("choices", []))
+            ],
+            usage=resp.get("usage"),
+        )
+
+    @classmethod
+    def chat_to_model(cls, payload, config):
+        # completions and chat endpoint contain the same parameters
+        return TogetherAIAdapter.completions_to_model(payload, config)
+
+    @classmethod
+    def chat_streaming_to_model(cls, payload, config):
+        # streaming and standard chat contain the same parameters
+        return TogetherAIAdapter.chat_to_model(payload, config)
+
+    @classmethod
+    def embeddings_to_model(cls, payload, config):
+        # Example request (https://docs.together.ai/reference/embeddings):
+        # curl --request POST \
+        #   --url https://api.together.xyz/v1/embeddings \
+        #   --header 'accept: application/json' \
+        #   --header 'content-type: application/json' \
+        #   --data '
+        #   {
+        #     "model": "togethercomputer/m2-bert-80M-8k-retrieval",
+        #     "input": "Our solar system orbits the Milky Way galaxy at about 515,000 mph"
+        #   }
+
+        # This is just to keep the interface consistent the adapter
+        # class is not needed here as the togetherai request similar
+        # to the openAI one.
+
+        return payload
+
+
+class TogetherAIProvider(BaseProvider):
+    NAME = "TogetherAI"
+    CONFIG_TYPE = TogetherAIConfig
+
+    def __init__(self, config: RouteConfig) -> None:
+        super().__init__(config)
+        if config.model.config is None or not isinstance(config.model.config, TogetherAIConfig):
+            # Should be unreachable
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid config type {config.model.config}"
+            )
+        self.togetherai_config: TogetherAIConfig = config.model.config
+
+    @property
+    def base_url(self):
+        # togetherai seems to support only this url
+        return "https://api.together.xyz/v1"
+
+    @property
+    def headers(self):
+        return {"Authorization": f"Bearer {self.togetherai_config.togetherai_api_key}"}
+
+    @property
+    def adapter_class(self) -> type[ProviderAdapter]:
+        return TogetherAIAdapter
+
+    def get_endpoint_url(self, route_type: str) -> str:
+        if route_type == "llm/v1/chat":
+            return f"{self.base_url}/chat/completions"
+        elif route_type == "llm/v1/completions":
+            return f"{self.base_url}/completions"
+        elif route_type == "llm/v1/embeddings":
+            return f"{self.base_url}/embeddings"
+        else:
+            raise ValueError(f"Invalid route type {route_type}")
+
+    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        return await send_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path=path,
+            payload=payload,
+        )
+
+    async def _stream_request(
+        self, path: str, payload: dict[str, Any]
+    ) -> AsyncGenerator[bytes, None]:
+        return send_stream_request(
+            headers=self.headers,
+            base_url=self.base_url,
+            path=path,
+            payload=payload,
+        )
+
+    async def embeddings(
+        self, payload: embeddings_schema.RequestPayload
+    ) -> embeddings_schema.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+
+        resp = await self._request(
+            path="embeddings",
+            payload=TogetherAIAdapter.embeddings_to_model(payload, self.config),
+        )
+
+        return TogetherAIAdapter.model_to_embeddings(resp, self.config)
+
+    async def completions_stream(
+        self, payload: completions_schema.RequestPayload
+    ) -> AsyncIterable[completions_schema.StreamResponsePayload]:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+
+        if not payload.get("max_tokens"):
+            raise AIGatewayException(
+                status_code=422,
+                detail=(
+                    "max_tokens is not present in payload."
+                    "It is a required parameter for TogetherAI completions."
+                ),
+            )
+
+        stream = await self._stream_request(
+            path="completions",
+            payload=TogetherAIAdapter.completions_streaming_to_model(payload, self.config),
+        )
+
+        async for chunk in stream:
+            chunk = chunk.strip()
+            if not chunk:
+                continue
+
+            chunk = strip_sse_prefix(chunk.decode("utf-8"))
+            if chunk == "[DONE]":
+                return
+
+            resp = json.loads(chunk)
+            yield TogetherAIAdapter.model_to_completions_streaming(resp, self.config)
+
+    async def completions(
+        self, payload: completions_schema.RequestPayload
+    ) -> completions_schema.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+
+        if not payload.get("max_tokens"):
+            raise AIGatewayException(
+                status_code=422,
+                detail=(
+                    "max_tokens is not present in payload."
+                    "It is a required parameter for TogetherAI completions."
+                ),
+            )
+
+        resp = await self._request(
+            path="completions", payload=TogetherAIAdapter.completions_to_model(payload, self.config)
+        )
+
+        return TogetherAIAdapter.model_to_completions(resp, self.config)
+
+    async def chat_stream(self, payload: chat_schema.RequestPayload) -> chat_schema.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+
+        stream = await self._stream_request(
+            path="chat/completions",
+            payload=TogetherAIAdapter.chat_streaming_to_model(payload, self.config),
+        )
+
+        async for chunk in stream:
+            chunk = chunk.strip()
+            if not chunk:
+                continue
+
+            chunk = strip_sse_prefix(chunk.decode("utf-8"))
+            if chunk == "[DONE]":
+                return
+
+            resp = json.loads(chunk)
+            yield TogetherAIAdapter.model_to_chat_streaming(resp, self.config)
+
+    async def chat(self, payload: chat_schema.RequestPayload) -> chat_schema.ResponsePayload:
+        from fastapi.encoders import jsonable_encoder
+
+        payload = jsonable_encoder(payload, exclude_none=True)
+
+        resp = await self._request(
+            path="chat/completions",
+            payload=TogetherAIAdapter.chat_to_model(payload, self.config),
+        )
+
+        return TogetherAIAdapter.model_to_chat(resp, self.config)

venv/lib/python3.9/site-packages/mlflow/gateway/providers/utils.py

@@ -0,0 +1,95 @@
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator
+
+import aiohttp
+
+from mlflow.gateway.constants import (
+    MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS,
+)
+from mlflow.utils.uri import append_to_uri_path
+
+
+@asynccontextmanager
+async def _aiohttp_post(headers: dict[str, str], base_url: str, path: str, payload: dict[str, Any]):
+    async with aiohttp.ClientSession(headers=headers) as session:
+        url = append_to_uri_path(base_url, path)
+        timeout = aiohttp.ClientTimeout(total=MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS)
+        async with session.post(url, json=payload, timeout=timeout) as response:
+            yield response
+
+
+async def send_request(headers: dict[str, str], base_url: str, path: str, payload: dict[str, Any]):
+    """
+    Send an HTTP request to a specific URL path with given headers and payload.
+
+    Args:
+        headers: The headers to include in the request.
+        base_url: The base URL where the request will be sent.
+        path: The specific path of the URL to which the request will be sent.
+        payload: The payload (or data) to be included in the request.
+
+    Returns:
+        The server's response as a JSON object.
+
+    Raises:
+        HTTPException if the HTTP request fails.
+    """
+    from fastapi import HTTPException
+
+    async with _aiohttp_post(headers, base_url, path, payload) as response:
+        content_type = response.headers.get("Content-Type")
+        if content_type and "application/json" in content_type:
+            js = await response.json()
+        elif content_type and "text/plain" in content_type:
+            js = {"message": await response.text()}
+        else:
+            raise HTTPException(
+                status_code=502,
+                detail=f"The returned data type from the route service is not supported. "
+                f"Received content type: {content_type}",
+            )
+        try:
+            response.raise_for_status()
+        except aiohttp.ClientResponseError as e:
+            detail = js.get("error", {}).get("message", e.message) if "error" in js else js
+            raise HTTPException(status_code=e.status, detail=detail)
+        return js
+
+
+async def send_stream_request(
+    headers: dict[str, str], base_url: str, path: str, payload: dict[str, Any]
+) -> AsyncGenerator[bytes, None]:
+    """
+    Send an HTTP request to a specific URL path with given headers and payload.
+
+    Args:
+        headers: The headers to include in the request.
+        base_url: The base URL where the request will be sent.
+        path: The specific path of the URL to which the request will be sent.
+        payload: The payload (or data) to be included in the request.
+
+    Returns:
+        The server's response as a JSON object.
+
+    Raises:
+        HTTPException if the HTTP request fails.
+    """
+    async with _aiohttp_post(headers, base_url, path, payload) as response:
+        async for line in response.content:
+            yield line
+
+
+def rename_payload_keys(payload: dict[str, Any], mapping: dict[str, str]) -> dict[str, Any]:
+    """Rename payload keys based on the specified mapping. If a key is not present in the
+    mapping, the key and its value will remain unchanged.
+
+    Args:
+        payload: The original dictionary to transform.
+        mapping: A dictionary where each key-value pair represents a mapping from the old
+            key to the new key.
+
+    Returns:
+        A new dictionary containing the transformed keys.
+
+    """
+    return {mapping.get(k, k): v for k, v in payload.items()}

venv/lib/python3.9/site-packages/mlflow/gateway/runner.py

@@ -0,0 +1,106 @@
+import logging
+import os
+import subprocess
+import sys
+from typing import Generator
+
+from watchfiles import watch
+
+from mlflow.environment_variables import MLFLOW_GATEWAY_CONFIG
+from mlflow.gateway import app
+from mlflow.gateway.config import _load_route_config
+from mlflow.gateway.utils import kill_child_processes
+
+_logger = logging.getLogger(__name__)
+
+
+def monitor_config(config_path: str) -> Generator[None, None, None]:
+    with open(config_path) as f:
+        prev_config = f.read()
+
+    for changes in watch(os.path.dirname(config_path)):
+        if not any((path == config_path) for _, path in changes):
+            continue
+
+        if not os.path.exists(config_path):
+            _logger.warning(f"{config_path} deleted")
+            continue
+
+        with open(config_path) as f:
+            config = f.read()
+        if config == prev_config:
+            continue
+
+        try:
+            _load_route_config(config_path)
+        except Exception as e:
+            _logger.warning("Invalid configuration: %s", e)
+            continue
+        else:
+            prev_config = config
+
+        yield
+
+
+class Runner:
+    def __init__(
+        self,
+        config_path: str,
+        host: str,
+        port: int,
+        workers: int,
+    ) -> None:
+        self.config_path = config_path
+        self.host = host
+        self.port = port
+        self.workers = workers
+        self.process = None
+
+    def start(self) -> None:
+        self.process = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "gunicorn",
+                "--bind",
+                f"{self.host}:{self.port}",
+                "--workers",
+                str(self.workers),
+                "--worker-class",
+                "uvicorn.workers.UvicornWorker",
+                f"{app.__name__}:create_app_from_env()",
+            ],
+            env={
+                **os.environ,
+                MLFLOW_GATEWAY_CONFIG.name: self.config_path,
+            },
+        )
+
+    def stop(self) -> None:
+        if self.process is not None:
+            self.process.terminate()
+            self.process.wait()
+            self.process = None
+
+    def reload(self) -> None:
+        kill_child_processes(self.process.pid)
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+
+
+def run_app(config_path: str, host: str, port: int, workers: int) -> None:
+    config_path = os.path.abspath(os.path.normpath(os.path.expanduser(config_path)))
+    with Runner(
+        config_path=config_path,
+        host=host,
+        port=port,
+        workers=workers,
+    ) as runner:
+        for _ in monitor_config(config_path):
+            _logger.info("Configuration updated, reloading workers")
+            runner.reload()

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/__init__.py

@@ -0,0 +1,3 @@
+from mlflow.gateway.schemas import chat, completions, embeddings
+
+__all__ = ["chat", "completions", "embeddings"]

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/chat.py

@@ -0,0 +1,142 @@
+"""
+This module defines the schemas for the MLflow AI Gateway's chat endpoint.
+
+The schemas must be compatible with OpenAI's Chat Completion API.
+https://platform.openai.com/docs/api-reference/chat
+
+NB: These Pydantic models just alias the models defined in mlflow.types.chat to avoid code
+    duplication, but with the addition of RequestModel and ResponseModel base classes.
+"""
+
+from typing import Literal, Optional
+
+from pydantic import Field
+
+from mlflow.gateway.base_models import RequestModel, ResponseModel
+
+# Import marked with noqa is for backward compatibility
+from mlflow.types.chat import (
+    ChatChoice,
+    ChatChoiceDelta,
+    ChatChunkChoice,
+    ChatCompletionChunk,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatMessage,
+    ChatUsage,  # noqa F401
+    Function,  # noqa F401
+    FunctionToolDefinition,
+    ToolCall,  # noqa F401
+)
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+# NB: `import x as y` does not work and will cause a Pydantic error.
+StreamDelta = ChatChoiceDelta
+StreamChoice = ChatChunkChoice
+RequestMessage = ChatMessage
+
+
+class UnityCatalogFunctionToolDefinition(RequestModel):
+    name: str
+
+
+class ChatToolWithUC(RequestModel):
+    """
+    A tool definition for the chat endpoint with Unity Catalog integration.
+    The Gateway request accepts a special tool type 'uc_function' for Unity Catalog integration.
+    https://mlflow.org/docs/latest/llms/deployments/uc_integration.html
+    """
+
+    type: Literal["function", "uc_function"]
+    function: Optional[FunctionToolDefinition] = None
+    uc_function: Optional[UnityCatalogFunctionToolDefinition] = None
+
+
+_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
+    "messages": [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Hello!"},
+    ],
+    "temperature": 0.0,
+    "max_tokens": 64,
+    "stop": ["END"],
+    "n": 1,
+    "stream": False,
+}
+
+
+class RequestPayload(ChatCompletionRequest, RequestModel):
+    messages: list[RequestMessage] = Field(..., min_items=1)
+    tools: Optional[list[ChatToolWithUC]] = None
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+
+
+_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "3cdb958c-e4cc-4834-b52b-1d1a7f324714",
+        "object": "chat.completion",
+        "created": 1700173217,
+        "model": "llama-2-70b-chat-hf",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hello! I am an AI assistant"},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
+    }
+}
+
+
+class ResponseMessage(ChatMessage, ResponseModel):
+    # Override the `tool_call_id` field to be excluded from the response.
+    # This is a band-aid solution to avoid exposing the tool_call_id in the response,
+    # while we use the same ChatMessage model for both request and response.
+    tool_call_id: Optional[str] = Field(None, exclude=True)
+
+
+class Choice(ChatChoice, ResponseModel):
+    # Override the `message` field to use the ResponseMessage model.
+    message: ResponseMessage
+
+
+class ResponsePayload(ChatCompletionResponse, ResponseModel):
+    # Override the `choices` field to use the Choice model
+    choices: list[Choice]
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+
+
+_STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "3cdb958c-e4cc-4834-b52b-1d1a7f324714",
+        "object": "chat.completion",
+        "created": 1700173217,
+        "model": "llama-2-70b-chat-hf",
+        "choices": [
+            {
+                "index": 6,
+                "finish_reason": "stop",
+                "delta": {"role": "assistant", "content": "you?"},
+            }
+        ],
+    }
+}
+
+
+class StreamResponsePayload(ChatCompletionChunk, ResponseModel):
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/completions.py

@@ -0,0 +1,109 @@
+from typing import Optional
+
+from mlflow.gateway.base_models import RequestModel, ResponseModel
+from mlflow.types.chat import BaseRequestPayload
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "prompt": "hello",
+        "temperature": 0.0,
+        "max_tokens": 64,
+        "stop": ["END"],
+        "n": 1,
+    }
+}
+
+
+class RequestPayload(BaseRequestPayload, RequestModel):
+    prompt: str
+    model: Optional[str] = None
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+
+
+class Choice(ResponseModel):
+    index: int
+    text: str
+    finish_reason: Optional[str] = None
+
+
+class CompletionsUsage(ResponseModel):
+    prompt_tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+
+
+_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "cmpl-123",
+        "object": "text_completion",
+        "created": 1589478378,
+        "model": "gpt-4",
+        "choices": [
+            {"text": "Hello! I am an AI Assistant!", "index": 0, "finish_reason": "length"}
+        ],
+        "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12},
+    }
+}
+
+
+class ResponsePayload(ResponseModel):
+    id: Optional[str] = None
+    object: str = "text_completion"
+    created: int
+    model: str
+    choices: list[Choice]
+    usage: CompletionsUsage
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+
+
+class StreamDelta(ResponseModel):
+    role: Optional[str] = None
+    content: Optional[str] = None
+
+
+class StreamChoice(ResponseModel):
+    index: int
+    finish_reason: Optional[str] = None
+    text: Optional[str] = None
+
+
+_STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "id": "cmpl-123",
+        "object": "text_completion",
+        "created": 1589478378,
+        "model": "gpt-4",
+        "choices": [
+            {
+                "index": 6,
+                "finish_reason": "stop",
+                "delta": {"role": "assistant", "content": "you?"},
+            }
+        ],
+    }
+}
+
+
+class StreamResponsePayload(ResponseModel):
+    id: Optional[str] = None
+    object: str = "text_completion_chunk"
+    created: int
+    model: str
+    choices: list[StreamChoice]
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _STREAM_RESPONSE_PAYLOAD_EXTRA_SCHEMA

venv/lib/python3.9/site-packages/mlflow/gateway/schemas/embeddings.py

@@ -0,0 +1,97 @@
+from typing import Optional, Union
+
+from mlflow.gateway.base_models import RequestModel, ResponseModel
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+_REQUEST_PAYLOAD_EXTRA_SCHEMA = {
+    "example": {
+        "input": ["hello", "world"],
+    }
+}
+
+
+class RequestPayload(RequestModel):
+    input: Union[str, list[str], list[int], list[list[int]]]
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _REQUEST_PAYLOAD_EXTRA_SCHEMA
+
+
+class EmbeddingObject(ResponseModel):
+    object: str = "embedding"
+    embedding: Union[list[float], str]
+    index: int
+
+
+class EmbeddingsUsage(ResponseModel):
+    prompt_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+
+
+_RESPONSE_PAYLOAD_EXTRA_SCHEMA = {
+    "object": "list",
+    "data": [
+        {
+            "object": "embedding",
+            "index": 0,
+            "embedding": [
+                0.017291732,
+                -0.017291732,
+                0.014577783,
+                -0.02902633,
+                -0.037271563,
+                0.019333655,
+                -0.023055641,
+                -0.007359971,
+                -0.015818445,
+                -0.030654699,
+                0.008348623,
+                0.018312693,
+                -0.017149571,
+                -0.0044424757,
+                -0.011165961,
+                0.01018377,
+            ],
+        },
+        {
+            "object": "embedding",
+            "index": 1,
+            "embedding": [
+                0.0060126893,
+                -0.008691099,
+                -0.0040095365,
+                0.019889368,
+                0.036211833,
+                -0.0013270887,
+                0.013401738,
+                -0.0036735237,
+                -0.0049594184,
+                0.035229642,
+                -0.03435084,
+                0.019798903,
+                -0.0006110424,
+                0.0073793563,
+                0.005657291,
+                0.022487005,
+            ],
+        },
+    ],
+    "model": "text-embedding-ada-002-v2",
+    "usage": {"prompt_tokens": 400, "total_tokens": 400},
+}
+
+
+class ResponsePayload(ResponseModel):
+    object: str = "list"
+    data: list[EmbeddingObject]
+    model: str
+    usage: EmbeddingsUsage
+
+    class Config:
+        if IS_PYDANTIC_V2_OR_NEWER:
+            json_schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA
+        else:
+            schema_extra = _RESPONSE_PAYLOAD_EXTRA_SCHEMA

venv/lib/python3.9/site-packages/mlflow/gateway/uc_function_utils.py

@@ -0,0 +1,321 @@
+# TODO: Move this in mlflow/gateway/utils/uc_functions.py
+
+import json
+import re
+from dataclasses import dataclass
+from io import StringIO
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+    from databricks.sdk.service.catalog import FunctionInfo, FunctionParameterInfo
+    from databricks.sdk.service.sql import StatementParameterListItem
+
+
+_UC_FUNCTION = "uc_function"
+
+
+def uc_type_to_json_schema_type(uc_type_json: Union[str, dict[str, Any]]) -> dict[str, Any]:
+    """
+    Converts the JSON representation of a Unity Catalog data type to the corresponding JSON schema
+    type. The conversion is lossy because we do not need to convert it back.
+    """
+    # See https://docs.databricks.com/en/sql/language-manual/sql-ref-datatypes.html
+    # The actual type name in type_json is different from the corresponding SQL type name.
+    spark_struct_field_mapping = {
+        "long": {"type": "integer"},
+        "binary": {"type": "string"},
+        "boolean": {"type": "boolean"},
+        "date": {"type": "string", "format": "date"},
+        "double": {"type": "number"},
+        "float": {"type": "number"},
+        "integer": {"type": "integer"},
+        "void": {"type": "null"},
+        "short": {"type": "integer"},
+        "string": {"type": "string"},
+        "timestamp": {"type": "string", "format": "date-time"},
+        "timestamp_ntz": {"type": "string", "format": "date-time"},
+        "byte": {"type": "integer"},
+    }
+    if isinstance(uc_type_json, str):
+        if t := spark_struct_field_mapping.get(uc_type_json):
+            return t
+        else:
+            if uc_type_json.startswith("decimal"):
+                return {"type": "number"}
+            elif uc_type_json.startswith("interval"):
+                raise TypeError(f"Type {uc_type_json} is not supported.")
+            else:
+                raise TypeError(f"Unknown type {uc_type_json}. Try upgrading this package.")
+    else:
+        assert isinstance(uc_type_json, dict)
+        type = uc_type_json["type"]
+        if type == "array":
+            element_type = uc_type_to_json_schema_type(uc_type_json["elementType"])
+            return {"type": "array", "items": element_type}
+        elif type == "map":
+            key_type = uc_type_json["keyType"]
+            if key_type != "string":
+                raise TypeError(f"Only support STRING key type for MAP but got {key_type}.")
+            value_type = uc_type_to_json_schema_type(uc_type_json["valueType"])
+            return {
+                "type": "object",
+                "additionalProperties": value_type,
+            }
+        elif type == "struct":
+            properties = {}
+            for field in uc_type_json["fields"]:
+                properties[field["name"]] = uc_type_to_json_schema_type(field["type"])
+            return {"type": "object", "properties": properties}
+        else:
+            raise TypeError(f"Unknown type {uc_type_json}. Try upgrading this package.")
+
+
+def extract_param_metadata(p: "FunctionParameterInfo") -> dict:
+    type_json = json.loads(p.type_json)["type"]
+    json_schema_type = uc_type_to_json_schema_type(type_json)
+    json_schema_type["name"] = p.name
+    json_schema_type["description"] = (
+        (p.comment or "") + f" (default: {p.parameter_default})" if p.parameter_default else ""
+    )
+    return json_schema_type
+
+
+def get_func_schema(func: "FunctionInfo") -> dict[str, Any]:
+    parameters = func.input_params.parameters if func.input_params else []
+    return {
+        "description": func.comment,
+        "name": _get_tool_name(func),
+        "parameters": {
+            "type": "object",
+            "properties": {p.name: extract_param_metadata(p) for p in parameters},
+            "required": [p.name for p in parameters if p.parameter_default is None],
+        },
+    }
+
+
+@dataclass
+class ParameterizedStatement:
+    statement: str
+    parameters: list["StatementParameterListItem"]
+
+
+@dataclass
+class FunctionExecutionResult:
+    """
+    Result of executing a function.
+    We always use a string to present the result value for AI model to consume.
+    """
+
+    error: Optional[str] = None
+    format: Optional[Literal["SCALAR", "CSV"]] = None
+    value: Optional[str] = None
+    truncated: Optional[bool] = None
+
+    def to_json(self) -> str:
+        data = {k: v for (k, v) in self.__dict__.items() if v is not None}
+        return json.dumps(data)
+
+
+def is_scalar(function: "FunctionInfo") -> bool:
+    """
+    Returns True if the function returns a single row instead of a table.
+    """
+    from databricks.sdk.service.catalog import ColumnTypeName
+
+    return function.data_type != ColumnTypeName.TABLE_TYPE
+
+
+def get_execute_function_sql_stmt(
+    function: "FunctionInfo",
+    json_params: dict[str, Any],
+) -> ParameterizedStatement:
+    from databricks.sdk.service.catalog import ColumnTypeName
+    from databricks.sdk.service.sql import StatementParameterListItem
+
+    parts = []
+    output_params = []
+    if is_scalar(function):
+        parts.append(f"SELECT {function.full_name}(")
+    else:
+        parts.append(f"SELECT * FROM {function.full_name}(")
+    if function.input_params is None or function.input_params.parameters is None:
+        assert not json_params, "Function has no parameters but parameters were provided."
+    else:
+        args = []
+        use_named_args = False
+        for p in function.input_params.parameters:
+            if p.name not in json_params:
+                if p.parameter_default is not None:
+                    use_named_args = True
+                else:
+                    raise ValueError(f"Parameter {p.name} is required but not provided.")
+            else:
+                arg_clause = ""
+                if use_named_args:
+                    arg_clause += f"{p.name} => "
+                json_value = json_params[p.name]
+                if p.type_name in (
+                    ColumnTypeName.ARRAY,
+                    ColumnTypeName.MAP,
+                    ColumnTypeName.STRUCT,
+                ):
+                    # Use from_json to restore values of complex types.
+                    json_value_str = json.dumps(json_value)
+                    # TODO: parametrize type
+                    arg_clause += f"from_json(:{p.name}, '{p.type_text}')"
+                    output_params.append(
+                        StatementParameterListItem(name=p.name, value=json_value_str)
+                    )
+                elif p.type_name == ColumnTypeName.BINARY:
+                    # Use ubbase64 to restore binary values.
+                    arg_clause += f"unbase64(:{p.name})"
+                    output_params.append(StatementParameterListItem(name=p.name, value=json_value))
+                else:
+                    arg_clause += f":{p.name}"
+                    output_params.append(
+                        StatementParameterListItem(name=p.name, value=json_value, type=p.type_text)
+                    )
+                args.append(arg_clause)
+        parts.append(",".join(args))
+    parts.append(")")
+    # TODO: check extra params in kwargs
+    statement = "".join(parts)
+    return ParameterizedStatement(statement=statement, parameters=output_params)
+
+
+def execute_function(
+    ws: "WorkspaceClient",
+    warehouse_id: str,
+    function: "FunctionInfo",
+    parameters: dict[str, Any],
+) -> FunctionExecutionResult:
+    """
+    Execute a function with the given arguments and return the result.
+    """
+    try:
+        import pandas as pd
+    except ImportError as e:
+        raise ImportError(
+            "Could not import pandas python package. Please install it with `pip install pandas`."
+        ) from e
+    from databricks.sdk.service.sql import StatementState
+
+    # TODO: async so we can run functions in parallel
+    parameterized_statement = get_execute_function_sql_stmt(function, parameters)
+    # TODO: make limits and wait timeout configurable
+    response = ws.statement_execution.execute_statement(
+        statement=parameterized_statement.statement,
+        warehouse_id=warehouse_id,
+        parameters=parameterized_statement.parameters,
+        wait_timeout="30s",
+        row_limit=100,
+        byte_limit=4096,
+    )
+    status = response.status
+    assert status is not None, f"Statement execution failed: {response}"
+    if status.state != StatementState.SUCCEEDED:
+        error = status.error
+        assert error is not None, "Statement execution failed but no error message was provided."
+        return FunctionExecutionResult(error=f"{error.error_code}: {error.message}")
+    manifest = response.manifest
+    assert manifest is not None
+    truncated = manifest.truncated
+    result = response.result
+    assert result is not None, "Statement execution succeeded but no result was provided."
+    data_array = result.data_array
+    if is_scalar(function):
+        value = None
+        if data_array and len(data_array) > 0 and len(data_array[0]) > 0:
+            value = str(data_array[0][0])  # type: ignore
+        return FunctionExecutionResult(format="SCALAR", value=value, truncated=truncated)
+    else:
+        schema = manifest.schema
+        assert schema is not None and schema.columns is not None, (
+            "Statement execution succeeded but no schema was provided."
+        )
+        columns = [c.name for c in schema.columns]
+        if data_array is None:
+            data_array = []
+        pdf = pd.DataFrame.from_records(data_array, columns=columns)
+        csv_buffer = StringIO()
+        pdf.to_csv(csv_buffer, index=False)
+        return FunctionExecutionResult(
+            format="CSV", value=csv_buffer.getvalue(), truncated=truncated
+        )
+
+
+def join_uc_functions(uc_functions: list[dict[str, Any]]):
+    calls = [
+        f"""
+<uc_function_call>
+{json.dumps(request, indent=2)}
+</uc_function_call>
+
+<uc_function_result>
+{json.dumps(result, indent=2)}
+</uc_function_result>
+""".strip()
+        for (request, result) in uc_functions
+    ]
+    return "\n\n".join(calls)
+
+
+def _get_tool_name(function: "FunctionInfo") -> str:
+    # The maximum function name length OpenAI supports is 64 characters.
+    return f"{function.catalog_name}__{function.schema_name}__{function.name}"[-64:]
+
+
+@dataclass
+class ParseResult:
+    tool_calls: list[dict[str, Any]]
+    tool_messages: list[dict[str, Any]]
+
+
+_UC_REGEX = re.compile(
+    r"""
+<uc_function_call>
+(?P<uc_function_call>.*?)
+</uc_function_call>
+
+<uc_function_result>
+(?P<uc_function_result>.*?)
+</uc_function_result>
+""",
+    re.DOTALL,
+)
+
+
+def parse_uc_functions(content) -> Optional[ParseResult]:
+    tool_calls = []
+    tool_messages = []
+    for m in _UC_REGEX.finditer(content):
+        c = m.group("uc_function_call")
+        g = m.group("uc_function_result")
+        tool_calls.append(json.loads(c))
+        tool_messages.append(json.loads(g))
+
+    return ParseResult(tool_calls, tool_messages) if tool_calls else None
+
+
+@dataclass
+class TokenUsageAccumulator:
+    prompt_tokens: int = 0
+    completions_tokens: int = 0
+    total_tokens: int = 0
+
+    def update(self, usage_dict):
+        self.prompt_tokens += usage_dict.get("prompt_tokens", 0)
+        self.completions_tokens += usage_dict.get("completion_tokens", 0)
+        self.total_tokens += usage_dict.get("total_tokens", 0)
+
+    def dict(self):
+        return {
+            "prompt_tokens": self.prompt_tokens,
+            "completion_tokens": self.completions_tokens,
+            "total_tokens": self.total_tokens,
+        }
+
+
+def prepend_uc_functions(content, uc_functions):
+    return join_uc_functions(uc_functions) + "\n\n" + content

venv/lib/python3.9/site-packages/mlflow/gateway/utils.py

@@ -0,0 +1,311 @@
+import base64
+import functools
+import inspect
+import json
+import logging
+import posixpath
+import re
+import textwrap
+import warnings
+from typing import Any, AsyncGenerator, Optional
+from urllib.parse import urlparse
+
+from mlflow.environment_variables import MLFLOW_GATEWAY_URI
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.constants import MLFLOW_AI_GATEWAY_MOSAICML_CHAT_SUPPORTED_MODEL_PREFIXES
+from mlflow.utils.uri import append_to_uri_path
+
+_logger = logging.getLogger(__name__)
+_gateway_uri: Optional[str] = None
+
+
+def is_valid_endpoint_name(name: str) -> bool:
+    """
+    Check whether a string contains any URL reserved characters, spaces, or characters other
+    than alphanumeric, underscore, hyphen, and dot.
+
+    Returns True if the string doesn't contain any of these characters.
+    """
+    return bool(re.fullmatch(r"[\w\-\.]+", name))
+
+
+def check_configuration_route_name_collisions(config):
+    routes = config.get("routes") or config.get("endpoints") or []
+    if len(routes) < 2:
+        return
+    names = [route["name"] for route in routes]
+    if len(names) != len(set(names)):
+        raise MlflowException.invalid_parameter_value(
+            "Duplicate names found in endpoint configurations. Please remove the duplicate endpoint"
+            " name from the configuration to ensure that endpoints are created properly."
+        )
+
+
+def check_configuration_deprecated_fields(config):
+    if "routes" in config:
+        warnings.warn(
+            "The 'routes' configuration key has been deprecated and will be removed in an"
+            " upcoming release. Use 'endpoints' instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+
+    routes = config.get("routes", []) or config.get("endpoints", [])
+    for route in routes:
+        if "route_type" in route:
+            warnings.warn(
+                "The 'route_type' configuration key has been deprecated and will be removed in an"
+                " upcoming release. Use 'endpoint_type' instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            break
+
+
+def kill_child_processes(parent_pid):
+    """
+    Gracefully terminate or kill child processes from a main process
+    """
+    import psutil
+
+    parent = psutil.Process(parent_pid)
+    for child in parent.children(recursive=True):
+        try:
+            child.terminate()
+        except psutil.NoSuchProcess:
+            pass
+    _, still_alive = psutil.wait_procs(parent.children(), timeout=3)
+    for p in still_alive:
+        p.kill()
+
+
+def _is_valid_uri(uri: str):
+    """
+    Evaluates the basic structure of a provided gateway uri to determine if the scheme and
+    netloc are provided
+    """
+    if uri == "databricks":
+        return True
+    try:
+        parsed = urlparse(uri)
+        return parsed.scheme == "databricks" or all([parsed.scheme, parsed.netloc])
+    except ValueError:
+        return False
+
+
+def _get_indent(s: str) -> str:
+    for l in s.splitlines():
+        if l.startswith(" "):
+            return " " * (len(l) - len(l.lstrip()))
+    return ""
+
+
+def _prepend(docstring: Optional[str], text: str) -> str:
+    if not docstring:
+        return text
+
+    indent = _get_indent(docstring)
+    return f"""
+{textwrap.indent(text, indent)}
+
+{docstring}
+"""
+
+
+def gateway_deprecated(obj):
+    msg = (
+        "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."
+    )
+    warning = f"""
+.. warning::
+
+    {msg}
+""".strip()
+    if inspect.isclass(obj):
+        original = obj.__init__
+
+        @functools.wraps(original)
+        def wrapper(*args, **kwargs):
+            warnings.warn(msg, FutureWarning, stacklevel=2)
+            return original(*args, **kwargs)
+
+        obj.__init__ = wrapper
+        obj.__init__.__doc__ = _prepend(obj.__init__.__doc__, warning)
+        return obj
+    else:
+
+        @functools.wraps(obj)
+        def wrapper(*args, **kwargs):
+            warnings.warn(msg, FutureWarning, stacklevel=2)
+            return obj(*args, **kwargs)
+
+        wrapper.__doc__ = _prepend(obj.__doc__, warning)
+
+        return wrapper
+
+
+@gateway_deprecated
+def set_gateway_uri(gateway_uri: str):
+    """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".
+    """
+    if not _is_valid_uri(gateway_uri):
+        raise MlflowException.invalid_parameter_value(
+            "The gateway uri provided is missing required elements. Ensure that the schema "
+            "and netloc are provided."
+        )
+
+    global _gateway_uri
+    _gateway_uri = gateway_uri
+
+
+@gateway_deprecated
+def get_gateway_uri() -> str:
+    """
+    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.
+    """
+    if _gateway_uri is not None:
+        return _gateway_uri
+    elif uri := MLFLOW_GATEWAY_URI.get():
+        return uri
+    else:
+        raise MlflowException(
+            "No Gateway server uri has been set. Please either set the MLflow Gateway URI via "
+            "`mlflow.gateway.set_gateway_uri()` or set the environment variable "
+            f"{MLFLOW_GATEWAY_URI} to the running Gateway API server's uri"
+        )
+
+
+def assemble_uri_path(paths: list[str]) -> str:
+    """Assemble a correct URI path from a list of path parts.
+
+    Args:
+        paths: A list of strings representing parts of a URI path.
+
+    Returns:
+        A string representing the complete assembled URI path.
+
+    """
+    stripped_paths = [path.strip("/").lstrip("/") for path in paths if path]
+    return "/" + posixpath.join(*stripped_paths) if stripped_paths else "/"
+
+
+def resolve_route_url(base_url: str, route: str) -> str:
+    """
+    Performs a validation on whether the returned value is a fully qualified url (as the case
+    with Databricks) or requires the assembly of a fully qualified url by appending the
+    Route return route_url to the base url of the AI Gateway server.
+
+    Args:
+        base_url: The base URL. Should include the scheme and domain, e.g.,
+            ``http://127.0.0.1:6000``.
+        route: The route to be appended to the base URL, e.g., ``/api/2.0/gateway/routes/`` or,
+            in the case of Databricks, the fully qualified url.
+
+    Returns:
+        The complete URL, either directly returned or formed and returned by joining the
+        base URL and the route path.
+    """
+    return route if _is_valid_uri(route) else append_to_uri_path(base_url, route)
+
+
+class SearchRoutesToken:
+    def __init__(self, index: int):
+        self._index = index
+
+    @property
+    def index(self):
+        return self._index
+
+    @classmethod
+    def decode(cls, encoded_token: str):
+        try:
+            decoded_token = base64.b64decode(encoded_token)
+            parsed_token = json.loads(decoded_token)
+            index = int(parsed_token.get("index"))
+        except Exception as e:
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid SearchRoutes token: {encoded_token}. The index is not defined as a "
+                "value that can be represented as a positive integer."
+            ) from e
+
+        if index < 0:
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid SearchRoutes token: {encoded_token}. The index cannot be negative."
+            )
+
+        return cls(index=index)
+
+    def encode(self) -> str:
+        token_json = json.dumps(
+            {
+                "index": self.index,
+            }
+        )
+        encoded_token_bytes = base64.b64encode(bytes(token_json, "utf-8"))
+        return encoded_token_bytes.decode("utf-8")
+
+
+def is_valid_mosiacml_chat_model(model_name: str) -> bool:
+    return any(
+        model_name.lower().startswith(supported)
+        for supported in MLFLOW_AI_GATEWAY_MOSAICML_CHAT_SUPPORTED_MODEL_PREFIXES
+    )
+
+
+def is_valid_ai21labs_model(model_name: str) -> bool:
+    return model_name in {"j2-ultra", "j2-mid", "j2-light"}
+
+
+def strip_sse_prefix(s: str) -> str:
+    # https://html.spec.whatwg.org/multipage/server-sent-events.html
+    return re.sub(r"^data:\s+", "", s)
+
+
+def to_sse_chunk(data: str) -> str:
+    # https://html.spec.whatwg.org/multipage/server-sent-events.html
+    return f"data: {data}\n\n"
+
+
+def _find_boundary(buffer: bytes) -> int:
+    try:
+        return buffer.index(b"\n")
+    except ValueError:
+        return -1
+
+
+async def handle_incomplete_chunks(
+    stream: AsyncGenerator[bytes, Any],
+) -> AsyncGenerator[bytes, Any]:
+    """
+    Wraps a streaming response and handles incomplete chunks from the server.
+    See https://community.openai.com/t/incomplete-stream-chunks-for-completions-api/383520
+    for more information.
+    """
+    buffer = b""
+    async for chunk in stream:
+        buffer += chunk
+        while (boundary := _find_boundary(buffer)) != -1:
+            yield buffer[:boundary]
+            buffer = buffer[boundary + 1 :]
+
+
+async def make_streaming_response(resp):
+    from starlette.responses import StreamingResponse
+
+    if isinstance(resp, AsyncGenerator):
+        return StreamingResponse(
+            (to_sse_chunk(d.json()) async for d in resp),
+            media_type="text/event-stream",
+        )
+    else:
+        return await resp