diff --git a/litellm/router.py b/litellm/router.py index 50f6d7f1b6a..5972b06f01a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -59,6 +59,7 @@ from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler @@ -141,10 +142,18 @@ from litellm.types.router import ( RoutingStrategy, ) from litellm.types.services import ServiceTypes -from litellm.types.utils import GenericBudgetConfigType, LiteLLMBatch +from litellm.types.utils import ( + CustomPricingLiteLLMParams, + GenericBudgetConfigType, + LiteLLMBatch, +) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage, CustomPricingLiteLLMParams +from litellm.types.utils import ( + ModelResponseStream, + StandardLoggingPayload, + Usage, +) from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -908,8 +917,13 @@ class Router: try: _deployment_copy = copy.deepcopy(deployment) litellm_params: dict = _deployment_copy["litellm_params"] - if "api_key" in litellm_params: + + if litellm.redact_user_api_key_info: + masker = SensitiveDataMasker(visible_prefix=2, visible_suffix=0) + _deployment_copy["litellm_params"] = masker.mask_dict(litellm_params) + elif "api_key" in litellm_params: litellm_params["api_key"] = litellm_params["api_key"][:2] + "*" * 10 + return _deployment_copy except Exception as e: verbose_router_logger.debug( diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 094df944bcc..a31c4d8210f 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -94,6 +94,33 @@ def test_print_deployment(model_list): assert 10 * "*" in printed_deployment["litellm_params"]["api_key"] +def test_print_deployment_with_redact_enabled(model_list): + """Test if sensitive credentials are masked when redact_user_api_key_info is enabled""" + import litellm + + router = Router(model_list=model_list) + deployment = { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-v2", + "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "aws_region_name": "us-west-2", + }, + } + + original_setting = litellm.redact_user_api_key_info + try: + litellm.redact_user_api_key_info = True + printed_deployment = router.print_deployment(deployment) + + assert "*" in printed_deployment["litellm_params"]["aws_access_key_id"] + assert "*" in printed_deployment["litellm_params"]["aws_secret_access_key"] + assert "us-west-2" == printed_deployment["litellm_params"]["aws_region_name"] + finally: + litellm.redact_user_api_key_info = original_setting + + def test_completion(model_list): """Test if the completion function is working correctly""" router = Router(model_list=model_list)