mlflow--mlflow/mlflow/server/gateway_api.py

"""
Database-backed Gateway API endpoints for MLflow Server.

This module provides dynamic gateway endpoints that are configured from the database
rather than from a static YAML configuration file. It integrates the AI Gateway
functionality directly into the MLflow tracking server.
"""

import functools
import logging
import sys
import time
from collections.abc import AsyncIterable, Callable
from typing import Any

from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse

from mlflow.entities.gateway_endpoint import GatewayModelLinkageType
from mlflow.exceptions import MlflowException
from mlflow.gateway.budget import check_budget_limit, make_budget_on_complete
from mlflow.gateway.config import (
    AmazonBedrockConfig,
    AnthropicConfig,
    EndpointConfig,
    EndpointType,
    GatewayRequestType,
    GeminiConfig,
    LiteLLMConfig,
    MistralConfig,
    OpenAIAPIType,
    OpenAIConfig,
    Provider,
    VertexAIConfig,
    _AuthConfigKey,
    _OpenAICompatibleConfig,
)
from mlflow.gateway.constants import MLFLOW_GATEWAY_CALLER_HEADER, GatewayCaller
from mlflow.gateway.guardrail_utils import (
    extract_auth_headers,
    load_guardrails,
    run_post_llm_guardrails,
    run_post_llm_guardrails_passthrough,
    run_pre_llm_guardrails,
)
from mlflow.gateway.guardrails import (
    _SANITIZE_BYPASS_HEADER,
    GuardrailViolation,
    JudgeGuardrail,
)
from mlflow.gateway.providers import get_provider
from mlflow.gateway.providers.base import (
    PASSTHROUGH_ROUTES,
    BaseProvider,
    FallbackProvider,
    PassthroughAction,
    TrafficRouteProvider,
)
from mlflow.gateway.providers.utils import provider_call_duration_ms
from mlflow.gateway.schemas import chat, embeddings
from mlflow.gateway.tracing_utils import (
    aggregate_anthropic_messages_stream_chunks,
    aggregate_chat_stream_chunks,
    aggregate_gemini_stream_generate_content_chunks,
    aggregate_openai_responses_stream_chunks,
    maybe_traced_gateway_call,
)
from mlflow.gateway.utils import safe_stream, to_sse_chunk, translate_http_exception
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
from mlflow.store.tracking.abstract_store import AbstractStore
from mlflow.store.tracking.gateway.config_resolver import get_endpoint_config
from mlflow.store.tracking.gateway.entities import (
    GatewayEndpointConfig,
    GatewayModelConfig,
    RoutingStrategy,
)
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
from mlflow.telemetry.events import GatewayInvocationEvent, GatewayInvocationType
from mlflow.telemetry.track import _record_event
from mlflow.tracing.constant import TraceMetadataKey
from mlflow.tracking._tracking_service.utils import _get_store
from mlflow.types.chat import ChatCompletionRequest
from mlflow.utils.provider_filter import is_provider_allowed, normalize_provider_name
from mlflow.utils.workspace_context import get_request_workspace

_logger = logging.getLogger(__name__)

gateway_router = APIRouter(prefix="/gateway", tags=["gateway"])


async def _get_request_body(request: Request) -> dict[str, Any]:
    """
    Get request body, using cached version if available.

    The auth middleware may have already parsed the request body for permission
    validation. Since Starlette request body can only be read once, we cache
    the parsed body in request.state.cached_body for reuse by route handlers.

    Args:
        request: The FastAPI Request object.

    Returns:
        Parsed JSON body as a dictionary.

    Raises:
        HTTPException: If the request body is not valid JSON.
    """
    # Check if body was already parsed by auth middleware
    cached_body = getattr(request.state, "cached_body", None)
    if isinstance(cached_body, dict):
        return cached_body

    # Otherwise parse it now
    try:
        return await request.json()
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid JSON payload: {e!s}")


def _get_user_metadata(request: Request) -> dict[str, Any]:
    """
    Extract user metadata from request state for tracing.

    The auth middleware stores the authenticated user's info in request.state.

    Args:
        request: The FastAPI Request object.

    Returns:
        Dictionary with user metadata (username and user_id if available).
    """
    metadata = {}
    if username := getattr(request.state, "username", None):
        metadata[TraceMetadataKey.AUTH_USERNAME] = username
    if user_id := getattr(request.state, "user_id", None):
        metadata[TraceMetadataKey.AUTH_USER_ID] = str(user_id)
    return metadata


