From 5ad698f2cc7b763b3b0a394a4cf796057a0079e6 Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 14:21:07 -0700 Subject: [PATCH 1/9] Revert "Fix - add safe divide by 0 for most places to prevent crash" This reverts commit 265d40e39051e148996b9fb7f354730c57ff23ac. --- litellm/litellm_core_utils/core_helpers.py | 21 ------ litellm/router_strategy/simple_shuffle.py | 7 +- .../litellm_core_utils/test_core_helpers.py | 72 +------------------ 3 files changed, 4 insertions(+), 96 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 4aeb9d4d640..13a2e554f12 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -37,27 +37,6 @@ def safe_divide_seconds( return float(seconds / denominator) -def safe_divide( - numerator: Union[int, float], - denominator: Union[int, float], - default: Union[int, float] = 0 -) -> Union[int, float]: - """ - Safely divide two numbers, returning a default value if denominator is zero. - - Args: - numerator: The number to divide - denominator: The number to divide by - default: Value to return if denominator is zero (defaults to 0) - - Returns: - The result of numerator/denominator, or default if denominator is zero - """ - if denominator == 0: - return default - return numerator / denominator - - def map_finish_reason( finish_reason: str, ): # openai supports 5 stop sequences - 'stop', 'length', 'function_call', 'content_filter', 'null' diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index bef145c44ba..da24c02f2e3 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -9,7 +9,6 @@ import random from typing import TYPE_CHECKING, Any, Dict, List, Union from litellm._logging import verbose_router_logger -from litellm.litellm_core_utils.core_helpers import safe_divide if TYPE_CHECKING: from litellm.router import Router as _Router @@ -47,7 +46,7 @@ def simple_shuffle( weights = [m["litellm_params"].get("weight", 0) for m in healthy_deployments] verbose_router_logger.debug(f"\nweight {weights}") total_weight = sum(weights) - weights = [safe_divide(weight, total_weight, 0) for weight in weights] + weights = [weight / total_weight for weight in weights] verbose_router_logger.debug(f"\n weights {weights}") # Perform weighted random pick selected_index = random.choices(range(len(weights)), weights=weights)[0] @@ -64,7 +63,7 @@ def simple_shuffle( rpms = [m["litellm_params"].get("rpm", 0) for m in healthy_deployments] verbose_router_logger.debug(f"\nrpms {rpms}") total_rpm = sum(rpms) - weights = [safe_divide(rpm, total_rpm, 0) for rpm in rpms] + weights = [rpm / total_rpm for rpm in rpms] verbose_router_logger.debug(f"\n weights {weights}") # Perform weighted random pick selected_index = random.choices(range(len(rpms)), weights=weights)[0] @@ -81,7 +80,7 @@ def simple_shuffle( tpms = [m["litellm_params"].get("tpm", 0) for m in healthy_deployments] verbose_router_logger.debug(f"\ntpms {tpms}") total_tpm = sum(tpms) - weights = [safe_divide(tpm, total_tpm, 0) for tpm in tpms] + weights = [tpm / total_tpm for tpm in tpms] verbose_router_logger.debug(f"\n weights {weights}") # Perform weighted random pick selected_index = random.choices(range(len(tpms)), weights=weights)[0] diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 32f3ad3f55c..d7869e6b800 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -9,7 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, safe_divide +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs def test_get_litellm_metadata_from_kwargs(): @@ -57,73 +57,3 @@ def test_preserve_upstream_non_openai_attributes(): ) assert model_response.test_key == "test_value" - - -def test_safe_divide_basic(): - """Test basic safe division functionality""" - # Normal division - result = safe_divide(10, 2) - assert result == 5.0, f"Expected 5.0, got {result}" - - # Division with float - result = safe_divide(7.5, 2.5) - assert result == 3.0, f"Expected 3.0, got {result}" - - # Division by zero with default - result = safe_divide(10, 0) - assert result == 0, f"Expected 0, got {result}" - - # Division by zero with custom default - result = safe_divide(10, 0, default=1) - assert result == 1, f"Expected 1, got {result}" - - # Division by zero with custom default as float - result = safe_divide(10, 0, default=0.5) - assert result == 0.5, f"Expected 0.5, got {result}" - - -def test_safe_divide_edge_cases(): - """Test edge cases for safe division""" - # Zero numerator - result = safe_divide(0, 5) - assert result == 0.0, f"Expected 0.0, got {result}" - - # Negative numbers - result = safe_divide(-10, 2) - assert result == -5.0, f"Expected -5.0, got {result}" - - # Negative denominator - result = safe_divide(10, -2) - assert result == -5.0, f"Expected -5.0, got {result}" - - # Both negative - result = safe_divide(-10, -2) - assert result == 5.0, f"Expected 5.0, got {result}" - - # Float division - result = safe_divide(1, 3) - assert abs(result - 0.3333333333333333) < 1e-10, f"Expected ~0.333..., got {result}" - - -def test_safe_divide_weight_scenario(): - """Test safe division in the context of weight calculations""" - # Simulate weight calculation scenario - weights = [3, 7, 0, 2] - total_weight = sum(weights) # 12 - - # Normal case - normalized_weights = [safe_divide(w, total_weight) for w in weights] - expected = [0.25, 7/12, 0.0, 1/6] - - for i, (actual, exp) in enumerate(zip(normalized_weights, expected)): - assert abs(actual - exp) < 1e-10, f"Weight {i}: Expected {exp}, got {actual}" - - # Zero total weight scenario (division by zero) - zero_weights = [0, 0, 0] - zero_total = sum(zero_weights) # 0 - - # Should return default values (0) for all weights - normalized_zero_weights = [safe_divide(w, zero_total) for w in zero_weights] - expected_zero = [0, 0, 0] - - assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}" From bfb0a3854ec1604bf6156b66489adccc32936d52 Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 14:21:22 -0700 Subject: [PATCH 2/9] Enhance logging in cost calculation tests to ensure DEBUG level captures are accurate. Updated tests to set logger level before assertions and restored original logger level after execution. This improves reliability of log level checks in both cost and batch cost calculation tests. --- .../test_cost_calculation_log_level.py | 144 ++++++++++-------- 1 file changed, 81 insertions(+), 63 deletions(-) diff --git a/tests/test_litellm/test_cost_calculation_log_level.py b/tests/test_litellm/test_cost_calculation_log_level.py index 4380ae8bf62..3925ea751af 100644 --- a/tests/test_litellm/test_cost_calculation_log_level.py +++ b/tests/test_litellm/test_cost_calculation_log_level.py @@ -17,46 +17,55 @@ def test_cost_calculation_uses_debug_level(caplog): This ensures cost calculation details don't appear in production logs. Part of fix for issue #9815. """ - # Create a mock completion response - mock_response = { - "id": "test", - "object": "chat.completion", - "created": 1234567890, - "model": "gpt-3.5-turbo", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop" - }], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "total_tokens": 30 + # Ensure verbose_logger is set to DEBUG level to capture the debug logs + from litellm._logging import verbose_logger + original_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) + + try: + # Create a mock completion response + mock_response = { + "id": "test", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-3.5-turbo", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Test response"}, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + } } - } - - # Test that cost calculation logs are at DEBUG level - with caplog.at_level(logging.DEBUG): - try: - cost = completion_cost( - completion_response=mock_response, - model="gpt-3.5-turbo" - ) - except Exception: - pass # Cost calculation may fail, but we're checking log levels - - # Find the cost calculation log records - cost_calc_records = [ - record for record in caplog.records - if "selected model name for cost calculation" in record.message - ] - - # Verify that cost calculation logs are at DEBUG level - assert len(cost_calc_records) > 0, "No cost calculation logs found" - - for record in cost_calc_records: - assert record.levelno == logging.DEBUG, \ - f"Cost calculation log should be DEBUG level, but was {record.levelname}" + + # Test that cost calculation logs are at DEBUG level + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + try: + cost = completion_cost( + completion_response=mock_response, + model="gpt-3.5-turbo" + ) + except Exception: + pass # Cost calculation may fail, but we're checking log levels + + # Find the cost calculation log records + cost_calc_records = [ + record for record in caplog.records + if "selected model name for cost calculation" in record.message + ] + + # Verify that cost calculation logs are at DEBUG level + assert len(cost_calc_records) > 0, "No cost calculation logs found" + + for record in cost_calc_records: + assert record.levelno == logging.DEBUG, \ + f"Cost calculation log should be DEBUG level, but was {record.levelname}" + finally: + # Restore original logger level + verbose_logger.setLevel(original_level) def test_batch_cost_calculation_uses_debug_level(caplog): @@ -65,29 +74,38 @@ def test_batch_cost_calculation_uses_debug_level(caplog): """ from litellm.cost_calculator import batch_cost_calculator from litellm.types.utils import Usage + from litellm._logging import verbose_logger - # Create a mock usage object - usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + # Ensure verbose_logger is set to DEBUG level to capture the debug logs + original_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) - # Test that batch cost calculation logs are at DEBUG level - with caplog.at_level(logging.DEBUG): - try: - batch_cost_calculator( - usage=usage, - model="gpt-3.5-turbo", - custom_llm_provider="openai" - ) - except Exception: - pass # May fail, but we're checking log levels - - # Find batch cost calculation log records - batch_cost_records = [ - record for record in caplog.records - if "Calculating batch cost per token" in record.message - ] - - # Verify logs exist and are at DEBUG level - if batch_cost_records: # May not always log depending on the code path - for record in batch_cost_records: - assert record.levelno == logging.DEBUG, \ - f"Batch cost calculation log should be DEBUG level, but was {record.levelname}" \ No newline at end of file + try: + # Create a mock usage object + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + + # Test that batch cost calculation logs are at DEBUG level + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + try: + batch_cost_calculator( + usage=usage, + model="gpt-3.5-turbo", + custom_llm_provider="openai" + ) + except Exception: + pass # May fail, but we're checking log levels + + # Find batch cost calculation log records + batch_cost_records = [ + record for record in caplog.records + if "Calculating batch cost per token" in record.message + ] + + # Verify logs exist and are at DEBUG level + if batch_cost_records: # May not always log depending on the code path + for record in batch_cost_records: + assert record.levelno == logging.DEBUG, \ + f"Batch cost calculation log should be DEBUG level, but was {record.levelname}" + finally: + # Restore original logger level + verbose_logger.setLevel(original_level) \ No newline at end of file From a6e55c0447b1a5ae455a64eac627318a7b6cd19b Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 14:23:06 -0700 Subject: [PATCH 3/9] Revert "Revert "Fix - add safe divide by 0 for most places to prevent crash"" This reverts commit 5ad698f2cc7b763b3b0a394a4cf796057a0079e6. --- litellm/litellm_core_utils/core_helpers.py | 21 ++++++ litellm/router_strategy/simple_shuffle.py | 7 +- .../litellm_core_utils/test_core_helpers.py | 72 ++++++++++++++++++- 3 files changed, 96 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 13a2e554f12..4aeb9d4d640 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -37,6 +37,27 @@ def safe_divide_seconds( return float(seconds / denominator) +def safe_divide( + numerator: Union[int, float], + denominator: Union[int, float], + default: Union[int, float] = 0 +) -> Union[int, float]: + """ + Safely divide two numbers, returning a default value if denominator is zero. + + Args: + numerator: The number to divide + denominator: The number to divide by + default: Value to return if denominator is zero (defaults to 0) + + Returns: + The result of numerator/denominator, or default if denominator is zero + """ + if denominator == 0: + return default + return numerator / denominator + + def map_finish_reason( finish_reason: str, ): # openai supports 5 stop sequences - 'stop', 'length', 'function_call', 'content_filter', 'null' diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index da24c02f2e3..bef145c44ba 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -9,6 +9,7 @@ import random from typing import TYPE_CHECKING, Any, Dict, List, Union from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.core_helpers import safe_divide if TYPE_CHECKING: from litellm.router import Router as _Router @@ -46,7 +47,7 @@ def simple_shuffle( weights = [m["litellm_params"].get("weight", 0) for m in healthy_deployments] verbose_router_logger.debug(f"\nweight {weights}") total_weight = sum(weights) - weights = [weight / total_weight for weight in weights] + weights = [safe_divide(weight, total_weight, 0) for weight in weights] verbose_router_logger.debug(f"\n weights {weights}") # Perform weighted random pick selected_index = random.choices(range(len(weights)), weights=weights)[0] @@ -63,7 +64,7 @@ def simple_shuffle( rpms = [m["litellm_params"].get("rpm", 0) for m in healthy_deployments] verbose_router_logger.debug(f"\nrpms {rpms}") total_rpm = sum(rpms) - weights = [rpm / total_rpm for rpm in rpms] + weights = [safe_divide(rpm, total_rpm, 0) for rpm in rpms] verbose_router_logger.debug(f"\n weights {weights}") # Perform weighted random pick selected_index = random.choices(range(len(rpms)), weights=weights)[0] @@ -80,7 +81,7 @@ def simple_shuffle( tpms = [m["litellm_params"].get("tpm", 0) for m in healthy_deployments] verbose_router_logger.debug(f"\ntpms {tpms}") total_tpm = sum(tpms) - weights = [tpm / total_tpm for tpm in tpms] + weights = [safe_divide(tpm, total_tpm, 0) for tpm in tpms] verbose_router_logger.debug(f"\n weights {weights}") # Perform weighted random pick selected_index = random.choices(range(len(tpms)), weights=weights)[0] diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index d7869e6b800..32f3ad3f55c 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -9,7 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, safe_divide def test_get_litellm_metadata_from_kwargs(): @@ -57,3 +57,73 @@ def test_preserve_upstream_non_openai_attributes(): ) assert model_response.test_key == "test_value" + + +def test_safe_divide_basic(): + """Test basic safe division functionality""" + # Normal division + result = safe_divide(10, 2) + assert result == 5.0, f"Expected 5.0, got {result}" + + # Division with float + result = safe_divide(7.5, 2.5) + assert result == 3.0, f"Expected 3.0, got {result}" + + # Division by zero with default + result = safe_divide(10, 0) + assert result == 0, f"Expected 0, got {result}" + + # Division by zero with custom default + result = safe_divide(10, 0, default=1) + assert result == 1, f"Expected 1, got {result}" + + # Division by zero with custom default as float + result = safe_divide(10, 0, default=0.5) + assert result == 0.5, f"Expected 0.5, got {result}" + + +def test_safe_divide_edge_cases(): + """Test edge cases for safe division""" + # Zero numerator + result = safe_divide(0, 5) + assert result == 0.0, f"Expected 0.0, got {result}" + + # Negative numbers + result = safe_divide(-10, 2) + assert result == -5.0, f"Expected -5.0, got {result}" + + # Negative denominator + result = safe_divide(10, -2) + assert result == -5.0, f"Expected -5.0, got {result}" + + # Both negative + result = safe_divide(-10, -2) + assert result == 5.0, f"Expected 5.0, got {result}" + + # Float division + result = safe_divide(1, 3) + assert abs(result - 0.3333333333333333) < 1e-10, f"Expected ~0.333..., got {result}" + + +def test_safe_divide_weight_scenario(): + """Test safe division in the context of weight calculations""" + # Simulate weight calculation scenario + weights = [3, 7, 0, 2] + total_weight = sum(weights) # 12 + + # Normal case + normalized_weights = [safe_divide(w, total_weight) for w in weights] + expected = [0.25, 7/12, 0.0, 1/6] + + for i, (actual, exp) in enumerate(zip(normalized_weights, expected)): + assert abs(actual - exp) < 1e-10, f"Weight {i}: Expected {exp}, got {actual}" + + # Zero total weight scenario (division by zero) + zero_weights = [0, 0, 0] + zero_total = sum(zero_weights) # 0 + + # Should return default values (0) for all weights + normalized_zero_weights = [safe_divide(w, zero_total) for w in zero_weights] + expected_zero = [0, 0, 0] + + assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}" From 5fc0803b945b48484926d8d18f67d7f370ce6f97 Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 14:40:49 -0700 Subject: [PATCH 4/9] Add mock user API key authentication in tag management tests This update introduces a helper function to create a mock user API key authentication object, which is utilized in the tag management endpoint tests. The mock authentication is integrated into the test cases for creating, updating, and deleting tags, enhancing the reliability of the tests by simulating user roles accurately. --- .../test_tag_management_endpoints.py | 33 +++++++++++++++++-- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index add08f55683..bd68618e56c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -14,16 +14,28 @@ from unittest.mock import patch import litellm from litellm.proxy.proxy_server import app +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest client = TestClient(app) +def create_mock_user_api_key_auth(): + """Helper function to create a mock auth object""" + return UserAPIKeyAuth( + user_id="test-user", + user_role=LitellmUserRoles.PROXY_ADMIN + ) + + @pytest.mark.asyncio async def test_create_and_get_tag(): """ Test creation of a new tag and retrieving its information """ + # Create a mock auth object + mock_user_api_key_auth = create_mock_user_api_key_auth() + # Mock the prisma client and _get_tags_config and _save_tags_config with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.llm_router" @@ -35,8 +47,11 @@ async def test_create_and_get_tag(): "litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment" ) as mock_add_tag, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" - ) as mock_get_models: + ) as mock_get_models, patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" + ) as mock_auth: # Setup mocks + mock_auth.return_value = mock_user_api_key_auth mock_get_tags.return_value = {} mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"} @@ -83,6 +98,9 @@ async def test_update_tag(): """ Test updating an existing tag """ + # Create a mock auth object + mock_user_api_key_auth = create_mock_user_api_key_auth() + # Mock the prisma client and _get_tags_config and _save_tags_config with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" @@ -90,8 +108,11 @@ async def test_update_tag(): "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" ) as mock_save_tags, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" - ) as mock_get_models: + ) as mock_get_models, patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" + ) as mock_auth: # Setup mocks for existing tag + mock_auth.return_value = mock_user_api_key_auth mock_get_tags.return_value = { "test-tag": { "name": "test-tag", @@ -129,13 +150,19 @@ async def test_delete_tag(): """ Test deleting a tag """ + # Create a mock auth object + mock_user_api_key_auth = create_mock_user_api_key_auth() + # Mock the prisma client and _get_tags_config and _save_tags_config with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" ) as mock_get_tags, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" - ) as mock_save_tags: + ) as mock_save_tags, patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" + ) as mock_auth: # Setup mocks for existing tag + mock_auth.return_value = mock_user_api_key_auth mock_get_tags.return_value = { "test-tag": { "name": "test-tag", From d21f467264d09c42fb756766dd7c6b0438b2c9f2 Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 14:58:49 -0700 Subject: [PATCH 5/9] Revert "Add mock user API key authentication in tag management tests" This reverts commit 5fc0803b945b48484926d8d18f67d7f370ce6f97. --- .../test_tag_management_endpoints.py | 33 ++----------------- 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index bd68618e56c..add08f55683 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -14,28 +14,16 @@ from unittest.mock import patch import litellm from litellm.proxy.proxy_server import app -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest client = TestClient(app) -def create_mock_user_api_key_auth(): - """Helper function to create a mock auth object""" - return UserAPIKeyAuth( - user_id="test-user", - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - @pytest.mark.asyncio async def test_create_and_get_tag(): """ Test creation of a new tag and retrieving its information """ - # Create a mock auth object - mock_user_api_key_auth = create_mock_user_api_key_auth() - # Mock the prisma client and _get_tags_config and _save_tags_config with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.llm_router" @@ -47,11 +35,8 @@ async def test_create_and_get_tag(): "litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment" ) as mock_add_tag, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" - ) as mock_get_models, patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" - ) as mock_auth: + ) as mock_get_models: # Setup mocks - mock_auth.return_value = mock_user_api_key_auth mock_get_tags.return_value = {} mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"} @@ -98,9 +83,6 @@ async def test_update_tag(): """ Test updating an existing tag """ - # Create a mock auth object - mock_user_api_key_auth = create_mock_user_api_key_auth() - # Mock the prisma client and _get_tags_config and _save_tags_config with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" @@ -108,11 +90,8 @@ async def test_update_tag(): "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" ) as mock_save_tags, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" - ) as mock_get_models, patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" - ) as mock_auth: + ) as mock_get_models: # Setup mocks for existing tag - mock_auth.return_value = mock_user_api_key_auth mock_get_tags.return_value = { "test-tag": { "name": "test-tag", @@ -150,19 +129,13 @@ async def test_delete_tag(): """ Test deleting a tag """ - # Create a mock auth object - mock_user_api_key_auth = create_mock_user_api_key_auth() - # Mock the prisma client and _get_tags_config and _save_tags_config with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" ) as mock_get_tags, patch( "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" - ) as mock_save_tags, patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" - ) as mock_auth: + ) as mock_save_tags: # Setup mocks for existing tag - mock_auth.return_value = mock_user_api_key_auth mock_get_tags.return_value = { "test-tag": { "name": "test-tag", From 45f188b04106a05fbdf2f441a1ddff9f9c17f0bf Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 15:03:21 -0700 Subject: [PATCH 6/9] Add mock user API key authentication in tag management tests This update integrates mock user API key authentication into the tag management endpoint tests, ensuring accurate simulation of user roles for creating, updating, and deleting tags. The changes enhance the reliability of the tests by properly setting up user authentication before executing test cases. --- .../test_tag_management_endpoints.py | 228 ++++++++++-------- 1 file changed, 134 insertions(+), 94 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index add08f55683..749ee4acd16 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -14,6 +14,7 @@ from unittest.mock import patch import litellm from litellm.proxy.proxy_server import app +from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest client = TestClient(app) @@ -24,58 +25,71 @@ async def test_create_and_get_tag(): """ Test creation of a new tag and retrieving its information """ - # Mock the prisma client and _get_tags_config and _save_tags_config - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.llm_router" - ) as mock_router, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" - ) as mock_get_tags, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" - ) as mock_save_tags, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment" - ) as mock_add_tag, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" - ) as mock_get_models: - # Setup mocks - mock_get_tags.return_value = {} - mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"} + # Mock the user authentication + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + # Mock the prisma client and _get_tags_config and _save_tags_config + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.llm_router" + ) as mock_router, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" + ) as mock_get_tags, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" + ) as mock_save_tags, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment" + ) as mock_add_tag, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" + ) as mock_get_models: + # Setup mocks + mock_get_tags.return_value = {} + mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"} - # Create a new tag - tag_data = { - "name": "test-tag", - "description": "Test tag for unit testing", - "models": ["model-1"], - } - - # Set admin access for the test - headers = {"Authorization": f"Bearer sk-1234"} - - # Test tag creation - response = client.post("/tag/new", json=tag_data, headers=headers) - print(f"response: {response.text}") - assert response.status_code == 200 - result = response.json() - assert result["message"] == "Tag test-tag created successfully" - assert result["tag"]["name"] == "test-tag" - assert result["tag"]["description"] == "Test tag for unit testing" - - # Mock updated tag config for the get request - mock_get_tags.return_value = { - "test-tag": { + # Create a new tag + tag_data = { "name": "test-tag", "description": "Test tag for unit testing", "models": ["model-1"], - "model_info": {"model-1": "gpt-3.5-turbo"}, } - } - # Test retrieving tag info - info_data = {"names": ["test-tag"]} - response = client.post("/tag/info", json=info_data, headers=headers) - assert response.status_code == 200 - result = response.json() - assert "test-tag" in result - assert result["test-tag"]["description"] == "Test tag for unit testing" + # Set admin access for the test + headers = {"Authorization": f"Bearer sk-1234"} + + # Test tag creation + response = client.post("/tag/new", json=tag_data, headers=headers) + print(f"response: {response.text}") + assert response.status_code == 200 + result = response.json() + assert result["message"] == "Tag test-tag created successfully" + assert result["tag"]["name"] == "test-tag" + assert result["tag"]["description"] == "Test tag for unit testing" + + # Mock updated tag config for the get request + mock_get_tags.return_value = { + "test-tag": { + "name": "test-tag", + "description": "Test tag for unit testing", + "models": ["model-1"], + "model_info": {"model-1": "gpt-3.5-turbo"}, + } + } + + # Test retrieving tag info + info_data = {"names": ["test-tag"]} + response = client.post("/tag/info", json=info_data, headers=headers) + assert response.status_code == 200 + result = response.json() + assert "test-tag" in result + assert result["test-tag"]["description"] == "Test tag for unit testing" + finally: + # Clean up dependency overrides + app.dependency_overrides.clear() @pytest.mark.asyncio @@ -83,16 +97,26 @@ async def test_update_tag(): """ Test updating an existing tag """ - # Mock the prisma client and _get_tags_config and _save_tags_config - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" - ) as mock_get_tags, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" - ) as mock_save_tags, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" - ) as mock_get_models: - # Setup mocks for existing tag - mock_get_tags.return_value = { + # Mock the user authentication + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + # Mock the prisma client and _get_tags_config and _save_tags_config + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" + ) as mock_get_tags, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" + ) as mock_save_tags, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names" + ) as mock_get_models: + # Setup mocks for existing tag + mock_get_tags.return_value = { "test-tag": { "name": "test-tag", "description": "Original description", @@ -101,27 +125,30 @@ async def test_update_tag(): "updated_at": "2023-01-01T00:00:00", "created_by": "user-123", } - } - mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"} + } + mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"} - # Update tag data - update_data = { - "name": "test-tag", - "description": "Updated description", - "models": ["model-1", "model-2"], - } + # Update tag data + update_data = { + "name": "test-tag", + "description": "Updated description", + "models": ["model-1", "model-2"], + } - # Set admin access for the test - headers = {"Authorization": f"Bearer sk-1234"} + # Set admin access for the test + headers = {"Authorization": f"Bearer sk-1234"} - # Test tag update - response = client.post("/tag/update", json=update_data, headers=headers) - assert response.status_code == 200 - result = response.json() - assert result["message"] == "Tag test-tag updated successfully" - assert result["tag"]["description"] == "Updated description" - assert len(result["tag"]["models"]) == 2 - assert "model-2" in result["tag"]["models"] + # Test tag update + response = client.post("/tag/update", json=update_data, headers=headers) + assert response.status_code == 200 + result = response.json() + assert result["message"] == "Tag test-tag updated successfully" + assert result["tag"]["description"] == "Updated description" + assert len(result["tag"]["models"]) == 2 + assert "model-2" in result["tag"]["models"] + finally: + # Clean up dependency overrides + app.dependency_overrides.clear() @pytest.mark.asyncio @@ -129,14 +156,24 @@ async def test_delete_tag(): """ Test deleting a tag """ - # Mock the prisma client and _get_tags_config and _save_tags_config - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" - ) as mock_get_tags, patch( - "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" - ) as mock_save_tags: - # Setup mocks for existing tag - mock_get_tags.return_value = { + # Mock the user authentication + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + # Mock the prisma client and _get_tags_config and _save_tags_config + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config" + ) as mock_get_tags, patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config" + ) as mock_save_tags: + # Setup mocks for existing tag + mock_get_tags.return_value = { "test-tag": { "name": "test-tag", "description": "Test tag for deletion", @@ -145,22 +182,25 @@ async def test_delete_tag(): "updated_at": "2023-01-01T00:00:00", "created_by": "user-123", } - } + } - # Delete tag data - delete_data = {"name": "test-tag"} + # Delete tag data + delete_data = {"name": "test-tag"} - # Set admin access for the test - headers = {"Authorization": f"Bearer sk-1234"} + # Set admin access for the test + headers = {"Authorization": f"Bearer sk-1234"} - # Test tag deletion - response = client.post("/tag/delete", json=delete_data, headers=headers) - assert response.status_code == 200 - result = response.json() - assert result["message"] == "Tag test-tag deleted successfully" + # Test tag deletion + response = client.post("/tag/delete", json=delete_data, headers=headers) + assert response.status_code == 200 + result = response.json() + assert result["message"] == "Tag test-tag deleted successfully" - # Verify _save_tags_config was called without the deleted tag - mock_save_tags.assert_called_once() + # Verify _save_tags_config was called without the deleted tag + mock_save_tags.assert_called_once() + finally: + # Clean up dependency overrides + app.dependency_overrides.clear() @pytest.mark.asyncio From f6e53deacd843b726beaab2ed9b6951e1b261c4d Mon Sep 17 00:00:00 2001 From: TomuHirata Date: Fri, 15 Aug 2025 07:20:50 +0900 Subject: [PATCH 7/9] Update mlflow logger usage span attributes (#13561) * test: sync mlflow request tags * fix test --- litellm/integrations/mlflow.py | 6 +-- .../test_litellm/integrations/test_mlflow.py | 41 ++++++++++++++++++- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index ea9051db4de..634d0c1fdc9 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -189,9 +189,9 @@ class MlflowLogger(CustomLogger): { "api_base": standard_obj.get("api_base"), "cache_hit": standard_obj.get("cache_hit"), - "usage": { - "completion_tokens": standard_obj.get("completion_tokens"), - "prompt_tokens": standard_obj.get("prompt_tokens"), + "mlflow.chat.tokenUsage": { + "input_tokens": standard_obj.get("prompt_tokens"), + "output_tokens": standard_obj.get("completion_tokens"), "total_tokens": standard_obj.get("total_tokens"), }, "raw_llm_response": standard_obj.get("response"), diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index 79a5fd3b791..f2ca8d992b9 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -71,5 +71,42 @@ async def test_mlflow_request_tags_functionality(): tags_param = call_args.kwargs.get('tags', {}) expected_tags = {"tag1": "", "tag2": "", "production": ""} assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}" - - print("✅ Request tags properly transformed and passed to MLflow trace") + + + +def test_mlflow_token_usage_attribute_structure(): + """Ensure token usage attributes are formatted with mlflow.chat.tokenUsage.""" + + mock_mlflow_tracking = MagicMock() + mock_mlflow_tracking.MlflowClient = MagicMock() + + with patch.dict( + "sys.modules", + { + "mlflow": MagicMock(), + "mlflow.tracking": mock_mlflow_tracking, + "mlflow.tracing.utils": MagicMock(), + }, + ): + from litellm.integrations.mlflow import MlflowLogger + + mlflow_logger = MlflowLogger() + + attrs = mlflow_logger._extract_attributes( # type: ignore + { + "litellm_call_id": "123", + "call_type": "completion", + "model": "gpt-3.5-turbo", + "standard_logging_object": { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12, + }, + } + ) + + assert attrs["mlflow.chat.tokenUsage"] == { + "input_tokens": 5, + "output_tokens": 7, + "total_tokens": 12, + } From 025ce175649574bd9c2d4ce91d24762d4ab8f77d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 14 Aug 2025 15:30:45 -0700 Subject: [PATCH 8/9] =?UTF-8?q?bump:=20version=201.75.5=20=E2=86=92=201.75?= =?UTF-8?q?.6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f88fc02c06a..225faf07298 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.75.5" +version = "1.75.6" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -155,7 +155,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.75.5" +version = "1.75.6" version_files = [ "pyproject.toml:^version" ] From 936c36bd5f0c3cbac2aa710a5182c72cd06e714e Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 15:41:58 -0700 Subject: [PATCH 9/9] Increase timeout for test-litellm workflow from 20 to 25 minutes to accommodate longer test execution times. --- .github/workflows/test-litellm.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 4ec3dcbb4cf..7e67aee8d73 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -7,7 +7,7 @@ on: jobs: test: runs-on: ubuntu-latest - timeout-minutes: 20 + timeout-minutes: 25 steps: - uses: actions/checkout@v4