diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 07a32530956..9b368978f33 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -374,13 +374,25 @@ def test_get_fallback_model_group_from_fallbacks(model_list): @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio -async def test_deployment_callback_on_success(model_list, sync_mode): +async def test_deployment_callback_on_success(sync_mode): """Test if the '_deployment_callback_on_success' function is working correctly""" import time + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": os.getenv("OPENAI_API_KEY"), + "rpm": 100, + }, + "model_info": {"id": "100"}, + } + ] router = Router(model_list=model_list) standard_logging_payload = create_standard_logging_payload() standard_logging_payload["total_tokens"] = 100 + standard_logging_payload["model_id"] = "100" kwargs = { "litellm_params": { "metadata": { @@ -1181,6 +1193,7 @@ def test_cached_get_model_group_info(model_list): def test_init_responses_api_endpoints(model_list): """Test if the '_init_responses_api_endpoints' function is working correctly""" from typing import Callable + router = Router(model_list=model_list) assert router.aget_responses is not None @@ -1215,29 +1228,29 @@ def test_mock_router_testing_params_str_to_bool_conversion( ): """Test if MockRouterTestingParams.from_kwargs correctly converts string values to booleans using str_to_bool""" from litellm.types.router import MockRouterTestingParams - + kwargs = { "mock_testing_fallbacks": mock_testing_fallbacks, "mock_testing_context_fallbacks": mock_testing_context_fallbacks, "mock_testing_content_policy_fallbacks": mock_testing_content_policy_fallbacks, "other_param": "should_remain", # This should not be affected } - + # Make a copy to verify kwargs are properly popped original_kwargs = kwargs.copy() - + mock_params = MockRouterTestingParams.from_kwargs(kwargs) - + # Verify the converted values assert mock_params.mock_testing_fallbacks == expected_fallbacks assert mock_params.mock_testing_context_fallbacks == expected_context assert mock_params.mock_testing_content_policy_fallbacks == expected_content_policy - + # Verify that the mock testing params were popped from kwargs assert "mock_testing_fallbacks" not in kwargs assert "mock_testing_context_fallbacks" not in kwargs assert "mock_testing_content_policy_fallbacks" not in kwargs - + # Verify other params remain unchanged assert kwargs["other_param"] == "should_remain" @@ -1245,50 +1258,49 @@ def test_mock_router_testing_params_str_to_bool_conversion( def test_is_auto_router_deployment(model_list): """Test if the '_is_auto_router_deployment' function correctly identifies auto-router deployments""" router = Router(model_list=model_list) - + # Test case 1: Model starts with "auto_router/" - should return True litellm_params_auto = LiteLLM_Params(model="auto_router/my-auto-router") assert router._is_auto_router_deployment(litellm_params_auto) is True - + # Test case 2: Model doesn't start with "auto_router/" - should return False litellm_params_regular = LiteLLM_Params(model="gpt-3.5-turbo") assert router._is_auto_router_deployment(litellm_params_regular) is False - + # Test case 3: Model is empty string - should return False litellm_params_empty = LiteLLM_Params(model="") assert router._is_auto_router_deployment(litellm_params_empty) is False - + # Test case 4: Model contains "auto_router/" but doesn't start with it - should return False litellm_params_contains = LiteLLM_Params(model="prefix_auto_router/something") assert router._is_auto_router_deployment(litellm_params_contains) is False - -@patch('litellm.router_strategy.auto_router.auto_router.AutoRouter') +@patch("litellm.router_strategy.auto_router.auto_router.AutoRouter") def test_init_auto_router_deployment_success(mock_auto_router, model_list): """Test if the 'init_auto_router_deployment' function successfully initializes auto-router when all params provided""" router = Router(model_list=model_list) - + # Create a mock AutoRouter instance mock_auto_router_instance = MagicMock() mock_auto_router.return_value = mock_auto_router_instance - + # Test case: All required parameters provided litellm_params = LiteLLM_Params( model="auto_router/test", auto_router_config_path="/path/to/config", auto_router_default_model="gpt-3.5-turbo", - auto_router_embedding_model="text-embedding-ada-002" + auto_router_embedding_model="text-embedding-ada-002", ) deployment = Deployment( - model_name="test-auto-router", + model_name="test-auto-router", litellm_params=litellm_params, - model_info={"id": "test-id"} + model_info={"id": "test-id"}, ) - + # Should not raise any exception router.init_auto_router_deployment(deployment) - + # Verify AutoRouter was called with correct parameters mock_auto_router.assert_called_once_with( model_name="test-auto-router", @@ -1298,86 +1310,96 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list): embedding_model="text-embedding-ada-002", litellm_router_instance=router, ) - + # Verify the auto-router was added to the router's auto_routers dict assert "test-auto-router" in router.auto_routers assert router.auto_routers["test-auto-router"] == mock_auto_router_instance -@patch('litellm.router_strategy.auto_router.auto_router.AutoRouter') +@patch("litellm.router_strategy.auto_router.auto_router.AutoRouter") def test_init_auto_router_deployment_duplicate_model_name(mock_auto_router, model_list): """Test if the 'init_auto_router_deployment' function raises ValueError when model_name already exists""" router = Router(model_list=model_list) - + # Create a mock AutoRouter instance mock_auto_router_instance = MagicMock() mock_auto_router.return_value = mock_auto_router_instance - + # Add an existing auto-router router.auto_routers["test-auto-router"] = mock_auto_router_instance - + # Try to add another auto-router with the same name litellm_params = LiteLLM_Params( model="auto_router/test", auto_router_config_path="/path/to/config", auto_router_default_model="gpt-3.5-turbo", - auto_router_embedding_model="text-embedding-ada-002" + auto_router_embedding_model="text-embedding-ada-002", ) deployment = Deployment( - model_name="test-auto-router", + model_name="test-auto-router", litellm_params=litellm_params, - model_info={"id": "test-id"} + model_info={"id": "test-id"}, ) - - with pytest.raises(ValueError, match="Auto-router deployment test-auto-router already exists"): + + with pytest.raises( + ValueError, match="Auto-router deployment test-auto-router already exists" + ): router.init_auto_router_deployment(deployment) def test_generate_model_id_with_deployment_model_name(model_list): """Test that _generate_model_id works correctly with deployment model_name and handles None values properly""" router = Router(model_list=model_list) - + # Test case 1: Normal case with valid model_group and litellm_params model_group = "gpt-4.1" litellm_params = { "model": "gpt-4.1", "api_key": "test_key", - "api_base": "https://api.openai.com/v1" + "api_base": "https://api.openai.com/v1", } - + try: - result = router._generate_model_id(model_group=model_group, litellm_params=litellm_params) + result = router._generate_model_id( + model_group=model_group, litellm_params=litellm_params + ) assert isinstance(result, str) assert len(result) > 0 print(f"✓ Success with valid model_group: {result}") except Exception as e: pytest.fail(f"Failed with valid model_group: {e}") - + # Test case 2: Edge case with None model_group (this should fail as expected - our fix prevents this from happening) try: - result = router._generate_model_id(model_group=None, litellm_params=litellm_params) - pytest.fail("Expected TypeError when model_group is None - this confirms our fix is needed") + result = router._generate_model_id( + model_group=None, litellm_params=litellm_params + ) + pytest.fail( + "Expected TypeError when model_group is None - this confirms our fix is needed" + ) except TypeError as e: assert "unsupported operand type(s) for +=" in str(e) print(f"✓ Correctly failed with None model_group (as expected): {e}") except Exception as e: pytest.fail(f"Unexpected error with None model_group: {e}") - + # Test case 3: Edge case with None key in litellm_params litellm_params_with_none_key = { "model": "gpt-4.1", "api_key": "test_key", - None: "should_be_skipped" # This should be handled gracefully + None: "should_be_skipped", # This should be handled gracefully } - + try: - result = router._generate_model_id(model_group=model_group, litellm_params=litellm_params_with_none_key) + result = router._generate_model_id( + model_group=model_group, litellm_params=litellm_params_with_none_key + ) assert isinstance(result, str) assert len(result) > 0 print(f"✓ Success with None key in litellm_params: {result}") except Exception as e: pytest.fail(f"Failed with None key in litellm_params: {e}") - + # Test case 4: Edge case with empty litellm_params try: result = router._generate_model_id(model_group=model_group, litellm_params={}) @@ -1386,216 +1408,213 @@ def test_generate_model_id_with_deployment_model_name(model_list): print(f"✓ Success with empty litellm_params: {result}") except Exception as e: pytest.fail(f"Failed with empty litellm_params: {e}") - + # Test case 5: Verify that the same inputs produce the same result (deterministic) - result1 = router._generate_model_id(model_group=model_group, litellm_params=litellm_params) - result2 = router._generate_model_id(model_group=model_group, litellm_params=litellm_params) + result1 = router._generate_model_id( + model_group=model_group, litellm_params=litellm_params + ) + result2 = router._generate_model_id( + model_group=model_group, litellm_params=litellm_params + ) assert result1 == result2, "Model ID generation should be deterministic" - + print("✓ All _generate_model_id tests passed!") def test_handle_clientside_credential_with_deployment_model_name(model_list): """Test that _handle_clientside_credential uses deployment model_name correctly""" router = Router(model_list=model_list) - + # Mock deployment with model_name deployment = { "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key" - } + "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, } - + # Mock kwargs with empty metadata (simulating the original issue) kwargs = { "metadata": {}, # Empty metadata, no model_group "litellm_params": { "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1" - } + "api_base": "https://api.openai.com/v1", + }, } - + # Mock dynamic_litellm_params that would be returned by get_dynamic_litellm_params dynamic_litellm_params = { "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1" + "api_base": "https://api.openai.com/v1", } - + # Test that the method doesn't fail when metadata is empty try: # This would normally call _generate_model_id internally # We're testing that the fix prevents the TypeError model_group = deployment["model_name"] # This is what our fix does assert model_group == "gpt-4.1" - + # Verify that _generate_model_id works with this model_group - result = router._generate_model_id(model_group=model_group, litellm_params=dynamic_litellm_params) + result = router._generate_model_id( + model_group=model_group, litellm_params=dynamic_litellm_params + ) assert isinstance(result, str) assert len(result) > 0 - + print(f"✓ Success with deployment model_name: {result}") except Exception as e: pytest.fail(f"Failed with deployment model_name: {e}") - + print("✓ _handle_clientside_credential test passed!") -@pytest.mark.parametrize("function_name, expected_metadata_key", [ - ("acompletion", "metadata"), - ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), - ("batch", "litellm_metadata"), - ("completion", "metadata"), - ("acreate_file", "litellm_metadata"), - ("aget_file", "litellm_metadata"), -]) -def test_handle_clientside_credential_metadata_loading(model_list, function_name, expected_metadata_key): +@pytest.mark.parametrize( + "function_name, expected_metadata_key", + [ + ("acompletion", "metadata"), + ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), + ("batch", "litellm_metadata"), + ("completion", "metadata"), + ("acreate_file", "litellm_metadata"), + ("aget_file", "litellm_metadata"), + ], +) +def test_handle_clientside_credential_metadata_loading( + model_list, function_name, expected_metadata_key +): """Test that _handle_clientside_credential correctly loads metadata based on function name""" router = Router(model_list=model_list) - + # Mock deployment deployment = { "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key" - }, - "model_info": { - "id": "original-id-123" - } + "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, + "model_info": {"id": "original-id-123"}, } - + # Mock kwargs with clientside credentials and metadata kwargs = { "api_key": "client_side_key", "api_base": "https://api.openai.com/v1", - expected_metadata_key: { - "model_group": "gpt-4.1", - "custom_field": "test_value" - } + expected_metadata_key: {"model_group": "gpt-4.1", "custom_field": "test_value"}, } - + # Call the function result_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs=kwargs, - function_name=function_name + deployment=deployment, kwargs=kwargs, function_name=function_name ) - + # Verify the result is a Deployment object assert isinstance(result_deployment, Deployment) - + # Verify the deployment has the correct model_name (should be the model_group from metadata) assert result_deployment.model_name == "gpt-4.1" - + # Verify the litellm_params contain the clientside credentials assert result_deployment.litellm_params.api_key == "client_side_key" assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - + # Verify the model_info has been updated with a new ID assert result_deployment.model_info.id != "original-id-123" assert result_deployment.model_info.original_model_id == "original-id-123" - + # Verify the deployment was added to the router assert len(router.model_list) == len(model_list) + 1 - + # Test that the function correctly uses the right metadata key # For acompletion, it should use "metadata" # For _ageneric_api_call_with_fallbacks/batch, it should use "litellm_metadata" if function_name == "acompletion": assert "metadata" in kwargs assert "litellm_metadata" not in kwargs - elif function_name in ["_ageneric_api_call_with_fallbacks", "batch", "acreate_file", "aget_file"]: + elif function_name in [ + "_ageneric_api_call_with_fallbacks", + "batch", + "acreate_file", + "aget_file", + ]: assert "litellm_metadata" in kwargs # Note: acompletion would not have litellm_metadata, but other functions might have both - - print(f"✓ Success with function_name '{function_name}' using '{expected_metadata_key}' metadata key") + + print( + f"✓ Success with function_name '{function_name}' using '{expected_metadata_key}' metadata key" + ) -@pytest.mark.parametrize("function_name, metadata_key", [ - ("acompletion", "metadata"), - ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), -]) -def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key): +@pytest.mark.parametrize( + "function_name, metadata_key", + [ + ("acompletion", "metadata"), + ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), + ], +) +def test_handle_clientside_credential_metadata_variable_name( + model_list, function_name, metadata_key +): """Test that _handle_clientside_credential uses the correct metadata variable name based on function name""" from litellm.router_utils.batch_utils import _get_router_metadata_variable_name - + router = Router(model_list=model_list) - + # Verify the metadata variable name is correct for each function - expected_metadata_key = _get_router_metadata_variable_name(function_name=function_name) + expected_metadata_key = _get_router_metadata_variable_name( + function_name=function_name + ) assert expected_metadata_key == metadata_key - + # Mock deployment deployment = { "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key" - }, - "model_info": { - "id": "original-id-456" - } + "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, + "model_info": {"id": "original-id-456"}, } - + # Mock kwargs with clientside credentials and the correct metadata key kwargs = { "api_key": "client_side_key", "api_base": "https://api.openai.com/v1", - metadata_key: { - "model_group": "gpt-4.1", - "test_field": "test_value" - } + metadata_key: {"model_group": "gpt-4.1", "test_field": "test_value"}, } - + # Call the function result_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs=kwargs, - function_name=function_name + deployment=deployment, kwargs=kwargs, function_name=function_name ) - + # Verify the function correctly extracted model_group from the right metadata key assert result_deployment.model_name == "gpt-4.1" - + # Verify the deployment was created with the correct metadata assert result_deployment.litellm_params.api_key == "client_side_key" assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - - print(f"✓ Success with function_name '{function_name}' correctly using '{metadata_key}' for metadata") + + print( + f"✓ Success with function_name '{function_name}' correctly using '{metadata_key}' for metadata" + ) def test_handle_clientside_credential_no_metadata(model_list): """Test that _handle_clientside_credential handles cases where no metadata is provided""" router = Router(model_list=model_list) - + # Mock deployment deployment = { "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key" - }, - "model_info": { - "id": "original-id-789" - } + "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, + "model_info": {"id": "original-id-789"}, } - + # Mock kwargs with clientside credentials but NO metadata kwargs = { "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1" + "api_base": "https://api.openai.com/v1", # No metadata key at all } - + # This should fail because there's no model_group in metadata # The function expects to find model_group in the metadata try: result_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs=kwargs, - function_name="acompletion" + deployment=deployment, kwargs=kwargs, function_name="acompletion" ) # If we get here, the function should have used deployment.model_name as fallback assert result_deployment.model_name == "gpt-4.1" @@ -1603,19 +1622,19 @@ def test_handle_clientside_credential_no_metadata(model_list): except Exception as e: # This is expected behavior - the function needs model_group to generate model_id print(f"✓ Correctly handled no metadata case: {e}") - + # Test with empty metadata kwargs_with_empty_metadata = { "api_key": "client_side_key", "api_base": "https://api.openai.com/v1", - "metadata": {} # Empty metadata + "metadata": {}, # Empty metadata } - + try: result_deployment = router._handle_clientside_credential( deployment=deployment, kwargs=kwargs_with_empty_metadata, - function_name="acompletion" + function_name="acompletion", ) # Should fail because empty metadata has no model_group pytest.fail("Expected failure with empty metadata") @@ -1626,36 +1645,31 @@ def test_handle_clientside_credential_no_metadata(model_list): def test_handle_clientside_credential_with_responses_function(model_list): """Test that _handle_clientside_credential works correctly with responses function name""" router = Router(model_list=model_list) - + # Mock deployment deployment = { "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key" - }, - "model_info": { - "id": "original-id-responses" - } + "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, + "model_info": {"id": "original-id-responses"}, } - + # Mock kwargs with clientside credentials and litellm_metadata (for responses function) kwargs = { "api_key": "client_side_key", "api_base": "https://api.openai.com/v1", "litellm_metadata": { "model_group": "gpt-4.1", - "responses_field": "responses_value" - } + "responses_field": "responses_value", + }, } - + # Call the function with _ageneric_api_call_with_fallbacks function name (which handles responses) result_deployment = router._handle_clientside_credential( deployment=deployment, kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks" + function_name="_ageneric_api_call_with_fallbacks", ) - + # Verify the result assert isinstance(result_deployment, Deployment) assert result_deployment.model_name == "gpt-4.1" @@ -1663,9 +1677,10 @@ def test_handle_clientside_credential_with_responses_function(model_list): assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" assert result_deployment.model_info.id != "original-id-responses" assert result_deployment.model_info.original_model_id == "original-id-responses" - + # Verify the deployment was added to the router assert len(router.model_list) == len(model_list) + 1 - - print("✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata") + print( + "✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata" + )