def _record_gateway_invocation(invocation_type: GatewayInvocationType) -> Callable[..., Any]:
    """
    Decorator for gateway invocation endpoints that records telemetry:
    success/failure status, duration, streaming mode, and caller.

    As a side effect, relays provider call duration to the gateway timing middleware by
    writing `request.state.gateway_provider_duration_ms`. This is required because
    Starlette's call_next() copies the ContextVar context for the handler task, so
    mutations to provider_call_duration_ms don't propagate back to the middleware.

    Timing headers (X-MLflow-Gateway-Duration-Ms, X-MLflow-Gateway-Overhead-Duration-Ms)
    are injected by gateway_timing_middleware in fastapi_app.py.

    Args:
        invocation_type: The type of invocation endpoint.
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            start_time = time.perf_counter()
            success = True
            result = None

            # Extract caller header from the Request object if present,
            # only accepting known caller values to avoid logging arbitrary input.
            caller = None
            request = next((a for a in (*args, *kwargs.values()) if isinstance(a, Request)), None)
            if request is not None:
                raw_caller = request.headers.get(MLFLOW_GATEWAY_CALLER_HEADER)
                if raw_caller in {e.value for e in GatewayCaller}:
                    caller = raw_caller

            try:
                result = await func(*args, **kwargs)
            except Exception:
                success = False
                raise
            finally:
                duration_ms = int((time.perf_counter() - start_time) * 1000)
                provider_duration = int(provider_call_duration_ms.get())
                is_streaming = isinstance(result, StreamingResponse)
                params = {
                    "is_streaming": is_streaming,
                    "invocation_type": invocation_type,
                }
                # provider_call_duration_ms is only updated by send_request()
                # (non-streaming); send_stream_request() never sets it, so
                # timing fields would always be 0 for streaming responses.
                if not is_streaming:
                    params["provider_duration_ms"] = provider_duration
                    params["gateway_overhead_ms"] = max(0, duration_ms - provider_duration)
                if caller:
                    params["caller"] = caller
                if request is not None:
                    params["has_traceparent"] = request.headers.get("traceparent") is not None
                    auth_mod = sys.modules.get("mlflow.server.auth")
                    params["auth_enabled"] = auth_mod.is_auth_enabled() if auth_mod else False
                    if endpoint_id := getattr(request.state, "endpoint_id", None):
                        params["endpoint_id"] = endpoint_id
                    # Prefer the actual provider from the response (set by
                    # BaseProvider after the call) over the endpoint config
                    # estimate, which may not reflect traffic-split/fallback.
                    actual_provider = getattr(result, "provider", None)
                    if provider := (actual_provider or getattr(request.state, "provider", None)):
                        params["provider"] = provider
                _record_event(
                    GatewayInvocationEvent,
                    params=params,
                    success=success,
                    duration_ms=duration_ms,
                )
                # Relay provider timing to the middleware via request.state.
                # ContextVar values set in the handler task don't propagate back
                # to the middleware task (Starlette copies the context for call_next).
                if request is not None:
                    request.state.gateway_provider_duration_ms = int(
                        provider_call_duration_ms.get()
                    )

            return result

        return wrapper

    return decorator


def _set_gateway_telemetry_state(request: Request, endpoint_config) -> None:
    """Set endpoint_id and provider on request.state for telemetry attribution."""
    request.state.endpoint_id = endpoint_config.endpoint_id
    if endpoint_config.models:
        primary_model = next(
            (
                m
                for m in endpoint_config.models
                if m.linkage_type == GatewayModelLinkageType.PRIMARY
            ),
            endpoint_config.models[0],
        )
        request.state.provider = str(primary_model.provider)


def _build_openai_compatible_config(model_config: "GatewayModelConfig"):
    """Build an _OpenAICompatibleConfig for providers that use the OpenAI API format."""
    auth_config = model_config.auth_config or {}
    return _OpenAICompatibleConfig(
        api_key=model_config.secret_value.get(_AuthConfigKey.API_KEY),
        api_base=auth_config.get(_AuthConfigKey.API_BASE),
    )


def _build_endpoint_config(
    endpoint_name: str,
    model_config: GatewayModelConfig,
    endpoint_type: EndpointType,
) -> EndpointConfig:
    """
    Build an EndpointConfig from model configuration.

    This function combines provider config building and endpoint config building
    into a single operation.

    Args:
        endpoint_name: The endpoint name.
        model_config: The model configuration object with decrypted secrets.
        endpoint_type: Endpoint type (chat or embeddings).

    Returns:
        EndpointConfig instance ready for provider instantiation.

    Raises:
        MlflowException: If provider configuration is invalid.
    """
    provider_name = model_config.provider
    if not is_provider_allowed(provider_name):
        _logger.debug(
            "Provider '%s' blocked by MLFLOW_GATEWAY_ALLOWED_PROVIDERS",
            provider_name,
        )
        raise MlflowException.invalid_parameter_value(
            f"Provider '{provider_name}' is not allowed by the current gateway provider policy."
        )

    provider_config = None

    if model_config.provider == Provider.OPENAI:
        auth_config = model_config.auth_config or {}
        openai_config = {
            "openai_api_key": model_config.secret_value.get(_AuthConfigKey.API_KEY),
        }

        # Check if this is Azure OpenAI (requires api_type, deployment_name, api_base, api_version)
        if "api_type" in auth_config and auth_config["api_type"] in ("azure", "azuread"):
            openai_config["openai_api_type"] = auth_config["api_type"]
            openai_config["openai_api_base"] = auth_config.get(_AuthConfigKey.API_BASE)
            openai_config["openai_deployment_name"] = auth_config.get("deployment_name")
            openai_config["openai_api_version"] = auth_config.get("api_version")
        else:
            # Standard OpenAI
            if _AuthConfigKey.API_BASE in auth_config:
                openai_config["openai_api_base"] = auth_config[_AuthConfigKey.API_BASE]
            if "organization" in auth_config:
                openai_config["openai_organization"] = auth_config["organization"]

        provider_config = OpenAIConfig(**openai_config)
    elif model_config.provider == Provider.AZURE:
        auth_config = model_config.auth_config or {}
        model_config.provider = Provider.OPENAI
        provider_config = OpenAIConfig(
            openai_api_type=OpenAIAPIType.AZURE,
            openai_api_key=model_config.secret_value.get(_AuthConfigKey.API_KEY),
            openai_api_base=auth_config.get(_AuthConfigKey.API_BASE),
            openai_deployment_name=model_config.model_name,
            openai_api_version=auth_config.get("api_version"),
        )
    elif model_config.provider == Provider.ANTHROPIC:
        anthropic_config = {
            "anthropic_api_key": model_config.secret_value.get(_AuthConfigKey.API_KEY),
        }
        if model_config.auth_config:
            if "version" in model_config.auth_config:
                anthropic_config["anthropic_version"] = model_config.auth_config["version"]
            if _AuthConfigKey.API_BASE in model_config.auth_config:
                anthropic_config["anthropic_api_base"] = model_config.auth_config[
                    _AuthConfigKey.API_BASE
                ]
        provider_config = AnthropicConfig(**anthropic_config)
    elif model_config.provider == Provider.MISTRAL:
        provider_config = MistralConfig(
            mistral_api_key=model_config.secret_value.get(_AuthConfigKey.API_KEY),
        )
    elif model_config.provider == Provider.GEMINI:
        provider_config = GeminiConfig(
            gemini_api_key=model_config.secret_value.get(_AuthConfigKey.API_KEY),
        )
    elif model_config.provider in {
        Provider.GROQ,
        Provider.DEEPSEEK,
        Provider.XAI,
        Provider.OPENROUTER,
        Provider.OLLAMA,
        Provider.PORTKEY,
    }:
        provider_config = _build_openai_compatible_config(model_config)
    elif normalize_provider_name(model_config.provider) == Provider.DATABRICKS:
        from mlflow.gateway.providers.databricks import DatabricksConfig

        auth_config = model_config.auth_config or {}
        auth_mode = auth_config.get(_AuthConfigKey.AUTH_MODE, "pat_token")
        config_kwargs = {}
        if api_base := auth_config.get(_AuthConfigKey.API_BASE):
            config_kwargs["host"] = api_base
        if auth_mode == "oauth_m2m":
            config_kwargs["client_id"] = auth_config.get("client_id")
            config_kwargs["client_secret"] = model_config.secret_value.get("client_secret")
        else:
            config_kwargs["token"] = model_config.secret_value.get(_AuthConfigKey.API_KEY)
        provider_config = DatabricksConfig(**config_kwargs)
        model_config.provider = Provider.DATABRICKS
    elif normalize_provider_name(model_config.provider) == Provider.BEDROCK:
        auth_config = model_config.auth_config or {}
        auth_mode = auth_config.get(_AuthConfigKey.AUTH_MODE, "api_key")
        if auth_mode == "api_key":
            # Bearer token auth — bypasses boto3 SigV4
            provider_config = AmazonBedrockConfig(
                aws_config={
                    "aws_bearer_token": model_config.secret_value.get(_AuthConfigKey.API_KEY),
                    "aws_region": auth_config.get("aws_region_name"),
                }
            )
        elif auth_mode == "access_keys":
            provider_config = AmazonBedrockConfig(
                aws_config={
                    "aws_access_key_id": model_config.secret_value.get("aws_access_key_id"),
                    "aws_secret_access_key": model_config.secret_value.get("aws_secret_access_key"),
                    "aws_region": auth_config.get("aws_region_name"),
                }
            )
        elif auth_mode == "iam_role":
            provider_config = AmazonBedrockConfig(
                aws_config={
                    "aws_role_arn": auth_config.get("aws_role_name"),
                    "aws_region": auth_config.get("aws_region_name"),
                }
            )
        else:
            # default_chain — boto3 resolves credentials from the
            # environment (env vars, ~/.aws/credentials, instance profile, etc.)
            aws_config = {"aws_region": auth_config.get("aws_region_name")}
            if role_arn := auth_config.get("aws_role_name"):
                aws_config["aws_role_arn"] = role_arn
            provider_config = AmazonBedrockConfig(aws_config=aws_config)
        model_config.provider = Provider.BEDROCK
    elif model_config.provider == Provider.VERTEX_AI:
        auth_config = model_config.auth_config or {}
        provider_config = VertexAIConfig(
            vertex_project=auth_config.get("vertex_project"),
            vertex_location=auth_config.get("vertex_location"),
            vertex_credentials=model_config.secret_value.get("vertex_credentials"),
        )
    else:
        # Use LiteLLM as fallback for unsupported providers
        # Store the original provider name for LiteLLM's provider/model format
        original_provider = model_config.provider
        auth_config = model_config.auth_config or {}
        # Merge auth_config with secret_value (secret_value contains api_key and other secrets)
        litellm_config = {
            "litellm_provider": original_provider,
            "litellm_auth_config": auth_config | model_config.secret_value,
        }
        provider_config = LiteLLMConfig(**litellm_config)
        model_config.provider = Provider.LITELLM

    # Build and return EndpointConfig
    return EndpointConfig(
        name=endpoint_name,
        endpoint_type=endpoint_type,
        model={
            "name": model_config.model_name,
            "provider": model_config.provider,
            "config": provider_config.model_dump(),
        },
    )


def _create_provider(
    endpoint_config: GatewayEndpointConfig,
    endpoint_type: EndpointType,
    enable_tracing: bool = False,
) -> BaseProvider:
    """
    Create a provider instance based on endpoint routing strategy.

    Fallback is independent of routing strategy - if fallback_config is present,
    the provider is wrapped with FallbackProvider.

    Args:
        endpoint_config: The endpoint configuration with model details and routing config.
        endpoint_type: Endpoint type (chat or embeddings).

    Returns:
        Provider instance (standard provider, TrafficRouteProvider, or FallbackProvider).

    Raises:
        MlflowException: If endpoint configuration is invalid or has no models.
    """
    # Get PRIMARY models
    primary_models = [
        model
        for model in endpoint_config.models
        if model.linkage_type == GatewayModelLinkageType.PRIMARY
    ]

    if not primary_models:
        raise MlflowException(
            f"Endpoint '{endpoint_config.endpoint_name}' has no PRIMARY models configured",
            error_code=RESOURCE_DOES_NOT_EXIST,
        )

    # Create base provider based on routing strategy
    if endpoint_config.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT:
        # Traffic split: distribute requests based on weights
        configs = []
        weights = []
        for model_config in primary_models:
            gateway_endpoint_config = _build_endpoint_config(
                endpoint_name=endpoint_config.endpoint_name,
                model_config=model_config,
                endpoint_type=endpoint_type,
            )
            configs.append(gateway_endpoint_config)
            weights.append(int(model_config.weight * 100))  # Convert to percentage

        primary_provider = TrafficRouteProvider(
            configs=configs,
            traffic_splits=weights,
            routing_strategy="TRAFFIC_SPLIT",
            enable_tracing=enable_tracing,
        )
    else:
        # Default: use the first PRIMARY model
        model_config = primary_models[0]
        gateway_endpoint_config = _build_endpoint_config(
            endpoint_config.endpoint_name, model_config, endpoint_type
        )
        provider_class = get_provider(model_config.provider)
        primary_provider = provider_class(gateway_endpoint_config, enable_tracing=enable_tracing)

    # Wrap with FallbackProvider if fallback configuration exists
    if endpoint_config.fallback_config:
        fallback_models = [
            model
            for model in endpoint_config.models
            if model.linkage_type == GatewayModelLinkageType.FALLBACK
        ]

        if not fallback_models:
            _logger.debug(
                f"Endpoint '{endpoint_config.endpoint_name}' has fallback_config "
                "but no FALLBACK models configured"
            )
            return primary_provider

        # Sort fallback models by fallback_order
        fallback_models.sort(
            key=lambda m: m.fallback_order if m.fallback_order is not None else float("inf")
        )

        fallback_providers = [
            get_provider(model_config.provider)(
                _build_endpoint_config(
                    endpoint_name=endpoint_config.endpoint_name,
                    model_config=model_config,
                    endpoint_type=endpoint_type,
                ),
                enable_tracing=enable_tracing,
            )
            for model_config in fallback_models
        ]

        max_attempts = endpoint_config.fallback_config.max_attempts or len(fallback_models)

        # FallbackProvider expects all providers (primary + fallback)
        all_providers = [primary_provider] + fallback_providers

        return FallbackProvider(
            providers=all_providers,
            max_attempts=max_attempts + 1,  # +1 to include primary
            strategy=endpoint_config.fallback_config.strategy,
            enable_tracing=enable_tracing,
        )

    return primary_provider


def _create_provider_from_endpoint_name(
    store: SqlAlchemyStore,
    endpoint_name: str,
    endpoint_type: EndpointType,
    enable_tracing: bool = True,
) -> tuple[BaseProvider, GatewayEndpointConfig]:
    """
    Create a provider from an endpoint name.

    Args:
        store: The SQLAlchemy store instance.
        endpoint_name: The endpoint name.
        endpoint_type: Endpoint type (chat or embeddings).
        enable_tracing: If True, enables MLflow tracing for provider calls.

    Returns:
        Tuple of (provider instance, endpoint config)
    """
    endpoint_config = get_endpoint_config(endpoint_name=endpoint_name, store=store)
    return _create_provider(
        endpoint_config, endpoint_type, enable_tracing=enable_tracing
    ), endpoint_config


def _validate_store(store: AbstractStore) -> None:
    if not isinstance(store, SqlAlchemyStore):
        raise HTTPException(
            status_code=500,
            detail="Gateway endpoints are only available with SqlAlchemyStore, "
            f"got {type(store).__name__}.",
        )


def _extract_endpoint_name_from_model(body: dict[str, Any]) -> str:
    """
    Extract and validate the endpoint name from the 'model' parameter in the request body.

    Args:
        body: The request body dictionary

    Returns:
        The endpoint name extracted from the 'model' parameter

    Raises:
        HTTPException: If the 'model' parameter is missing
    """
    endpoint_name = body.get("model")
    if not endpoint_name:
        raise HTTPException(
            status_code=400,
            detail="Missing required 'model' parameter in request body",
        )
    return endpoint_name


def _get_guardrails_and_auth(
    store, endpoint_config, request: Request
) -> tuple[list[JudgeGuardrail], dict[str, str]]:
    """Load guardrails and extract auth headers, skipping guardrails for internal bypass calls."""
    headers = dict(request.headers)
    bypass = headers.get(_SANITIZE_BYPASS_HEADER) == "1"
    guardrails = [] if bypass else load_guardrails(store, endpoint_config, request)
    return guardrails, extract_auth_headers(headers)


@gateway_router.post("/{endpoint_name}/mlflow/invocations", response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.MLFLOW_INVOCATIONS)
async def invocations(endpoint_name: str, request: Request):
    """
    Unified invocations endpoint handler that supports both chat and embeddings.

    The handler automatically detects the request type based on the payload structure:
    - If payload has "messages" field -> chat endpoint
    - If payload has "input" field -> embeddings endpoint
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)
    headers = dict(request.headers)

    store = _get_store()
    workspace = get_request_workspace()

    _validate_store(store)
    endpoint_config = get_endpoint_config(endpoint_name=endpoint_name, store=store)
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    # Detect request type based on payload structure
    if "messages" in body:
        # Chat request
        endpoint_type = EndpointType.LLM_V1_CHAT
        try:
            payload = chat.RequestPayload(**body)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid chat payload: {e!s}")

        provider, endpoint_config = _create_provider_from_endpoint_name(
            store, endpoint_name, endpoint_type
        )

        if payload.stream:
            # Post-LLM guardrails are not applied to streaming responses.
            # Pre-LLM guardrails run inside the trace as child spans; violations
            # are surfaced as SSE error chunks via safe_stream.
            async def _guarded_stream(
                payload: chat.RequestPayload,
            ):
                request_dict = await run_pre_llm_guardrails(
                    guardrails,
                    payload.model_dump(),
                    auth_headers=auth_headers,
                    usage_tracking=endpoint_config.usage_tracking,
                    payload_schema=ChatCompletionRequest.model_json_schema(),
                )
                async for chunk in provider.chat_stream(chat.RequestPayload(**request_dict)):
                    yield chunk

            stream = maybe_traced_gateway_call(
                _guarded_stream,
                endpoint_config,
                user_metadata,
                output_reducer=aggregate_chat_stream_chunks,
                request_headers=headers,
                request_type=GatewayRequestType.UNIFIED_CHAT,
                on_complete=make_budget_on_complete(store, workspace),
            )(payload)
            return StreamingResponse(
                safe_stream(to_sse_chunk(chunk.model_dump_json()) async for chunk in stream),
                media_type="text/event-stream",
            )
        else:

            async def _guarded_chat(
                payload: chat.RequestPayload,
            ) -> chat.ResponsePayload:
                request_dict = await run_pre_llm_guardrails(
                    guardrails,
                    payload.model_dump(),
                    auth_headers=auth_headers,
                    usage_tracking=endpoint_config.usage_tracking,
                    payload_schema=ChatCompletionRequest.model_json_schema(),
                )
                modified_payload = chat.RequestPayload(**request_dict)
                response = await provider.chat(modified_payload)
                return await run_post_llm_guardrails(
                    guardrails,
                    request_dict,
                    response,
                    auth_headers=auth_headers,
                    usage_tracking=endpoint_config.usage_tracking,
                )

            try:
                return await maybe_traced_gateway_call(
                    _guarded_chat,
                    endpoint_config,
                    user_metadata,
                    request_headers=headers,
                    request_type=GatewayRequestType.UNIFIED_CHAT,
                    on_complete=make_budget_on_complete(store, workspace),
                )(payload)
            except GuardrailViolation as e:
                raise HTTPException(status_code=400, detail=str(e))

    elif "input" in body:
        # Embeddings request
        endpoint_type = EndpointType.LLM_V1_EMBEDDINGS
        try:
            payload = embeddings.RequestPayload(**body)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid embeddings payload: {e!s}")

        provider, endpoint_config = _create_provider_from_endpoint_name(
            store, endpoint_name, endpoint_type
        )

        return await maybe_traced_gateway_call(
            provider.embeddings,
            endpoint_config,
            user_metadata,
            request_headers=headers,
            request_type=GatewayRequestType.UNIFIED_EMBEDDINGS,
            on_complete=make_budget_on_complete(store, workspace),
        )(payload)

    else:
        raise HTTPException(
            status_code=400,
            detail="Invalid request: payload format must be either chat or embeddings",
        )


