star

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