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 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/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" ] 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, + } 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 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