@gateway_router.post("/mlflow/v1/chat/completions", response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.MLFLOW_CHAT_COMPLETIONS)
async def chat_completions(request: Request):
    """
    OpenAI-compatible chat completions endpoint.

    This endpoint follows the OpenAI API format where the endpoint name is specified
    via the "model" parameter in the request body, allowing clients to use the
    standard OpenAI SDK.

    Example:
        POST /gateway/mlflow/v1/chat/completions
        {
            "model": "my-endpoint-name",
            "messages": [{"role": "user", "content": "Hello"}]
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)
    headers = dict(request.headers)

    # Extract endpoint name from "model" parameter
    endpoint_name = _extract_endpoint_name_from_model(body)
    body.pop("model")

    store = _get_store()
    workspace = get_request_workspace()

    _validate_store(store)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    try:
        payload = chat.RequestPayload(**body)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid chat payload: {e!s}")

    if payload.stream:
        # Post-LLM guardrails are not applied to streaming responses.
        # Pre-LLM guardrails run inside the trace as child spans; violations
        # are surfaced as SSE error chunks via safe_stream.
        async def _guarded_stream(
            payload: chat.RequestPayload,
        ):
            request_dict = await run_pre_llm_guardrails(
                guardrails,
                payload.model_dump(),
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
                payload_schema=ChatCompletionRequest.model_json_schema(),
            )
            async for chunk in provider.chat_stream(chat.RequestPayload(**request_dict)):
                yield chunk

        stream = maybe_traced_gateway_call(
            _guarded_stream,
            endpoint_config,
            user_metadata,
            output_reducer=aggregate_chat_stream_chunks,
            request_headers=headers,
            request_type=GatewayRequestType.UNIFIED_CHAT,
            on_complete=make_budget_on_complete(store, workspace),
        )(payload)
        return StreamingResponse(
            safe_stream(to_sse_chunk(chunk.model_dump_json()) async for chunk in stream),
            media_type="text/event-stream",
        )
    else:

        async def _guarded_chat(
            payload: chat.RequestPayload,
        ) -> chat.ResponsePayload:
            request_dict = await run_pre_llm_guardrails(
                guardrails,
                payload.model_dump(),
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
                payload_schema=ChatCompletionRequest.model_json_schema(),
            )
            modified_payload = chat.RequestPayload(**request_dict)
            response = await provider.chat(modified_payload)
            return await run_post_llm_guardrails(
                guardrails,
                request_dict,
                response,
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
            )

        try:
            return await maybe_traced_gateway_call(
                _guarded_chat,
                endpoint_config,
                user_metadata,
                request_headers=headers,
                request_type=GatewayRequestType.UNIFIED_CHAT,
                on_complete=make_budget_on_complete(store, workspace),
            )(payload)
        except GuardrailViolation as e:
            raise HTTPException(status_code=400, detail=str(e))


@gateway_router.post(PASSTHROUGH_ROUTES[PassthroughAction.OPENAI_CHAT], response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.OPENAI_PASSTHROUGH_CHAT)
async def openai_passthrough_chat(request: Request):
    """
    OpenAI passthrough endpoint for chat completions.

    This endpoint accepts raw OpenAI API format and passes it through to the
    OpenAI provider with the configured API key and model. The 'model' parameter
    in the request specifies which MLflow endpoint to use.

    Supports streaming responses when the 'stream' parameter is set to true.

    Example:
        POST /gateway/openai/v1/chat/completions
        {
            "model": "my-openai-endpoint",
            "messages": [{"role": "user", "content": "Hello"}],
            "temperature": 0.7,
            "stream": true
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)

    endpoint_name = _extract_endpoint_name_from_model(body)
    body.pop("model")
    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    if body.get("stream", False):
        # Post-LLM guardrails are not applied to streaming responses.
        async def _guarded_stream(body: dict[str, Any]):
            request_dict = await run_pre_llm_guardrails(
                guardrails,
                body,
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
            )
            stream = await provider.passthrough(
                action=PassthroughAction.OPENAI_CHAT, payload=request_dict, headers=headers
            )
            async for chunk in stream:
                yield chunk

        traced_stream = maybe_traced_gateway_call(
            _guarded_stream,
            endpoint_config,
            user_metadata,
            request_headers=headers,
            request_type=GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_CHAT,
            on_complete=make_budget_on_complete(store, workspace),
        )
        return StreamingResponse(
            safe_stream(traced_stream(body), as_bytes=True), media_type="text/event-stream"
        )

    async def _guarded_passthrough(body: dict[str, Any]) -> dict[str, Any]:
        body = await run_pre_llm_guardrails(
            guardrails,
            body,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )
        response = await provider.passthrough(
            action=PassthroughAction.OPENAI_CHAT, payload=body, headers=headers
        )
        return await run_post_llm_guardrails_passthrough(
            guardrails,
            body,
            response,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )

    try:
        return await maybe_traced_gateway_call(
            _guarded_passthrough,
            endpoint_config,
            user_metadata,
            request_headers=headers,
            request_type=GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_CHAT,
            on_complete=make_budget_on_complete(store, workspace),
        )(body)
    except GuardrailViolation as e:
        raise HTTPException(status_code=400, detail=str(e))


