import json from pathlib import Path from typing import Any from unittest import mock from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI, HTTPException from fastapi.responses import StreamingResponse from starlette.testclient import TestClient import mlflow from mlflow.entities import ( FallbackConfig, FallbackStrategy, GatewayEndpointModelConfig, GatewayModelLinkageType, RoutingStrategy, SpanType, ) from mlflow.entities.gateway_guardrail import GuardrailAction, GuardrailStage from mlflow.entities.trace_state import TraceState from mlflow.exceptions import MlflowException from mlflow.gateway.config import ( EndpointType, GatewayRequestType, GeminiConfig, LiteLLMConfig, MistralConfig, OpenAIAPIType, OpenAIConfig, ) from mlflow.gateway.constants import MLFLOW_GATEWAY_DURATION_HEADER, MLFLOW_GATEWAY_OVERHEAD_HEADER from mlflow.gateway.guardrails import _SANITIZE_BYPASS_HEADER, JudgeGuardrail from mlflow.gateway.providers.anthropic import AnthropicProvider from mlflow.gateway.providers.base import ( FallbackProvider, TrafficRouteProvider, ) from mlflow.gateway.providers.databricks import DatabricksConfig, DatabricksProvider from mlflow.gateway.providers.gemini import GeminiProvider from mlflow.gateway.providers.litellm import LiteLLMProvider from mlflow.gateway.providers.mistral import MistralProvider from mlflow.gateway.providers.openai import OpenAIProvider from mlflow.gateway.providers.utils import provider_call_duration_ms from mlflow.gateway.schemas import chat, embeddings from mlflow.server.fastapi_app import add_gateway_timing_middleware from mlflow.server.gateway_api import ( _build_endpoint_config, _create_provider_from_endpoint_name, anthropic_passthrough_messages, chat_completions, gateway_router, gemini_passthrough_generate_content, gemini_passthrough_stream_generate_content, invocations, openai_passthrough_chat, openai_passthrough_embeddings, openai_passthrough_responses, openai_passthrough_responses_compact, ) from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig, GatewayModelConfig from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore from mlflow.tracing.client import TracingClient from mlflow.tracing.constant import ( SpanAttributeKey, TokenUsageKey, TraceMetadataKey, ) pytestmark = pytest.mark.notrackingurimock TEST_PASSPHRASE = "test-passphrase-for-gateway-api-tests" @pytest.fixture(autouse=True) def set_kek_passphrase(monkeypatch): monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", TEST_PASSPHRASE) @pytest.fixture def store(tmp_path: Path, db_uri: str): artifact_uri = tmp_path / "artifacts" artifact_uri.mkdir(exist_ok=True) mlflow.set_tracking_uri(db_uri) yield SqlAlchemyStore(db_uri, artifact_uri.as_uri()) mlflow.set_tracking_uri(None) def create_mock_request( cached_body: dict[str, Any] | None = None, username: str | None = None, user_id: int | str | None = None, ) -> MagicMock: """Create a mock request with proper state attributes for gateway tests.""" mock_request = MagicMock() mock_request.state.cached_body = cached_body mock_request.state.username = username mock_request.state.user_id = user_id return mock_request def _make_model_config(provider="openai", model_name="gpt-4o"): return GatewayModelConfig( model_definition_id="md-test", provider=provider, model_name=model_name, secret_value={"api_key": "sk-test"}, ) def test_build_endpoint_config_rejects_provider_not_in_allowed_list(monkeypatch): monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "anthropic") with pytest.raises(MlflowException, match="not allowed"): _build_endpoint_config("test-ep", _make_model_config("openai"), EndpointType.LLM_V1_CHAT) def test_build_endpoint_config_allows_provider_in_allowed_list(monkeypatch): monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "openai") config = _build_endpoint_config( "test-ep", _make_model_config("openai"), EndpointType.LLM_V1_CHAT ) assert config.name == "test-ep" assert isinstance(config.model.config, OpenAIConfig) def test_build_endpoint_config_allows_provider_when_no_filter(): config = _build_endpoint_config( "test-ep", _make_model_config("openai"), EndpointType.LLM_V1_CHAT ) assert config.name == "test-ep" def test_create_provider_from_endpoint_name_openai(store: SqlAlchemyStore): # Create test data secret = store.create_gateway_secret( secret_name="openai-key", secret_value={"api_key": "sk-test-123"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="gpt-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) endpoint = store.create_gateway_endpoint( name="test-openai-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, OpenAIProvider) assert isinstance(provider.config.model.config, OpenAIConfig) assert provider.config.model.config.openai_api_key == "sk-test-123" def test_create_provider_from_endpoint_name_azure_openai(store: SqlAlchemyStore): # Test Azure OpenAI configuration secret = store.create_gateway_secret( secret_name="azure-openai-key", secret_value={"api_key": "azure-api-key-test"}, provider="openai", auth_config={ "api_type": "azure", "api_base": "https://my-resource.openai.azure.com", "api_version": "2024-02-01", }, ) model_def = store.create_gateway_model_definition( name="azure-gpt-model", secret_id=secret.secret_id, provider="azure", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-azure-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, OpenAIProvider) assert isinstance(provider.config.model.config, OpenAIConfig) assert provider.config.model.config.openai_api_type == OpenAIAPIType.AZURE assert provider.config.model.config.openai_api_base == "https://my-resource.openai.azure.com" assert provider.config.model.config.openai_deployment_name == "gpt-4" assert provider.config.model.config.openai_api_version == "2024-02-01" assert provider.config.model.config.openai_api_key == "azure-api-key-test" def test_create_provider_from_endpoint_name_azure_openai_with_azuread(store: SqlAlchemyStore): # Test Azure OpenAI with AzureAD authentication secret = store.create_gateway_secret( secret_name="azuread-openai-key", secret_value={"api_key": "azuread-api-key-test"}, provider="openai", auth_config={ "api_type": "azuread", "api_base": "https://my-resource-ad.openai.azure.com", "deployment_name": "gpt-4-deployment-ad", "api_version": "2024-02-01", }, ) model_def = store.create_gateway_model_definition( name="azuread-gpt-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-azuread-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, OpenAIProvider) assert isinstance(provider.config.model.config, OpenAIConfig) assert provider.config.model.config.openai_api_type == OpenAIAPIType.AZUREAD assert provider.config.model.config.openai_api_base == "https://my-resource-ad.openai.azure.com" assert provider.config.model.config.openai_deployment_name == "gpt-4-deployment-ad" assert provider.config.model.config.openai_api_version == "2024-02-01" assert provider.config.model.config.openai_api_key == "azuread-api-key-test" def test_create_provider_from_endpoint_name_anthropic(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="anthropic-key", secret_value={"api_key": "sk-ant-test"}, provider="anthropic", ) model_def = store.create_gateway_model_definition( name="claude-model", secret_id=secret.secret_id, provider="anthropic", model_name="claude-3-sonnet", ) endpoint = store.create_gateway_endpoint( name="test-anthropic-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, AnthropicProvider) assert provider.config.model.config.anthropic_api_key == "sk-ant-test" assert provider.base_url == "https://api.anthropic.com/v1" def test_create_provider_from_endpoint_name_anthropic_with_api_base(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="anthropic-proxy-key", secret_value={"api_key": "sk-ant-proxy-test"}, provider="anthropic", auth_config={"api_base": "http://localhost:6655/anthropic/v1"}, ) model_def = store.create_gateway_model_definition( name="claude-proxy-model", secret_id=secret.secret_id, provider="anthropic", model_name="claude-3-7-sonnet", ) endpoint = store.create_gateway_endpoint( name="test-anthropic-proxy-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, AnthropicProvider) assert provider.config.model.config.anthropic_api_key == "sk-ant-proxy-test" assert provider.config.model.config.anthropic_api_base == "http://localhost:6655/anthropic/v1" assert provider.base_url == "http://localhost:6655/anthropic/v1" def test_create_provider_from_endpoint_name_mistral(store: SqlAlchemyStore): # Test Mistral provider secret = store.create_gateway_secret( secret_name="mistral-key", secret_value={"api_key": "mistral-test-key"}, provider="mistral", ) model_def = store.create_gateway_model_definition( name="mistral-model", secret_id=secret.secret_id, provider="mistral", model_name="mistral-large-latest", ) endpoint = store.create_gateway_endpoint( name="test-mistral-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, MistralProvider) assert isinstance(provider.config.model.config, MistralConfig) assert provider.config.model.config.mistral_api_key == "mistral-test-key" def test_create_provider_from_endpoint_name_gemini(store: SqlAlchemyStore): # Test Gemini provider secret = store.create_gateway_secret( secret_name="gemini-key", secret_value={"api_key": "gemini-test-key"}, provider="gemini", ) model_def = store.create_gateway_model_definition( name="gemini-model", secret_id=secret.secret_id, provider="gemini", model_name="gemini-1.5-pro", ) endpoint = store.create_gateway_endpoint( name="test-gemini-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, GeminiProvider) assert isinstance(provider.config.model.config, GeminiConfig) assert provider.config.model.config.gemini_api_key == "gemini-test-key" def test_create_provider_from_endpoint_name_litellm(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="litellm-key", secret_value={"api_key": "litellm-test-key"}, provider="litellm", ) model_def = store.create_gateway_model_definition( name="litellm-model", secret_id=secret.secret_id, provider="litellm", model_name="claude-3-5-sonnet-20241022", ) endpoint = store.create_gateway_endpoint( name="test-litellm-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, LiteLLMProvider) assert isinstance(provider.config.model.config, LiteLLMConfig) assert provider.config.model.config.litellm_auth_config["api_key"] == "litellm-test-key" assert provider.config.model.config.litellm_provider == "litellm" # get_provider_name() returns the actual provider name for tracing/metrics assert provider.get_provider_name() == "litellm" def test_create_provider_from_endpoint_name_litellm_with_api_base(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="litellm-custom-key", secret_value={"api_key": "litellm-custom-key"}, provider="litellm", auth_config={"api_base": "https://custom-api.example.com"}, ) model_def = store.create_gateway_model_definition( name="litellm-custom-model", secret_id=secret.secret_id, provider="litellm", model_name="custom-model", ) endpoint = store.create_gateway_endpoint( name="test-litellm-custom-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, LiteLLMProvider) assert isinstance(provider.config.model.config, LiteLLMConfig) assert provider.config.model.config.litellm_auth_config["api_key"] == "litellm-custom-key" assert ( provider.config.model.config.litellm_auth_config["api_base"] == "https://custom-api.example.com" ) assert provider.config.model.config.litellm_provider == "litellm" @pytest.mark.parametrize( "input_url", [ "https://my-workspace.databricks.com", "https://my-workspace.databricks.com/serving-endpoints", ], ) def test_create_provider_from_endpoint_name_databricks_normalizes_base_url( store: SqlAlchemyStore, input_url: str ): secret = store.create_gateway_secret( secret_name="databricks-key", secret_value={"api_key": "databricks-token-123"}, provider="databricks", auth_config={"api_base": input_url}, ) model_def = store.create_gateway_model_definition( name="databricks-model", secret_id=secret.secret_id, provider="databricks", model_name="databricks-dbrx-instruct", ) endpoint = store.create_gateway_endpoint( name="test-databricks-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, DatabricksProvider) assert isinstance(provider.config.model.config, DatabricksConfig) # Verify the base URL was normalized to include /serving-endpoints assert provider._api_base == "https://my-workspace.databricks.com/serving-endpoints" def test_api_key_not_read_from_file(store: SqlAlchemyStore, tmp_path: Path, monkeypatch): monkeypatch.delenv("MLFLOW_GATEWAY_RESOLVE_API_KEY_FROM_FILE", raising=False) # Create a file whose path will be used as the "api_key" value secret_file = tmp_path / "secret.txt" secret_file.write_text("file-content-should-not-appear") secret = store.create_gateway_secret( secret_name="lfi-test-key", # Use the file path as the api_key — the gateway must NOT read the file secret_value={"api_key": str(secret_file)}, provider="openai", ) model_def = store.create_gateway_model_definition( name="lfi-test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) endpoint = store.create_gateway_endpoint( name="lfi-test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) # The key must be the literal file path string, NOT the file contents assert provider.config.model.config.openai_api_key == str(secret_file) assert provider.config.model.config.openai_api_key != "file-content-should-not-appear" def test_create_provider_from_endpoint_name_nonexistent_endpoint(store: SqlAlchemyStore): with pytest.raises(MlflowException, match="not found"): _create_provider_from_endpoint_name(store, "nonexistent-id", EndpointType.LLM_V1_CHAT) @pytest.mark.asyncio async def test_invocations_handler_chat(store: SqlAlchemyStore): # Create test data secret = store.create_gateway_secret( secret_name="chat-key", secret_value={"api_key": "sk-test-chat"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="chat-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="chat-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock the provider's chat method mock_response = chat.ResponsePayload( id="test-id", object="chat.completion", created=1234567890, model="gpt-4", choices=[ chat.Choice( index=0, message=chat.ResponseMessage(role="assistant", content="Hello!"), finish_reason="stop", ) ], usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15), ) # Create a mock request with chat payload mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "messages": [{"role": "user", "content": "Hi"}], "temperature": 0.7, "stream": False, } ) # Patch the provider creation to return a mocked provider with patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider: mock_provider = MagicMock() mock_provider.chat = AsyncMock(return_value=mock_response) mock_endpoint_config = GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) # Call the handler response = await invocations(endpoint.name, mock_request) # Verify assert response.id == "test-id" assert response.choices[0].message.content == "Hello!" assert mock_provider.chat.called @pytest.mark.asyncio async def test_invocations_handler_embeddings(store: SqlAlchemyStore): # Create test data secret = store.create_gateway_secret( secret_name="embed-key", secret_value={"api_key": "sk-test-embed"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="embed-model", secret_id=secret.secret_id, provider="openai", model_name="text-embedding-ada-002", ) endpoint = store.create_gateway_endpoint( name="embed-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock the provider's embeddings method mock_response = embeddings.ResponsePayload( object="list", data=[embeddings.EmbeddingObject(embedding=[0.1, 0.2, 0.3], index=0)], model="text-embedding-ada-002", usage=embeddings.EmbeddingsUsage(prompt_tokens=5, total_tokens=5), ) # Create a mock request with embeddings payload mock_request = create_mock_request() mock_request.json = AsyncMock(return_value={"input": "test text"}) # Patch the provider creation to return a mocked provider with patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider: mock_provider = MagicMock() mock_provider.embeddings = AsyncMock(return_value=mock_response) mock_endpoint_config = GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) # Call the handler response = await invocations(endpoint.name, mock_request) # Verify assert response.object == "list" assert len(response.data) == 1 assert response.data[0].embedding == [0.1, 0.2, 0.3] assert mock_provider.embeddings.called def test_gateway_router_initialization(): assert gateway_router is not None assert gateway_router.prefix == "/gateway" @pytest.mark.asyncio async def test_invocations_handler_invalid_json(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="test-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock request that raises exception when parsing JSON mock_request = create_mock_request() mock_request.json = AsyncMock(side_effect=ValueError("Invalid JSON")) with pytest.raises(HTTPException, match="Invalid JSON payload: Invalid JSON") as exc_info: await invocations(endpoint.name, mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_invocations_handler_missing_fields(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="test-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create request with neither messages nor input mock_request = create_mock_request() mock_request.json = AsyncMock(return_value={"temperature": 0.7}) with patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider: mock_provider = MagicMock() mock_endpoint_config = GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) with pytest.raises( HTTPException, match="Invalid request: payload format must be either chat or embeddings" ) as exc_info: await invocations(endpoint.name, mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_invocations_handler_invalid_chat_payload(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="test-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create request with invalid messages structure mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "messages": "not a list", # Should be a list "stream": False, } ) with patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider: mock_provider = MagicMock() mock_endpoint_config = GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) with pytest.raises(HTTPException, match="Invalid chat payload") as exc_info: await invocations(endpoint.name, mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_invocations_handler_invalid_embeddings_payload(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="test-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="test-model", secret_id=secret.secret_id, provider="openai", model_name="text-embedding-ada-002", ) endpoint = store.create_gateway_endpoint( name="test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create request with invalid input structure mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "input": 123, # Should be string or list of strings } ) with patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider: mock_provider = MagicMock() mock_endpoint_config = GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) with pytest.raises(HTTPException, match="Invalid embeddings payload") as exc_info: await invocations(endpoint.name, mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_invocations_handler_streaming(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="test-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create streaming request mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "messages": [{"role": "user", "content": "Hi"}], "stream": True, } ) # Mock streaming chunks async def mock_stream(): yield chat.StreamResponsePayload( id="test-id", object="chat.completion.chunk", created=1234567890, model="gpt-4", choices=[ chat.StreamChoice( index=0, delta=chat.StreamDelta(role="assistant", content="Hello"), finish_reason=None, ) ], ) with patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider: mock_provider = MagicMock() mock_provider.chat_stream = MagicMock(return_value=mock_stream()) mock_endpoint_config = GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) response = await invocations(endpoint.name, mock_request) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" # chat_stream is inside a lazy async generator; consume the body to trigger execution async for _ in response.body_iterator: pass assert mock_provider.chat_stream.called def test_create_provider_from_endpoint_name_no_models(store: SqlAlchemyStore): # Create a minimal endpoint to get an endpoint_name secret = store.create_gateway_secret( secret_name="test-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock get_endpoint_config to return an empty models list with patch( "mlflow.server.gateway_api.get_endpoint_config", return_value=GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name="test-endpoint", models=[] ), ): with pytest.raises(MlflowException, match="has no PRIMARY models configured"): _create_provider_from_endpoint_name(store, endpoint.name, EndpointType.LLM_V1_CHAT) # ============================================================================= # OpenAI-compatible chat completions endpoint tests # ============================================================================= @pytest.mark.asyncio async def test_chat_completions_endpoint(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-compat-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-compat-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) store.create_gateway_endpoint( name="my-chat-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock the provider's chat method mock_response = chat.ResponsePayload( id="test-id", object="chat.completion", created=1234567890, model="gpt-4", choices=[ chat.Choice( index=0, message=chat.ResponseMessage(role="assistant", content="Hello from OpenAI!"), finish_reason="stop", ) ], usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15), ) # Create a mock request with OpenAI-compatible format mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "my-chat-endpoint", # Endpoint name via model parameter "messages": [{"role": "user", "content": "Hi"}], "temperature": 0.7, "stream": False, } ) # Patch the provider creation to return a mocked provider with ( patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider, patch("mlflow.server.gateway_api.load_guardrails", return_value=[]), ): mock_provider = MagicMock() mock_provider.chat = AsyncMock(return_value=mock_response) mock_endpoint_config = GatewayEndpointConfig( endpoint_id="test-endpoint-id", endpoint_name="my-chat-endpoint", models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) # Call the handler response = await chat_completions(mock_request) # Verify assert response.id == "test-id" assert response.choices[0].message.content == "Hello from OpenAI!" assert mock_provider.chat.called def test_response_timing_headers(store: SqlAlchemyStore): app = FastAPI() app.include_router(gateway_router) add_gateway_timing_middleware(app) mock_response = chat.ResponsePayload( id="test-id", object="chat.completion", created=1234567890, model="gpt-4", choices=[ chat.Choice( index=0, message=chat.ResponseMessage(role="assistant", content="Hello!"), finish_reason="stop", ) ], usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15), ) mock_endpoint_config = GatewayEndpointConfig( endpoint_id="test-endpoint-id", endpoint_name="my-endpoint", models=[] ) async def _mock_chat_with_provider_timing(payload): # Simulate a real provider call by setting the ContextVar as send_request would. provider_call_duration_ms.set(50.0) return mock_response with ( patch("mlflow.server.gateway_api._get_store", return_value=store), patch("mlflow.server.gateway_api.get_request_workspace", return_value=None), patch("mlflow.server.gateway_api.check_budget_limit"), patch("mlflow.server.gateway_api.load_guardrails", return_value=[]), patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider, ): mock_provider = MagicMock() mock_provider.chat = _mock_chat_with_provider_timing mock_create_provider.return_value = (mock_provider, mock_endpoint_config) client = TestClient(app) response = client.post( "/gateway/mlflow/v1/chat/completions", json={"model": "my-endpoint", "messages": [{"role": "user", "content": "Hi"}]}, ) assert response.status_code == 200 duration = int(response.headers[MLFLOW_GATEWAY_DURATION_HEADER]) overhead = int(response.headers[MLFLOW_GATEWAY_OVERHEAD_HEADER]) assert duration >= 0 assert 0 <= overhead <= duration def test_response_timing_headers_streaming(store: SqlAlchemyStore): app = FastAPI() app.include_router(gateway_router) add_gateway_timing_middleware(app) mock_endpoint_config = GatewayEndpointConfig( endpoint_id="test-endpoint-id", endpoint_name="my-endpoint", models=[] ) async def _mock_chat_stream(payload): yield chat.StreamResponsePayload( id="test-id", object="chat.completion.chunk", created=1234567890, model="gpt-4", choices=[ chat.StreamChoice( index=0, delta=chat.StreamDelta(role="assistant", content="Hello"), finish_reason=None, ) ], ) with ( patch("mlflow.server.gateway_api._get_store", return_value=store), patch("mlflow.server.gateway_api.get_request_workspace", return_value=None), patch("mlflow.server.gateway_api.check_budget_limit"), patch("mlflow.server.gateway_api.load_guardrails", return_value=[]), patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider, ): mock_provider = MagicMock() mock_provider.chat_stream = MagicMock(return_value=_mock_chat_stream(None)) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) client = TestClient(app) response = client.post( "/gateway/mlflow/v1/chat/completions", json={ "model": "my-endpoint", "messages": [{"role": "user", "content": "Hi"}], "stream": True, }, ) assert response.status_code == 200 assert MLFLOW_GATEWAY_DURATION_HEADER in response.headers assert int(response.headers[MLFLOW_GATEWAY_DURATION_HEADER]) >= 0 # Overhead header is omitted for streaming since provider_call_duration_ms is not set. assert MLFLOW_GATEWAY_OVERHEAD_HEADER not in response.headers def test_response_timing_headers_error(store: SqlAlchemyStore): app = FastAPI() app.include_router(gateway_router) add_gateway_timing_middleware(app) mock_endpoint_config = GatewayEndpointConfig( endpoint_id="test-endpoint-id", endpoint_name="my-error-endpoint", models=[] ) async def _mock_chat_raises(payload): provider_call_duration_ms.set(30.0) raise HTTPException(status_code=502, detail="Upstream provider error") with ( patch("mlflow.server.gateway_api._get_store", return_value=store), patch("mlflow.server.gateway_api.get_request_workspace", return_value=None), patch("mlflow.server.gateway_api.check_budget_limit"), patch("mlflow.server.gateway_api.load_guardrails", return_value=[]), patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider, ): mock_provider = MagicMock() mock_provider.chat = _mock_chat_raises mock_create_provider.return_value = (mock_provider, mock_endpoint_config) client = TestClient(app) response = client.post( "/gateway/mlflow/v1/chat/completions", json={"model": "my-error-endpoint", "messages": [{"role": "user", "content": "Hi"}]}, ) assert response.status_code == 502 assert MLFLOW_GATEWAY_DURATION_HEADER in response.headers duration = int(response.headers[MLFLOW_GATEWAY_DURATION_HEADER]) assert duration >= 0 assert MLFLOW_GATEWAY_OVERHEAD_HEADER in response.headers overhead = int(response.headers[MLFLOW_GATEWAY_OVERHEAD_HEADER]) assert 0 <= overhead <= duration @pytest.mark.asyncio async def test_chat_completions_endpoint_streaming(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="stream-key", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="stream-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) store.create_gateway_endpoint( name="stream-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create streaming request mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "stream-endpoint", "messages": [{"role": "user", "content": "Hi"}], "stream": True, } ) # Mock streaming chunks async def mock_stream(): yield chat.StreamResponsePayload( id="test-id", object="chat.completion.chunk", created=1234567890, model="gpt-4", choices=[ chat.StreamChoice( index=0, delta=chat.StreamDelta(role="assistant", content="Hello"), finish_reason=None, ) ], ) with ( patch( "mlflow.server.gateway_api._create_provider_from_endpoint_name" ) as mock_create_provider, patch("mlflow.server.gateway_api.load_guardrails", return_value=[]), ): mock_provider = MagicMock() mock_provider.chat_stream = MagicMock(return_value=mock_stream()) mock_endpoint_config = GatewayEndpointConfig( endpoint_id="test-endpoint-id", endpoint_name="stream-endpoint", models=[] ) mock_create_provider.return_value = (mock_provider, mock_endpoint_config) response = await chat_completions(mock_request) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" # chat_stream is inside a lazy async generator; consume the body to trigger execution async for _ in response.body_iterator: pass assert mock_provider.chat_stream.called @pytest.mark.asyncio async def test_chat_completions_endpoint_missing_model_parameter(store: SqlAlchemyStore): # Create request without model parameter mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "messages": [{"role": "user", "content": "Hi"}], } ) with pytest.raises(HTTPException, match="Missing required 'model' parameter") as exc_info: await chat_completions(mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_chat_completions_endpoint_missing_messages(store: SqlAlchemyStore): # Create test endpoint first so we can test payload validation secret = store.create_gateway_secret( secret_name="chat-missing-msg-key", secret_value={"api_key": "sk-test-key"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="gpt-missing-msg-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) store.create_gateway_endpoint( name="my-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create request without messages mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "my-endpoint", "temperature": 0.7, } ) with pytest.raises(HTTPException, match="Invalid chat payload") as exc_info: await chat_completions(mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_chat_completions_endpoint_invalid_json(store: SqlAlchemyStore): mock_request = create_mock_request() mock_request.json = AsyncMock(side_effect=ValueError("Invalid JSON")) with pytest.raises(HTTPException, match="Invalid JSON payload: Invalid JSON") as exc_info: await chat_completions(mock_request) assert exc_info.value.status_code == 400 @pytest.mark.asyncio async def test_openai_passthrough_chat(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-passthrough-key", secret_value={"api_key": "sk-test-passthrough"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-passthrough-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) store.create_gateway_endpoint( name="openai-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock OpenAI API response mock_response = { "id": "chatcmpl-123", "object": "chat.completion", "created": 1234567890, "model": "gpt-4o", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "Hello from passthrough!"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, } # Create mock request mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-passthrough-endpoint", "messages": [{"role": "user", "content": "Hello"}], } ) # Mock send_request directly with mock.patch( "mlflow.gateway.providers.openai.send_request", return_value=mock_response ) as mock_send: response = await openai_passthrough_chat(mock_request) # Verify send_request was called assert mock_send.called call_kwargs = mock_send.call_args[1] assert call_kwargs["path"] == "chat/completions" assert call_kwargs["payload"]["model"] == "gpt-4o" assert call_kwargs["payload"]["messages"] == [{"role": "user", "content": "Hello"}] # Verify response is raw OpenAI format assert response["id"] == "chatcmpl-123" assert response["model"] == "gpt-4o" assert response["choices"][0]["message"]["content"] == "Hello from passthrough!" @pytest.mark.asyncio async def test_openai_passthrough_embeddings(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-embed-passthrough-key", secret_value={"api_key": "sk-test-embed"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-embed-passthrough-model", secret_id=secret.secret_id, provider="openai", model_name="text-embedding-3-small", ) store.create_gateway_endpoint( name="openai-embed-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock OpenAI API response mock_response = { "object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], "model": "text-embedding-3-small", "usage": {"prompt_tokens": 5, "total_tokens": 5}, } # Create mock request mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-embed-passthrough-endpoint", "input": "Test input", } ) # Mock send_request directly with mock.patch( "mlflow.gateway.providers.openai.send_request", return_value=mock_response ) as mock_send: response = await openai_passthrough_embeddings(mock_request) # Verify send_request was called assert mock_send.called call_kwargs = mock_send.call_args[1] assert call_kwargs["path"] == "embeddings" assert call_kwargs["payload"]["model"] == "text-embedding-3-small" assert call_kwargs["payload"]["input"] == "Test input" # Verify response is raw OpenAI format assert response["model"] == "text-embedding-3-small" assert response["data"][0]["embedding"] == [0.1, 0.2, 0.3] @pytest.mark.asyncio async def test_openai_passthrough_responses(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-responses-key", secret_value={"api_key": "sk-test-responses"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-responses-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) store.create_gateway_endpoint( name="openai-responses-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock OpenAI Responses API response (using correct Responses API schema) mock_response = { "id": "resp-123", "object": "response", "created": 1234567890, "model": "gpt-4o", "status": "completed", "output": [ { "role": "assistant", "content": [{"type": "output_text", "text": "Response from Responses API"}], } ], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, } # Create mock request mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-responses-endpoint", "input": [{"role": "user", "content": "Hello"}], "instructions": "You are a helpful assistant", "response_format": {"type": "text"}, } ) # Mock send_request directly with mock.patch( "mlflow.gateway.providers.openai.send_request", return_value=mock_response ) as mock_send: response = await openai_passthrough_responses(mock_request) # Verify send_request was called assert mock_send.called call_kwargs = mock_send.call_args[1] assert call_kwargs["path"] == "responses" assert call_kwargs["payload"]["model"] == "gpt-4o" assert call_kwargs["payload"]["input"] == [{"role": "user", "content": "Hello"}] assert call_kwargs["payload"]["instructions"] == "You are a helpful assistant" assert call_kwargs["payload"]["response_format"] == {"type": "text"} # Verify response is raw OpenAI Responses API format assert response["id"] == "resp-123" assert response["object"] == "response" assert response["status"] == "completed" assert response["output"][0]["content"][0]["text"] == "Response from Responses API" @pytest.mark.asyncio async def test_openai_passthrough_responses_compact(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-responses-compact-key", secret_value={"api_key": "sk-test-responses-compact"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-responses-compact-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) store.create_gateway_endpoint( name="openai-responses-compact-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Mock OpenAI Responses-compact response: same shape as /responses output mock_response = { "id": "resp-compact-123", "object": "response", "created": 1234567890, "model": "gpt-4o", "status": "completed", "output": [ { "role": "assistant", "content": [{"type": "output_text", "text": "Compacted summary"}], } ], "usage": {"input_tokens": 50, "output_tokens": 10, "total_tokens": 60}, } # Compaction request — typical body carries `previous_response_id` and `model` mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-responses-compact-endpoint", "previous_response_id": "resp_abc123", } ) with mock.patch( "mlflow.gateway.providers.openai.send_request", return_value=mock_response ) as mock_send: response = await openai_passthrough_responses_compact(mock_request) # Verify send_request was called with the /compact upstream path assert mock_send.called call_kwargs = mock_send.call_args[1] assert call_kwargs["path"] == "responses/compact" assert call_kwargs["payload"]["model"] == "gpt-4o" assert call_kwargs["payload"]["previous_response_id"] == "resp_abc123" # Verify response is the raw OpenAI Responses-compact format assert response["id"] == "resp-compact-123" assert response["object"] == "response" assert response["status"] == "completed" assert response["output"][0]["content"][0]["text"] == "Compacted summary" @pytest.mark.asyncio async def test_openai_passthrough_responses_compact_rejects_stream(): """``/responses/compact`` is unary upstream; the handler must reject a client-supplied ``stream=true`` with HTTP 400 before invoking the provider (whose passthrough machinery treats all non-embeddings actions as stream-capable and would otherwise open an SSE stream against an upstream endpoint that does not support it). """ mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-responses-compact-endpoint", "previous_response_id": "resp_abc123", "stream": True, } ) # send_request should never be called — the handler rejects before # reaching the provider. with ( mock.patch("mlflow.gateway.providers.openai.send_request") as mock_send, pytest.raises(HTTPException, match="stream=true is not supported") as exc_info, ): await openai_passthrough_responses_compact(mock_request) assert exc_info.value.status_code == 400 assert not mock_send.called @pytest.mark.asyncio async def test_openai_passthrough_chat_streaming(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-stream-passthrough-key", secret_value={"api_key": "sk-test-stream"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-stream-passthrough-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) store.create_gateway_endpoint( name="openai-stream-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create mock request with streaming enabled mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-stream-passthrough-endpoint", "messages": [{"role": "user", "content": "Hello"}], "stream": True, } ) # Mock streaming response chunks mock_stream_chunks = [ b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}\n\n', # noqa: E501 b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}\n\n', # noqa: E501 b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', # noqa: E501 ] async def mock_stream_generator(): for chunk in mock_stream_chunks: yield chunk with mock.patch( "mlflow.gateway.providers.openai.send_stream_request", return_value=mock_stream_generator(), ) as mock_send_stream: response = await openai_passthrough_chat(mock_request) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" chunks = [chunk async for chunk in response.body_iterator] assert mock_send_stream.called assert len(chunks) == 3 assert b"Hello" in chunks[0] assert b"world" in chunks[1] assert b"stop" in chunks[2] @pytest.mark.asyncio async def test_openai_passthrough_responses_streaming(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-responses-stream-key", secret_value={"api_key": "sk-test-responses-stream"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-responses-stream-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) store.create_gateway_endpoint( name="openai-responses-stream-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) # Create mock request with streaming enabled mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "openai-responses-stream-endpoint", "input": [{"type": "text", "text": "Hello"}], "instructions": "You are a helpful assistant", "stream": True, } ) # Mock streaming response chunks for Responses API mock_stream_chunks = [ b'data: {"type":"response.created","response":{"id":"resp_1","object":"response","created_at":1741290958,"status":"in_progress","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":null,"user":null,"metadata":{}}}\n\n', # noqa: E501 b'data: {"type":"response.output_item.added","output_index":0,"item":{"id":"msg_1","type":"message","status":"in_progress","role":"assistant","content":[]}}\n\n', # noqa: E501 b'data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":"","annotations":[]}}\n\n', # noqa: E501 b'data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hi"}\n\n', # noqa: E501 b'data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hi there! How can I assist you today?"}\n\n', # noqa: E501 b'data: {"type":"response.content_part.done","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}}\n\n', # noqa: E501 b'data: {"type":"response.output_item.done","output_index":0,"item":{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}]}}\n\n', # noqa: E501 b'data: {"type":"response.completed","response":{"id":"resp_1","object":"response","created_at":1741290958,"status":"completed","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}]}],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":{"input_tokens":37,"output_tokens":11,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":48},"user":null,"metadata":{}}}\n\n', # noqa: E501 ] async def mock_stream_generator(): for chunk in mock_stream_chunks: yield chunk with mock.patch( "mlflow.gateway.providers.openai.send_stream_request", return_value=mock_stream_generator(), ) as mock_send_stream: response = await openai_passthrough_responses(mock_request) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" chunks = [chunk async for chunk in response.body_iterator] assert mock_send_stream.called assert len(chunks) == 8 assert b"response.created" in chunks[0] assert b"response.output_item.added" in chunks[1] assert b"response.content_part.added" in chunks[2] assert b"response.output_text.delta" in chunks[3] assert b"response.output_text.done" in chunks[4] assert b"response.content_part.done" in chunks[5] assert b"response.output_item.done" in chunks[6] assert b"response.completed" in chunks[7] # ============================================================================= # Anthropic Messages API passthrough endpoint tests # ============================================================================= @pytest.mark.asyncio async def test_anthropic_passthrough_messages(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="anthropic-passthrough-key", secret_value={"api_key": "sk-ant-test"}, provider="anthropic", ) model_def = store.create_gateway_model_definition( name="anthropic-passthrough-model", secret_id=secret.secret_id, provider="anthropic", model_name="claude-3-5-sonnet-20241022", ) store.create_gateway_endpoint( name="anthropic-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "anthropic-passthrough-endpoint", "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 1024, } ) mock_response = { "id": "msg_01XFDUDYJgAACzvnptvVoYEL", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello! How can I assist you today?"}], "model": "claude-3-5-sonnet-20241022", "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 20}, } with mock.patch( "mlflow.gateway.providers.anthropic.send_request", return_value=mock_response ) as mock_send: response = await anthropic_passthrough_messages(mock_request) assert mock_send.called call_args = mock_send.call_args assert call_args[1]["path"] == "messages" assert call_args[1]["payload"]["model"] == "claude-3-5-sonnet-20241022" assert call_args[1]["payload"]["messages"] == [{"role": "user", "content": "Hello"}] assert call_args[1]["payload"]["max_tokens"] == 1024 assert response["id"] == "msg_01XFDUDYJgAACzvnptvVoYEL" assert response["model"] == "claude-3-5-sonnet-20241022" assert response["content"][0]["text"] == "Hello! How can I assist you today?" @pytest.mark.asyncio async def test_anthropic_passthrough_messages_streaming(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="anthropic-stream-passthrough-key", secret_value={"api_key": "sk-ant-test-stream"}, provider="anthropic", ) model_def = store.create_gateway_model_definition( name="anthropic-stream-passthrough-model", secret_id=secret.secret_id, provider="anthropic", model_name="claude-3-5-sonnet-20241022", ) store.create_gateway_endpoint( name="anthropic-stream-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": "anthropic-stream-passthrough-endpoint", "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 1024, "stream": True, } ) mock_stream_chunks = [ b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_01XFDUDYJgAACzvnptvVoYEL","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet-20241022","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}\n\n', # noqa: E501 b'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n', # noqa: E501 b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}\n\n', # noqa: E501 b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"!"}}\n\n', # noqa: E501 b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":20}}\n\n', # noqa: E501 b'event: message_stop\ndata: {"type":"message_stop"}\n\n', ] async def mock_stream_generator(): for chunk in mock_stream_chunks: yield chunk with mock.patch( "mlflow.gateway.providers.anthropic.send_stream_request", return_value=mock_stream_generator(), ) as mock_send_stream: response = await anthropic_passthrough_messages(mock_request) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" chunks = [chunk async for chunk in response.body_iterator] assert mock_send_stream.called assert len(chunks) == 7 assert b"message_start" in chunks[0] assert b"content_block_start" in chunks[1] assert b"content_block_delta" in chunks[2] assert b"Hello" in chunks[2] assert b"content_block_delta" in chunks[3] assert b"!" in chunks[3] assert b"content_block_stop" in chunks[4] assert b"message_delta" in chunks[5] assert b"message_stop" in chunks[6] # ============================================================================= # Gemini generateContent/streamGenerateContent passthrough endpoint tests # ============================================================================= @pytest.mark.asyncio async def test_gemini_passthrough_generate_content(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="gemini-passthrough-key", secret_value={"api_key": "test-key"}, provider="gemini", ) model_def = store.create_gateway_model_definition( name="gemini-passthrough-model", secret_id=secret.secret_id, provider="gemini", model_name="gemini-2.0-flash", ) store.create_gateway_endpoint( name="gemini-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "contents": [ { "role": "user", "parts": [{"text": "Hello"}], } ] } ) mock_response = { "candidates": [ { "content": { "parts": [{"text": "Hello! How can I assist you today?"}], "role": "model", }, "finishReason": "STOP", } ], "usageMetadata": { "promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15, }, } with mock.patch( "mlflow.gateway.providers.gemini.send_request", return_value=mock_response ) as mock_send: response = await gemini_passthrough_generate_content( "gemini-passthrough-endpoint", mock_request ) assert mock_send.called call_args = mock_send.call_args assert call_args[1]["path"] == "gemini-2.0-flash:generateContent" assert call_args[1]["payload"]["contents"] == [ {"role": "user", "parts": [{"text": "Hello"}]} ] assert ( response["candidates"][0]["content"]["parts"][0]["text"] == "Hello! How can I assist you today?" ) assert response["usageMetadata"]["totalTokenCount"] == 15 @pytest.mark.asyncio async def test_gemini_passthrough_stream_generate_content(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="gemini-stream-passthrough-key", secret_value={"api_key": "test-stream-key"}, provider="gemini", ) model_def = store.create_gateway_model_definition( name="gemini-stream-passthrough-model", secret_id=secret.secret_id, provider="gemini", model_name="gemini-2.0-flash", ) store.create_gateway_endpoint( name="gemini-stream-passthrough-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "contents": [ { "role": "user", "parts": [{"text": "Hello"}], } ] } ) mock_stream_chunks = [ b'data: {"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"}}]}\n\n', b'data: {"candidates":[{"content":{"parts":[{"text":"!"}],"role":"model"}}]}\n\n', b'data: {"candidates":[{"content":{"parts":[{"text":" How can I help you?"}],"role":"model"},"finishReason":"STOP"}]}\n\n', # noqa: E501 ] async def mock_stream_generator(): for chunk in mock_stream_chunks: yield chunk with mock.patch( "mlflow.gateway.providers.gemini.send_stream_request", return_value=mock_stream_generator(), ) as mock_send_stream: response = await gemini_passthrough_stream_generate_content( "gemini-stream-passthrough-endpoint", mock_request ) assert isinstance(response, StreamingResponse) assert response.media_type == "text/event-stream" chunks = [chunk async for chunk in response.body_iterator] assert mock_send_stream.called assert len(chunks) == 3 assert b"Hello" in chunks[0] assert b"!" in chunks[1] assert b"How can I help you?" in chunks[2] assert b"STOP" in chunks[2] def test_create_fallback_provider_single_model(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-fallback-key", secret_value={"api_key": "sk-test-key"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="gpt-fallback-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-fallback-single-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.FALLBACK, weight=1.0, fallback_order=0, ), ], routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT, fallback_config=FallbackConfig( strategy=FallbackStrategy.SEQUENTIAL, max_attempts=1, ), ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, FallbackProvider) assert len(provider._providers) == 2 assert isinstance(provider._providers[0], TrafficRouteProvider) assert isinstance(provider._providers[1], OpenAIProvider) assert provider._max_attempts == 2 def test_create_fallback_provider_multiple_models(store: SqlAlchemyStore): secret1 = store.create_gateway_secret( secret_name="openai-primary-key", secret_value={"api_key": "sk-primary-key"}, provider="openai", ) model_def1 = store.create_gateway_model_definition( name="gpt-primary-model", secret_id=secret1.secret_id, provider="openai", model_name="gpt-4", ) secret2 = store.create_gateway_secret( secret_name="anthropic-fallback-key", secret_value={"api_key": "sk-ant-fallback"}, provider="anthropic", ) model_def2 = store.create_gateway_model_definition( name="claude-fallback-model", secret_id=secret2.secret_id, provider="anthropic", model_name="claude-3-sonnet", ) endpoint = store.create_gateway_endpoint( name="test-fallback-multi-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def1.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), GatewayEndpointModelConfig( model_definition_id=model_def2.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), GatewayEndpointModelConfig( model_definition_id=model_def1.model_definition_id, linkage_type=GatewayModelLinkageType.FALLBACK, weight=1.0, fallback_order=0, ), GatewayEndpointModelConfig( model_definition_id=model_def2.model_definition_id, linkage_type=GatewayModelLinkageType.FALLBACK, weight=1.0, fallback_order=1, ), ], routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT, fallback_config=FallbackConfig( strategy=FallbackStrategy.SEQUENTIAL, max_attempts=2, ), ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, FallbackProvider) assert len(provider._providers) == 3 primary = provider._providers[0] assert isinstance(primary, TrafficRouteProvider) assert isinstance(primary._providers[0], OpenAIProvider) assert isinstance(primary._providers[1], AnthropicProvider) assert isinstance(provider._providers[1], OpenAIProvider) assert isinstance(provider._providers[2], AnthropicProvider) assert provider._max_attempts == 3 def test_create_fallback_provider_max_attempts_exceeds_providers(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-fallback-key", secret_value={"api_key": "sk-test-key"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="gpt-fallback-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-fallback-max-attempts-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.FALLBACK, weight=1.0, fallback_order=0, ), ], routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT, fallback_config=FallbackConfig( strategy=FallbackStrategy.SEQUENTIAL, max_attempts=10, ), ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) # FallbackProvider is the outer provider, individual providers inside are wrapped assert isinstance(provider, FallbackProvider) assert provider._max_attempts == 2 def test_create_fallback_provider_no_max_attempts(store: SqlAlchemyStore): secret1 = store.create_gateway_secret( secret_name="openai-primary-key", secret_value={"api_key": "sk-primary-key"}, provider="openai", ) model_def1 = store.create_gateway_model_definition( name="gpt-primary-model", secret_id=secret1.secret_id, provider="openai", model_name="gpt-4", ) secret2 = store.create_gateway_secret( secret_name="anthropic-fallback-key", secret_value={"api_key": "sk-ant-fallback"}, provider="anthropic", ) model_def2 = store.create_gateway_model_definition( name="claude-fallback-model", secret_id=secret2.secret_id, provider="anthropic", model_name="claude-3-sonnet", ) endpoint = store.create_gateway_endpoint( name="test-fallback-no-max-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def1.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), GatewayEndpointModelConfig( model_definition_id=model_def2.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), GatewayEndpointModelConfig( model_definition_id=model_def1.model_definition_id, linkage_type=GatewayModelLinkageType.FALLBACK, weight=1.0, fallback_order=0, ), GatewayEndpointModelConfig( model_definition_id=model_def2.model_definition_id, linkage_type=GatewayModelLinkageType.FALLBACK, weight=1.0, fallback_order=1, ), ], routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT, fallback_config=FallbackConfig( strategy=FallbackStrategy.SEQUENTIAL, max_attempts=None, ), ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) # FallbackProvider is the outer provider, individual providers inside are wrapped assert isinstance(provider, FallbackProvider) assert len(provider._providers) == 3 assert provider._max_attempts == 3 def test_create_provider_default_routing_single_model(store: SqlAlchemyStore): secret = store.create_gateway_secret( secret_name="openai-default-key", secret_value={"api_key": "sk-test-key"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="gpt-default-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name="test-default-routing-endpoint", model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], ) provider, _ = _create_provider_from_endpoint_name( store, endpoint.name, EndpointType.LLM_V1_CHAT ) assert isinstance(provider, OpenAIProvider) assert not isinstance(provider, FallbackProvider) # ============================================================================= # Gateway Tracing Tests # ============================================================================= async def _call_invocations(endpoint_name: str, request, payload: dict[str, Any]): # invocations doesn't use "model" field - endpoint is in URL payload_without_model = {k: v for k, v in payload.items() if k != "model"} request.json = AsyncMock(return_value=payload_without_model) return await invocations(endpoint_name, request) async def _call_chat_completions(endpoint_name: str, request, payload: dict[str, Any]): # chat_completions uses "model" field to specify endpoint request.json = AsyncMock(return_value=payload) return await chat_completions(request) @pytest.mark.asyncio @pytest.mark.parametrize( "handler", [_call_invocations, _call_chat_completions], ids=["invocations", "chat_completions"] ) async def test_gateway_creates_trace_with_usage(store: SqlAlchemyStore, handler): endpoint_name = "tracing-test-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") # Create endpoint with usage tracking enabled secret = store.create_gateway_secret( secret_name="tracing-test-key", secret_value={"api_key": "sk-test-tracing"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="tracing-test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() payload = { "model": endpoint_name, "messages": [{"role": "user", "content": "Hi"}], "stream": False, } # Mock the OpenAI send_request to return our mock response with mock.patch( "mlflow.gateway.providers.openai.send_request", return_value={ "id": "test-id", "object": "chat.completion", "created": 1234567890, "model": "gpt-4", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, }, ): response = await handler(endpoint_name, mock_request, payload) assert response.id == "test-id" assert response.choices[0].message.content == "Hello!" # Verify trace was created traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] assert trace.info.state == TraceState.OK # Verify gateway metadata is present in trace assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.UNIFIED_CHAT ) # Verify span has provider information (provider name is lowercased for cost lookup alignment) span_names = {span.name for span in trace.data.spans} assert "provider/openai/gpt-4" in span_names # Find the provider span and check attributes provider_span = next( (span for span in trace.data.spans if span.name == "provider/openai/gpt-4"), None ) assert provider_span is not None assert provider_span.attributes.get(SpanAttributeKey.MODEL_PROVIDER) == "openai" assert provider_span.attributes.get(SpanAttributeKey.MODEL) == "gpt-4" # Verify token usage is captured on the provider span token_usage = provider_span.attributes.get(SpanAttributeKey.CHAT_USAGE) assert token_usage is not None assert token_usage[TokenUsageKey.INPUT_TOKENS] == 10 assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5 assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 15 # Verify trace metadata has aggregated token usage (auto-generated from span attributes) trace_token_usage = json.loads(trace.info.trace_metadata.get(TraceMetadataKey.TOKEN_USAGE)) assert trace_token_usage[TokenUsageKey.INPUT_TOKENS] == 10 assert trace_token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5 assert trace_token_usage[TokenUsageKey.TOTAL_TOKENS] == 15 @pytest.mark.asyncio @pytest.mark.parametrize( "handler", [_call_invocations, _call_chat_completions], ids=["invocations", "chat_completions"] ) async def test_gateway_streaming_creates_trace(store: SqlAlchemyStore, handler): endpoint_name = "stream-tracing-test-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") # Create endpoint with usage tracking enabled secret = store.create_gateway_secret( secret_name="stream-tracing-test-key", secret_value={"api_key": "sk-test-stream-tracing"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="stream-tracing-test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() payload = { "model": endpoint_name, "messages": [{"role": "user", "content": "Hi"}], "stream": True, } # Mock streaming response chunks with usage in the final chunk mock_stream_chunks = [ b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}\n\n', # noqa: E501 b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]}\n\n', # noqa: E501 b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}\n\n', # noqa: E501 b"data: [DONE]\n\n", ] async def mock_stream_generator(): for chunk in mock_stream_chunks: yield chunk with mock.patch( "mlflow.gateway.providers.openai.send_stream_request", return_value=mock_stream_generator(), ): response = await handler(endpoint_name, mock_request, payload) # Verify streaming response is returned assert isinstance(response, StreamingResponse) # Consume the response chunks = [chunk async for chunk in response.body_iterator] assert len(chunks) > 0 # Verify trace was created for the gateway invocation traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] assert trace.info.state == TraceState.OK # Verify gateway metadata is present in trace assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.UNIFIED_CHAT ) # Verify gateway span exists with correct attributes gateway_span = next( (span for span in trace.data.spans if span.name == f"gateway/{endpoint_name}"), None ) assert gateway_span is not None assert gateway_span.attributes.get("endpoint_name") == endpoint_name # Verify that streaming output is aggregated into a ChatCompletion-like response output = gateway_span.outputs assert output is not None assert output["object"] == "chat.completion" assert output["id"] == "chatcmpl-123" assert output["model"] == "gpt-4" assert len(output["choices"]) == 1 assert output["choices"][0]["index"] == 0 assert output["choices"][0]["message"]["role"] == "assistant" assert output["choices"][0]["message"]["content"] == "Hello!" assert output["choices"][0]["finish_reason"] == "stop" assert output["usage"]["prompt_tokens"] == 10 assert output["usage"]["completion_tokens"] == 5 assert output["usage"]["total_tokens"] == 15 @pytest.mark.asyncio @pytest.mark.parametrize( "handler", [_call_invocations, _call_chat_completions], ids=["invocations", "chat_completions"] ) async def test_gateway_trace_includes_user_metadata(store: SqlAlchemyStore, handler): endpoint_name = "user-metadata-tracing-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") # Create endpoint with usage tracking enabled secret = store.create_gateway_secret( secret_name="user-metadata-tracing-key", secret_value={"api_key": "sk-test-user-metadata"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="user-metadata-tracing-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) # Create mock request with user metadata set (as auth middleware would do) mock_request = create_mock_request(username="test_user", user_id=42) payload = { "model": endpoint_name, "messages": [{"role": "user", "content": "Hi"}], "stream": False, } # Mock the OpenAI send_request to return our mock response with mock.patch( "mlflow.gateway.providers.openai.send_request", return_value={ "id": "test-id", "object": "chat.completion", "created": 1234567890, "model": "gpt-4", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, }, ): response = await handler(endpoint_name, mock_request, payload) assert response.id == "test-id" assert response.choices[0].message.content == "Hello!" # Verify trace was created traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] assert trace.info.state == TraceState.OK # Verify user metadata is present in trace info assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USERNAME) == "test_user" assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USER_ID) == "42" # Verify gateway metadata is present alongside user metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.UNIFIED_CHAT ) # Verify span attributes still include endpoint info gateway_span = next( (span for span in trace.data.spans if span.name == f"gateway/{endpoint_name}"), None ) assert gateway_span is not None assert gateway_span.attributes.get("endpoint_name") == endpoint_name # ============================================================================= # Passthrough Token Usage Tracking Tests # ============================================================================= @pytest.mark.asyncio async def test_openai_passthrough_chat_token_usage_tracking(store: SqlAlchemyStore): endpoint_name = "openai-passthrough-usage-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name="openai-passthrough-usage-key", secret_value={"api_key": "sk-test-usage"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-passthrough-usage-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": endpoint_name, "messages": [{"role": "user", "content": "Hello"}], } ) mock_response = { "id": "chatcmpl-123", "object": "chat.completion", "created": 1234567890, "model": "gpt-4o", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, } with mock.patch("mlflow.gateway.providers.openai.send_request", return_value=mock_response): response = await openai_passthrough_chat(mock_request) assert response["usage"]["prompt_tokens"] == 10 assert response["usage"]["completion_tokens"] == 5 assert response["usage"]["total_tokens"] == 15 # Verify trace was created with token usage traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] assert trace.info.state == TraceState.OK # Verify gateway metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_CHAT ) # Find the passthrough span and check token usage attributes passthrough_span = next( (span for span in trace.data.spans if "action" in span.attributes), None ) assert passthrough_span is not None assert passthrough_span.attributes.get("action") == "openai_chat" token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE) assert token_usage is not None assert token_usage[TokenUsageKey.INPUT_TOKENS] == 10 assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5 assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 15 @pytest.mark.asyncio async def test_openai_passthrough_embeddings_token_usage_tracking(store: SqlAlchemyStore): endpoint_name = "openai-embed-usage-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name="openai-embed-usage-key", secret_value={"api_key": "sk-test-embed-usage"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-embed-usage-model", secret_id=secret.secret_id, provider="openai", model_name="text-embedding-3-small", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": endpoint_name, "input": "Test text for embedding", } ) mock_response = { "object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], "model": "text-embedding-3-small", "usage": {"prompt_tokens": 5, "total_tokens": 5}, } with mock.patch("mlflow.gateway.providers.openai.send_request", return_value=mock_response): response = await openai_passthrough_embeddings(mock_request) assert response["usage"]["prompt_tokens"] == 5 assert response["usage"]["total_tokens"] == 5 # Verify trace was created with token usage traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] # Verify gateway metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_EMBEDDINGS ) # Find the passthrough span and check token usage attributes passthrough_span = next( (span for span in trace.data.spans if "action" in span.attributes), None ) assert passthrough_span is not None assert passthrough_span.attributes.get("action") == "openai_embeddings" token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE) assert token_usage is not None assert token_usage[TokenUsageKey.INPUT_TOKENS] == 5 assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 5 @pytest.mark.asyncio async def test_openai_passthrough_responses_token_usage_tracking(store: SqlAlchemyStore): endpoint_name = "openai-responses-usage-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name="openai-responses-usage-key", secret_value={"api_key": "sk-test-responses-usage"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-responses-usage-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": endpoint_name, "input": [{"role": "user", "content": "Hello"}], } ) mock_response = { "id": "resp-123", "object": "response", "created": 1234567890, "model": "gpt-4o", "status": "completed", "output": [ { "role": "assistant", "content": [{"type": "output_text", "text": "Hello!"}], } ], "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, } with mock.patch("mlflow.gateway.providers.openai.send_request", return_value=mock_response): response = await openai_passthrough_responses(mock_request) assert response["usage"]["input_tokens"] == 10 assert response["usage"]["output_tokens"] == 5 assert response["usage"]["total_tokens"] == 15 # Verify trace was created with token usage traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] assert trace.info.state == TraceState.OK # Verify gateway metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_RESPONSES ) # Find the passthrough span and check token usage attributes passthrough_span = next( (span for span in trace.data.spans if "action" in span.attributes), None ) assert passthrough_span is not None assert passthrough_span.attributes.get("action") == "openai_responses" token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE) assert token_usage is not None assert token_usage[TokenUsageKey.INPUT_TOKENS] == 10 assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5 assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 15 @pytest.mark.asyncio async def test_anthropic_passthrough_messages_token_usage_tracking(store: SqlAlchemyStore): endpoint_name = "anthropic-usage-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name="anthropic-usage-key", secret_value={"api_key": "sk-ant-usage"}, provider="anthropic", ) model_def = store.create_gateway_model_definition( name="anthropic-usage-model", secret_id=secret.secret_id, provider="anthropic", model_name="claude-3-5-sonnet-20241022", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": endpoint_name, "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 1024, } ) mock_response = { "id": "msg_01XFDUDYJgAACzvnptvVoYEL", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Hello!"}], "model": "claude-3-5-sonnet-20241022", "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 12, "output_tokens": 8}, } with mock.patch("mlflow.gateway.providers.anthropic.send_request", return_value=mock_response): response = await anthropic_passthrough_messages(mock_request) assert response["usage"]["input_tokens"] == 12 assert response["usage"]["output_tokens"] == 8 # Verify trace was created with token usage traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] # Verify gateway metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.PASSTHROUGH_MODEL_ANTHROPIC_MESSAGES ) # Find the passthrough span and check token usage attributes passthrough_span = next( (span for span in trace.data.spans if "action" in span.attributes), None ) assert passthrough_span is not None assert passthrough_span.attributes.get("action") == "anthropic_messages" token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE) assert token_usage is not None assert token_usage[TokenUsageKey.INPUT_TOKENS] == 12 assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 8 # Anthropic doesn't provide total_tokens, so we calculate it assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 20 @pytest.mark.asyncio async def test_gemini_passthrough_generate_content_token_usage_tracking(store: SqlAlchemyStore): endpoint_name = "gemini-usage-endpoint" # Create experiment for tracing experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name="gemini-usage-key", secret_value={"api_key": "test-gemini-usage-key"}, provider="gemini", ) model_def = store.create_gateway_model_definition( name="gemini-usage-model", secret_id=secret.secret_id, provider="gemini", model_name="gemini-2.0-flash", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "contents": [ { "role": "user", "parts": [{"text": "Hello"}], } ] } ) mock_response = { "candidates": [ { "content": { "parts": [{"text": "Hello! How can I help?"}], "role": "model", }, "finishReason": "STOP", } ], "usageMetadata": { "promptTokenCount": 7, "candidatesTokenCount": 9, "totalTokenCount": 16, }, } with mock.patch("mlflow.gateway.providers.gemini.send_request", return_value=mock_response): response = await gemini_passthrough_generate_content(endpoint_name, mock_request) assert response["usageMetadata"]["promptTokenCount"] == 7 assert response["usageMetadata"]["candidatesTokenCount"] == 9 assert response["usageMetadata"]["totalTokenCount"] == 16 # Verify trace was created with token usage traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] # Verify gateway metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.PASSTHROUGH_MODEL_GEMINI_GENERATE_CONTENT ) # Find the passthrough span and check token usage attributes passthrough_span = next( (span for span in trace.data.spans if "action" in span.attributes), None ) assert passthrough_span is not None assert passthrough_span.attributes.get("action") == "gemini_generate_content" token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE) assert token_usage is not None assert token_usage[TokenUsageKey.INPUT_TOKENS] == 7 assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 9 assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 16 @pytest.mark.asyncio async def test_openai_passthrough_streaming_captures_chunks(store: SqlAlchemyStore): endpoint_name = "openai-passthrough-streaming-chunks" experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name="openai-stream-chunks-key", secret_value={"api_key": "sk-test-stream-chunks"}, provider="openai", ) model_def = store.create_gateway_model_definition( name="openai-stream-chunks-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4o", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ), ], usage_tracking=True, experiment_id=experiment_id, ) mock_request = create_mock_request() mock_request.json = AsyncMock( return_value={ "model": endpoint_name, "messages": [{"role": "user", "content": "Hello"}], "stream": True, } ) mock_request.headers = {} mock_stream_chunks = [ b'data: {"id":"chatcmpl-123","choices":[{"delta":{"content":"Hi"}}]}\n\n', b'data: {"id":"chatcmpl-123","choices":[{"delta":{"content":"!"}}]}\n\n', b"data: [DONE]\n\n", ] async def mock_stream_generator(): for chunk in mock_stream_chunks: yield chunk with mock.patch( "mlflow.gateway.providers.openai.send_stream_request", return_value=mock_stream_generator(), ): response = await openai_passthrough_chat(mock_request) assert isinstance(response, StreamingResponse) chunks = [chunk async for chunk in response.body_iterator] assert len(chunks) == len(mock_stream_chunks) # Verify trace was created traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 trace = traces[0] assert trace.info.state == TraceState.OK # Verify gateway metadata assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID) == endpoint.endpoint_id ) assert ( trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE) == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_CHAT ) gateway_span = next( (span for span in trace.data.spans if span.name == f"gateway/{endpoint_name}"), None ) assert gateway_span is not None # Verify streaming chunks are captured in outputs (raw SSE bytes decoded to strings) assert gateway_span.outputs is not None assert len(gateway_span.outputs) == len(mock_stream_chunks) # Verify the outputs contain actual SSE data (not async generator object repr) assert "data:" in gateway_span.outputs[0] assert "chatcmpl-123" in gateway_span.outputs[0] # ─── Guardrail end-to-end scenarios ────────────────────────────────────────── class _SimpleScorer: """Minimal scorer that returns 'yes' or 'no' and tracks call count.""" def __init__(self, *, passing: bool = True) -> None: self.call_count = 0 self._passing = passing def __call__(self, **kwargs) -> str: self.call_count += 1 return "yes" if self._passing else "no" def _make_guardrail_judge(stage, action=GuardrailAction.VALIDATION, *, passing=True): scorer = _SimpleScorer(passing=passing) return JudgeGuardrail( scorer=scorer, stage=GuardrailStage(stage), action=GuardrailAction(action), name=f"test-{stage.lower()}", ) def _make_guardrail_chat_response(content: str = "Hello!") -> chat.ResponsePayload: return chat.ResponsePayload( id="resp-id", object="chat.completion", created=1234567890, model="gpt-4", choices=[ chat.Choice( index=0, message=chat.ResponseMessage(role="assistant", content=content), finish_reason="stop", ) ], usage=chat.ChatUsage(prompt_tokens=5, completion_tokens=5, total_tokens=10), ) def _make_guardrail_mock_request(body: dict[str, Any], headers: dict[str, str] | None = None): req = MagicMock() req.state.cached_body = None req.state.username = None req.state.user_id = None req.json = AsyncMock(return_value=body) req.headers = headers or {} req.base_url = "http://localhost:5000/" return req _GUARDRAIL_SERIALIZED_SCORER = json.dumps({"name": "safety", "builtin_scorer_class": "Safety"}) def _setup_db_guardrail( store: SqlAlchemyStore, endpoint_name: str, stage: str, action: str, action_endpoint_name: str | None = None, execution_order: int | None = None, name: str | None = None, ): """Create scorer + guardrail in DB and attach it to the endpoint.""" guardrail_name = name or f"guardrail-{endpoint_name}-{stage}" experiment_id = store.create_experiment(f"exp-{guardrail_name}") scorer_ver = store.register_scorer( experiment_id, f"scorer-{guardrail_name}", _GUARDRAIL_SERIALIZED_SCORER ) action_endpoint_id = None if action_endpoint_name: action_endpoint_id = store.get_gateway_endpoint(name=action_endpoint_name).endpoint_id guardrail = store.create_gateway_guardrail( name=guardrail_name, scorer_id=scorer_ver.scorer_id, scorer_version=scorer_ver.scorer_version, stage=GuardrailStage(stage), action=GuardrailAction(action), action_endpoint_id=action_endpoint_id, ) endpoint = store.get_gateway_endpoint(name=endpoint_name) store.add_guardrail_to_endpoint( endpoint.endpoint_id, guardrail.guardrail_id, execution_order=execution_order ) return guardrail, scorer_ver def _setup_guardrail_endpoint(store: SqlAlchemyStore, name: str): secret = store.create_gateway_secret( secret_name=f"key-{name}", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name=f"model-{name}", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) return store.create_gateway_endpoint( name=name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ) ], ) @pytest.mark.asyncio async def test_invocations_bypass_header_skips_guardrails(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-bypass") mock_response = _make_guardrail_chat_response("Bypass response") mock_request = _make_guardrail_mock_request( {"messages": [{"role": "user", "content": "hello"}]}, headers={_SANITIZE_BYPASS_HEADER: "1"}, ) with ( patch("mlflow.server.gateway_api._create_provider_from_endpoint_name") as mock_create, patch("mlflow.server.gateway_api.load_guardrails") as mock_load, ): mock_provider = MagicMock() mock_provider.chat = AsyncMock(return_value=mock_response) mock_create.return_value = ( mock_provider, GatewayEndpointConfig( endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[] ), ) response = await invocations(endpoint.name, mock_request) mock_load.assert_not_called() assert response.choices[0].message.content == "Bypass response" @pytest.mark.asyncio async def test_invocations_bypass_header_wrong_value_runs_guardrails(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-bypass-wrong-val") _setup_db_guardrail(store, "ep-bypass-wrong-val", "BEFORE", "VALIDATION") mock_request = _make_guardrail_mock_request( {"messages": [{"role": "user", "content": "hello"}]}, headers={_SANITIZE_BYPASS_HEADER: "true"}, # wrong value — must not bypass ) blocking_scorer = _SimpleScorer(passing=False) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=blocking_scorer), patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock()), ): with pytest.raises(HTTPException, match="400"): await invocations(endpoint.name, mock_request) assert blocking_scorer.call_count == 1 @pytest.mark.asyncio async def test_real_db_pre_llm_guardrail_passes(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "real-ep-pre-llm-pass") _setup_db_guardrail(store, "real-ep-pre-llm-pass", "BEFORE", "VALIDATION") mock_response = _make_guardrail_chat_response("Safe response") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "hello"}] }) passing_scorer = _SimpleScorer(passing=True) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=passing_scorer), patch( "mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock(return_value=mock_response), ), ): response = await invocations(endpoint.name, mock_request) assert response.choices[0].message.content == "Safe response" assert passing_scorer.call_count == 1 @pytest.mark.asyncio async def test_real_db_pre_llm_guardrail_blocks(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "real-ep-pre-llm-block") _setup_db_guardrail(store, "real-ep-pre-llm-block", "BEFORE", "VALIDATION") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "bad input"}] }) blocking_scorer = _SimpleScorer(passing=False) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=blocking_scorer), patch( "mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock(), ) as mock_chat, ): with pytest.raises(HTTPException, match="400"): await invocations(endpoint.name, mock_request) assert not mock_chat.called assert blocking_scorer.call_count == 1 @pytest.mark.asyncio async def test_real_db_post_llm_guardrail_blocks(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "real-ep-post-llm-block") _setup_db_guardrail(store, "real-ep-post-llm-block", "AFTER", "VALIDATION") mock_response = _make_guardrail_chat_response("Unsafe output") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "hello"}] }) blocking_scorer = _SimpleScorer(passing=False) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=blocking_scorer), patch( "mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock(return_value=mock_response), ) as mock_chat, ): with pytest.raises(HTTPException, match="400"): await invocations(endpoint.name, mock_request) assert mock_chat.called assert blocking_scorer.call_count == 1 @pytest.mark.asyncio async def test_invocations_before_sanitize_rewrites_request(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-sanitize-before") sanitizer = _setup_guardrail_endpoint(store, "ep-sanitizer") _setup_db_guardrail( store, "ep-sanitize-before", "BEFORE", "SANITIZATION", action_endpoint_name=sanitizer.name ) sanitized_body = {"messages": [{"role": "user", "content": "cleaned input"}]} mock_response = _make_guardrail_chat_response("Response to cleaned input") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "bad input"}] }) failing_scorer = _SimpleScorer(passing=False) captured_payloads: list[Any] = [] async def fake_chat(payload): captured_payloads.append(payload) return mock_response with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer), patch( "mlflow.gateway.guardrails.send_request", AsyncMock( return_value={"choices": [{"message": {"content": json.dumps(sanitized_body)}}]} ), ), patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", side_effect=fake_chat), ): response = await invocations(endpoint.name, mock_request) assert response.choices[0].message.content == "Response to cleaned input" assert failing_scorer.call_count == 1 assert captured_payloads[0].messages[0].content == "cleaned input" @pytest.mark.asyncio async def test_invocations_after_sanitize_rewrites_response(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-sanitize-after") sanitizer = _setup_guardrail_endpoint(store, "ep-sanitizer-after") _setup_db_guardrail( store, "ep-sanitize-after", "AFTER", "SANITIZATION", action_endpoint_name=sanitizer.name ) sanitized_response = { "id": "resp-sanitized", "object": "chat.completion", "created": 1234567890, "model": "gpt-4", "choices": [ { "index": 0, "message": {"role": "assistant", "content": "polite output"}, "finish_reason": "stop", } ], "usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10}, } mock_response = _make_guardrail_chat_response("rude output") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "hello"}] }) failing_scorer = _SimpleScorer(passing=False) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer), patch( "mlflow.gateway.guardrails.send_request", AsyncMock( return_value={"choices": [{"message": {"content": json.dumps(sanitized_response)}}]} ), ), patch( "mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock(return_value=mock_response), ), ): response = await invocations(endpoint.name, mock_request) assert response.choices[0].message.content == "polite output" assert failing_scorer.call_count == 1 @pytest.mark.asyncio async def test_invocations_sanitize_no_action_endpoint_blocks(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-sanitize-no-ep") _setup_db_guardrail(store, "ep-sanitize-no-ep", "BEFORE", "SANITIZATION") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "bad input"}] }) failing_scorer = _SimpleScorer(passing=False) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer), patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock()), ): with pytest.raises(HTTPException, match="400"): await invocations(endpoint.name, mock_request) @pytest.mark.asyncio async def test_chat_completions_before_sanitize_rewrites_request(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-cc-sanitize-before") sanitizer = _setup_guardrail_endpoint(store, "ep-cc-sanitizer") _setup_db_guardrail( store, "ep-cc-sanitize-before", "BEFORE", "SANITIZATION", action_endpoint_name=sanitizer.name, ) sanitized_body = {"messages": [{"role": "user", "content": "cleaned input"}]} mock_response = _make_guardrail_chat_response("Response to cleaned input") mock_request = _make_guardrail_mock_request({ "model": endpoint.name, "messages": [{"role": "user", "content": "bad input"}], }) failing_scorer = _SimpleScorer(passing=False) captured_payloads: list[Any] = [] async def fake_chat(payload): captured_payloads.append(payload) return mock_response with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer), patch( "mlflow.gateway.guardrails.send_request", AsyncMock( return_value={"choices": [{"message": {"content": json.dumps(sanitized_body)}}]} ), ), patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", side_effect=fake_chat), ): from mlflow.server.gateway_api import chat_completions response = await chat_completions(mock_request) assert response.choices[0].message.content == "Response to cleaned input" assert failing_scorer.call_count == 1 assert captured_payloads[0].messages[0].content == "cleaned input" @pytest.mark.asyncio async def test_guardrails_run_in_execution_order(store: SqlAlchemyStore): endpoint = _setup_guardrail_endpoint(store, "ep-order-test") # Register in reverse order (5→1) to ensure DB insertion order != execution order. for i in range(5, 0, -1): _setup_db_guardrail( store, "ep-order-test", "BEFORE", "VALIDATION", execution_order=i, name=f"g-order-{i}", ) mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "hello"}] }) call_order: list[str] = [] def make_scorer(label: str, passing: bool): def scorer(**kwargs): call_order.append(label) return "yes" if passing else "no" return scorer # Guardrails 1-4 pass; guardrail 5 blocks — so all 5 must run in order 1→2→3→4→5. scorers = [make_scorer(f"order-{i}", passing=(i < 5)) for i in range(1, 6)] call_count = {"n": 0} def model_validate_side_effect(serialized): scorer = scorers[call_count["n"]] call_count["n"] += 1 return scorer with ( patch( "mlflow.genai.scorers.base.Scorer.model_validate", side_effect=model_validate_side_effect, ), patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock()), ): with pytest.raises(HTTPException, match="400"): await invocations(endpoint.name, mock_request) assert call_order == ["order-1", "order-2", "order-3", "order-4", "order-5"] @pytest.mark.asyncio async def test_guardrail_spans_created_when_usage_tracking_on(store: SqlAlchemyStore): endpoint_name = "ep-guardrail-tracing" experiment_id = store.create_experiment(f"gateway/{endpoint_name}") secret = store.create_gateway_secret( secret_name=f"key-{endpoint_name}", secret_value={"api_key": "sk-test"}, provider="openai", ) model_def = store.create_gateway_model_definition( name=f"model-{endpoint_name}", secret_id=secret.secret_id, provider="openai", model_name="gpt-4", ) endpoint = store.create_gateway_endpoint( name=endpoint_name, model_configs=[ GatewayEndpointModelConfig( model_definition_id=model_def.model_definition_id, linkage_type=GatewayModelLinkageType.PRIMARY, weight=1.0, ) ], usage_tracking=True, experiment_id=experiment_id, ) _setup_db_guardrail(store, endpoint_name, "BEFORE", "VALIDATION", name="safety-check") mock_response = _make_guardrail_chat_response("Safe response") mock_request = _make_guardrail_mock_request({ "messages": [{"role": "user", "content": "hello"}] }) passing_scorer = _SimpleScorer(passing=True) with ( patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=passing_scorer), patch( "mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock(return_value=mock_response), ), ): response = await invocations(endpoint.name, mock_request) assert response.choices[0].message.content == "Safe response" traces = TracingClient().search_traces(locations=[experiment_id]) assert len(traces) == 1 span_map = {s.name: s for s in traces[0].data.spans} assert "guardrail/safety-check" in span_map assert "judge" in span_map gspan = span_map["guardrail/safety-check"] jspan = span_map["judge"] assert gspan.span_type == SpanType.GUARDRAIL assert jspan.span_type == SpanType.EVALUATOR assert jspan.outputs["passed"] is True assert jspan.parent_id == gspan.span_id