@gateway_router.post(PASSTHROUGH_ROUTES[PassthroughAction.OPENAI_EMBEDDINGS], response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.OPENAI_PASSTHROUGH_EMBEDDINGS)
async def openai_passthrough_embeddings(request: Request):
    """
    OpenAI passthrough endpoint for embeddings.

    This endpoint accepts raw OpenAI API format and passes it through to the
    OpenAI provider with the configured API key and model. The 'model' parameter
    in the request specifies which MLflow endpoint to use.

    Example:
        POST /gateway/openai/v1/embeddings
        {
            "model": "my-openai-endpoint",
            "input": "The food was delicious and the waiter..."
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)

    endpoint_name = _extract_endpoint_name_from_model(body)
    body.pop("model")
    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_EMBEDDINGS
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    try:
        body = await run_pre_llm_guardrails(
            guardrails,
            body,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )
    except GuardrailViolation as e:
        raise HTTPException(status_code=400, detail=str(e))

    traced_passthrough = maybe_traced_gateway_call(
        provider.passthrough,
        endpoint_config,
        user_metadata,
        request_headers=headers,
        request_type=GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_EMBEDDINGS,
        on_complete=make_budget_on_complete(store, workspace),
    )
    # Post-LLM guardrails are skipped for embeddings: responses are float vectors
    # that content judges cannot meaningfully evaluate.
    return await traced_passthrough(
        action=PassthroughAction.OPENAI_EMBEDDINGS, payload=body, headers=headers
    )


async def _openai_responses_passthrough_unary(
    request: Request,
    body: dict[str, Any],
    action: PassthroughAction,
    request_type: GatewayRequestType,
) -> dict[str, Any]:
    """Shared body for the unary (non-streaming) OpenAI Responses-shaped
    passthrough routes — the non-streaming branch of ``/responses`` and the
    unary-only ``/responses/compact``.

    Resolves the configured endpoint from ``body["model"]``, runs the usual
    telemetry / budget / guardrail setup, wraps the upstream call in pre/post
    guardrails, and returns the raw upstream response.
    ``GuardrailViolation`` is translated to HTTP 400. Each call site is
    responsible for reading the request body (so callers can also reject
    ``stream=true`` when their endpoint doesn't support streaming, as
    ``/responses/compact`` does).
    """
    user_metadata = _get_user_metadata(request)
    endpoint_name = _extract_endpoint_name_from_model(body)
    body.pop("model")
    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    async def _guarded_passthrough(body: dict[str, Any]) -> dict[str, Any]:
        body = await run_pre_llm_guardrails(
            guardrails,
            body,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )
        response = await provider.passthrough(action=action, payload=body, headers=headers)
        return await run_post_llm_guardrails_passthrough(
            guardrails,
            body,
            response,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )

    try:
        return await maybe_traced_gateway_call(
            _guarded_passthrough,
            endpoint_config,
            user_metadata,
            request_headers=headers,
            request_type=request_type,
            on_complete=make_budget_on_complete(store, workspace),
        )(body)
    except GuardrailViolation as e:
        raise HTTPException(status_code=400, detail=str(e))


@gateway_router.post(PASSTHROUGH_ROUTES[PassthroughAction.OPENAI_RESPONSES], response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.OPENAI_PASSTHROUGH_RESPONSES)
async def openai_passthrough_responses(request: Request):
    """
    OpenAI passthrough endpoint for the Responses API.

    This endpoint accepts raw OpenAI Responses API format and passes it through to the
    OpenAI provider with the configured API key and model. The 'model' parameter
    in the request specifies which MLflow endpoint to use.

    Supports streaming responses when the 'stream' parameter is set to true.

    Example:
        POST /gateway/openai/v1/responses
        {
            "model": "my-openai-endpoint",
            "input": [{"type": "text", "text": "Hello"}],
            "instructions": "You are a helpful assistant",
            "stream": true
        }
    """
    body = await _get_request_body(request)

    if body.get("stream", False):
        # Post-LLM guardrails are not applied to streaming responses. Streaming
        # keeps its own inline setup because the closure below captures these
        # variables directly; the unary path delegates to the shared helper.
        user_metadata = _get_user_metadata(request)
        endpoint_name = _extract_endpoint_name_from_model(body)
        body.pop("model")
        store = _get_store()
        workspace = get_request_workspace()
        _validate_store(store)
        headers = dict(request.headers)
        provider, endpoint_config = _create_provider_from_endpoint_name(
            store, endpoint_name, EndpointType.LLM_V1_CHAT
        )
        _set_gateway_telemetry_state(request, endpoint_config)
        check_budget_limit(store, endpoint_config, workspace=workspace)
        guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

        async def _guarded_stream(body: dict[str, Any]):
            request_dict = await run_pre_llm_guardrails(
                guardrails,
                body,
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
            )
            stream = await provider.passthrough(
                action=PassthroughAction.OPENAI_RESPONSES, payload=request_dict, headers=headers
            )
            async for chunk in stream:
                yield chunk

        traced_stream = maybe_traced_gateway_call(
            _guarded_stream,
            endpoint_config,
            user_metadata,
            output_reducer=aggregate_openai_responses_stream_chunks,
            request_headers=headers,
            request_type=GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_RESPONSES,
            on_complete=make_budget_on_complete(store, workspace),
        )
        return StreamingResponse(
            safe_stream(traced_stream(body), as_bytes=True), media_type="text/event-stream"
        )

    return await _openai_responses_passthrough_unary(
        request,
        body,
        action=PassthroughAction.OPENAI_RESPONSES,
        request_type=GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_RESPONSES,
    )


@gateway_router.post(
    PASSTHROUGH_ROUTES[PassthroughAction.OPENAI_RESPONSES_COMPACT], response_model=None
)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.OPENAI_PASSTHROUGH_RESPONSES)
async def openai_passthrough_responses_compact(request: Request):
    """
    OpenAI passthrough endpoint for the Responses API ``/compact`` route.

    Mirrors :func:`openai_passthrough_responses` for session compaction calls
    (e.g. those issued by the OpenAI Agents SDK's ``OpenAIResponsesCompactionSession``).
    Compaction is a unary request — there is no streaming variant. A request
    that sets ``stream=true`` is rejected with HTTP 400 to avoid the underlying
    provider passthrough (which treats all non-embeddings actions as
    stream-capable) attempting an SSE stream against an upstream endpoint that
    does not support it.

    Example:
        POST /gateway/openai/v1/responses/compact
        {
            "model": "my-openai-endpoint",
            "previous_response_id": "resp_abc123"
        }
    """
    body = await _get_request_body(request)

    # /responses/compact is unary upstream; explicitly reject `stream=true`
    # before the provider's passthrough machinery tries to open an SSE stream.
    if body.get("stream", False):
        raise HTTPException(
            status_code=400,
            detail=(
                "stream=true is not supported on /responses/compact; compaction is a unary request."
            ),
        )

    return await _openai_responses_passthrough_unary(
        request,
        body,
        action=PassthroughAction.OPENAI_RESPONSES_COMPACT,
        request_type=GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_RESPONSES,
    )


@gateway_router.post(PASSTHROUGH_ROUTES[PassthroughAction.ANTHROPIC_MESSAGES], response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.ANTHROPIC_PASSTHROUGH_MESSAGES)
async def anthropic_passthrough_messages(request: Request):
    """
    Anthropic passthrough endpoint for the Messages API.

    This endpoint accepts raw Anthropic API format and passes it through to the
    Anthropic provider with the configured API key and model. The 'model' parameter
    in the request specifies which MLflow endpoint to use.

    Supports streaming responses when the 'stream' parameter is set to true.

    Example:
        POST /gateway/anthropic/v1/messages
        {
            "model": "my-anthropic-endpoint",
            "messages": [{"role": "user", "content": "Hello"}],
            "max_tokens": 1024,
            "stream": true
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)

    endpoint_name = _extract_endpoint_name_from_model(body)
    body.pop("model")
    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    if body.get("stream", False):
        # Post-LLM guardrails are not applied to streaming responses.
        async def _guarded_stream(body: dict[str, Any]):
            request_dict = await run_pre_llm_guardrails(
                guardrails,
                body,
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
            )
            stream = await provider.passthrough(
                action=PassthroughAction.ANTHROPIC_MESSAGES, payload=request_dict, headers=headers
            )
            async for chunk in stream:
                yield chunk

        traced_stream = maybe_traced_gateway_call(
            _guarded_stream,
            endpoint_config,
            user_metadata,
            output_reducer=aggregate_anthropic_messages_stream_chunks,
            request_headers=headers,
            request_type=GatewayRequestType.PASSTHROUGH_MODEL_ANTHROPIC_MESSAGES,
            on_complete=make_budget_on_complete(store, workspace),
            message_format="anthropic",
        )
        return StreamingResponse(
            safe_stream(traced_stream(body), as_bytes=True), media_type="text/event-stream"
        )

    async def _guarded_passthrough(body: dict[str, Any]) -> dict[str, Any]:
        body = await run_pre_llm_guardrails(
            guardrails,
            body,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )
        response = await provider.passthrough(
            action=PassthroughAction.ANTHROPIC_MESSAGES, payload=body, headers=headers
        )
        return await run_post_llm_guardrails_passthrough(
            guardrails,
            body,
            response,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )

    try:
        return await maybe_traced_gateway_call(
            _guarded_passthrough,
            endpoint_config,
            user_metadata,
            request_headers=headers,
            request_type=GatewayRequestType.PASSTHROUGH_MODEL_ANTHROPIC_MESSAGES,
            on_complete=make_budget_on_complete(store, workspace),
            message_format="anthropic",
        )(body)
    except GuardrailViolation as e:
        raise HTTPException(status_code=400, detail=str(e))


@gateway_router.post(
    PASSTHROUGH_ROUTES[PassthroughAction.GEMINI_GENERATE_CONTENT], response_model=None
)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.GEMINI_PASSTHROUGH_GENERATE_CONTENT)
async def gemini_passthrough_generate_content(endpoint_name: str, request: Request):
    """
    Gemini passthrough endpoint for generateContent API (non-streaming).

    This endpoint accepts raw Gemini API format and passes it through to the
    Gemini provider with the configured API key. The endpoint_name in the URL path
    specifies which MLflow endpoint to use.

    Example:
        POST /gateway/gemini/v1beta/models/my-gemini-endpoint:generateContent
        {
            "contents": [
                {
                    "role": "user",
                    "parts": [{"text": "Hello"}]
                }
            ]
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)

    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    async def _guarded_passthrough(body: dict[str, Any]) -> dict[str, Any]:
        body = await run_pre_llm_guardrails(
            guardrails,
            body,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )
        response = await provider.passthrough(
            action=PassthroughAction.GEMINI_GENERATE_CONTENT, payload=body, headers=headers
        )
        return await run_post_llm_guardrails_passthrough(
            guardrails,
            body,
            response,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )

    try:
        return await maybe_traced_gateway_call(
            _guarded_passthrough,
            endpoint_config,
            user_metadata,
            request_headers=headers,
            request_type=GatewayRequestType.PASSTHROUGH_MODEL_GEMINI_GENERATE_CONTENT,
            on_complete=make_budget_on_complete(store, workspace),
            message_format="gemini",
        )(body)
    except GuardrailViolation as e:
        raise HTTPException(status_code=400, detail=str(e))


@gateway_router.post(
    PASSTHROUGH_ROUTES[PassthroughAction.GEMINI_STREAM_GENERATE_CONTENT], response_model=None
)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.GEMINI_PASSTHROUGH_STREAM_GENERATE_CONTENT)
async def gemini_passthrough_stream_generate_content(endpoint_name: str, request: Request):
    """
    Gemini passthrough endpoint for streamGenerateContent API (streaming).

    This endpoint accepts raw Gemini API format and passes it through to the
    Gemini provider with the configured API key. The endpoint_name in the URL path
    specifies which MLflow endpoint to use.

    Example:
        POST /gateway/gemini/v1beta/models/my-gemini-endpoint:streamGenerateContent
        {
            "contents": [
                {
                    "role": "user",
                    "parts": [{"text": "Hello"}]
                }
            ]
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)

    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    # Post-LLM guardrails are not applied to streaming responses.
    async def _guarded_stream(body: dict[str, Any]):
        request_dict = await run_pre_llm_guardrails(
            guardrails,
            body,
            auth_headers=auth_headers,
            usage_tracking=endpoint_config.usage_tracking,
        )
        stream = await provider.passthrough(
            action=PassthroughAction.GEMINI_STREAM_GENERATE_CONTENT,
            payload=request_dict,
            headers=headers,
        )
        async for chunk in stream:
            yield chunk

    traced_stream = maybe_traced_gateway_call(
        _guarded_stream,
        endpoint_config,
        user_metadata,
        output_reducer=aggregate_gemini_stream_generate_content_chunks,
        request_headers=headers,
        request_type=GatewayRequestType.PASSTHROUGH_MODEL_GEMINI_GENERATE_CONTENT,
        on_complete=make_budget_on_complete(store, workspace),
        message_format="gemini",
    )
    return StreamingResponse(
        safe_stream(traced_stream(body), as_bytes=True), media_type="text/event-stream"
    )


@gateway_router.post("/proxy/{endpoint_name}/{path:path}", response_model=None)
@translate_http_exception
@_record_gateway_invocation(GatewayInvocationType.RAW_PROXY)
async def raw_proxy(endpoint_name: str, path: str, request: Request):
    """
    Raw proxy endpoint.

    Routes the request payload as-is to the upstream provider at
    <provider_base_url>/<path>, using the credentials configured for the named
    endpoint. Unlike the typed passthrough routes, the ``model`` field in the
    payload is NOT replaced with the value from the endpoint config.

    Streaming is detected automatically from the response Content-Type
    (``text/event-stream`` or ``application/x-ndjson``), so this endpoint
    supports providers like Gemini that signal streaming via Content-Type rather
    than a ``"stream": true`` request flag.

    Example:
        POST /gateway/proxy/my-openai-endpoint/chat/completions
        {
            "model": "gpt-4o",
            "messages": [{"role": "user", "content": "Hello"}]
        }
    """
    body = await _get_request_body(request)
    user_metadata = _get_user_metadata(request)
    store = _get_store()
    workspace = get_request_workspace()
    _validate_store(store)
    headers = dict(request.headers)
    # Preserve upstream query parameters (e.g. Gemini's ?alt=sse for SSE streaming).
    if qs := request.url.query:
        path = f"{path}?{qs}"
    # The proxy route is type-agnostic; EndpointType is only used to satisfy
    # EndpointConfig construction and has no effect on proxy behavior (the proxy
    # bypasses typed routing entirely). LLM_V1_CHAT is used as the default since
    # the DB does not store a per-endpoint type for DB-backed gateway endpoints.
    provider, endpoint_config = _create_provider_from_endpoint_name(
        store, endpoint_name, EndpointType.LLM_V1_CHAT
    )
    _set_gateway_telemetry_state(request, endpoint_config)
    check_budget_limit(store, endpoint_config, workspace=workspace)
    guardrails, auth_headers = _get_guardrails_and_auth(store, endpoint_config, request)

    # _do_proxy is always an async generator so maybe_traced_gateway_call can wrap it
    # consistently regardless of whether the upstream responds with a streaming or
    # non-streaming body.  For streaming responses it yields raw bytes; for
    # non-streaming responses it yields a single dict.  The route handler peeks at
    # the first item to decide how to respond to the client.
    async def _do_proxy(body: dict[str, Any]):
        try:
            body = await run_pre_llm_guardrails(
                guardrails,
                body,
                auth_headers=auth_headers,
                usage_tracking=endpoint_config.usage_tracking,
            )
        except GuardrailViolation as e:
            raise HTTPException(status_code=400, detail=str(e))

        result = await provider.proxy(path, body, headers)

        if isinstance(result, AsyncIterable):
            # Post-LLM guardrails are not applied to streaming responses.
            async for chunk in result:
                yield chunk
        else:
            try:
                result = await run_post_llm_guardrails_passthrough(
                    guardrails,
                    body,
                    result,
                    auth_headers=auth_headers,
                    usage_tracking=endpoint_config.usage_tracking,
                )
            except GuardrailViolation as e:
                raise HTTPException(status_code=400, detail=str(e))
            yield result

    traced_proxy = maybe_traced_gateway_call(
        _do_proxy,
        endpoint_config,
        user_metadata,
        request_headers=headers,
        request_type=GatewayRequestType.RAW_PROXY,
        on_complete=make_budget_on_complete(store, workspace),
    )

    gen = traced_proxy(body)
    try:
        first = await gen.__anext__()
    except StopAsyncIteration:
        # _do_proxy raised before yielding (e.g. guardrail violation, provider 501).
        # The original HTTPException already propagated; re-raise a generic 500 as fallback.
        raise HTTPException(status_code=500, detail="Proxy handler exited without a response.")
    if isinstance(first, bytes):

        async def _prepend(first_chunk, rest):
            yield first_chunk
            async for chunk in rest:
                yield chunk

        return StreamingResponse(
            safe_stream(_prepend(first, gen), as_bytes=True), media_type="text/event-stream"
        )
    return first