From 6a0cf3db50623de48a7f9f9c02afc87fe5b3a6c9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 24 Mar 2025 13:45:39 -0700 Subject: [PATCH 1/5] fix(litellm_logging.py): always log the api base Fixes issue where api base missing from spend logs due to refactor --- litellm/litellm_core_utils/litellm_logging.py | 25 +++++++++------- .../litellm_core_utils/model_param_helper.py | 15 ++++++++-- litellm/proxy/_new_secret_config.yaml | 5 +++- .../test_spend_management_endpoints.py | 30 +++++++++++++++++++ 4 files changed, 61 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 67511968e21..3565c4468c5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -518,6 +518,16 @@ class Logging(LiteLLMLoggingBaseClass): } return data + def _get_masked_api_base(self, api_base: str) -> str: + if "key=" in api_base: + # Find the position of "key=" in the string + key_index = api_base.find("key=") + 4 + # Mask the last 5 characters after "key=" + masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] + else: + masked_api_base = api_base + return str(masked_api_base) + def _pre_call(self, input, api_key, model=None, additional_args={}): """ Common helper function across the sync + async pre-call function @@ -531,6 +541,9 @@ class Logging(LiteLLMLoggingBaseClass): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model + self.model_call_details["litellm_params"]["api_base"] = ( + self._get_masked_api_base(additional_args.get("api_base", "")) + ) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 @@ -714,15 +727,6 @@ class Logging(LiteLLMLoggingBaseClass): headers = {} data = additional_args.get("complete_input_dict", {}) api_base = str(additional_args.get("api_base", "")) - if "key=" in api_base: - # Find the position of "key=" in the string - key_index = api_base.find("key=") + 4 - # Mask the last 5 characters after "key=" - masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] - else: - masked_api_base = api_base - self.model_call_details["litellm_params"]["api_base"] = masked_api_base - curl_command = self._get_request_curl_command( api_base=api_base, headers=headers, @@ -737,11 +741,12 @@ class Logging(LiteLLMLoggingBaseClass): def _get_request_curl_command( self, api_base: str, headers: Optional[dict], additional_args: dict, data: dict ) -> str: + masked_api_base = self._get_masked_api_base(api_base) if headers is None: headers = {} curl_command = "\n\nPOST Request Sent from LiteLLM:\n" curl_command += "curl -X POST \\\n" - curl_command += f"{api_base} \\\n" + curl_command += f"{masked_api_base} \\\n" masked_headers = self._get_masked_headers(headers) formatted_headers = " ".join( [f"-H '{k}: {v}'" for k, v in masked_headers.items()] diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index d792ede2822..b7d8fc19d14 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -138,13 +138,22 @@ class ModelParamHelper: TranscriptionCreateParamsNonStreaming, TranscriptionCreateParamsStreaming, ) - non_streaming_kwargs = set(getattr(TranscriptionCreateParamsNonStreaming, "__annotations__", {}).keys()) - streaming_kwargs = set(getattr(TranscriptionCreateParamsStreaming, "__annotations__", {}).keys()) + + non_streaming_kwargs = set( + getattr( + TranscriptionCreateParamsNonStreaming, "__annotations__", {} + ).keys() + ) + streaming_kwargs = set( + getattr( + TranscriptionCreateParamsStreaming, "__annotations__", {} + ).keys() + ) all_transcription_kwargs = non_streaming_kwargs.union(streaming_kwargs) return all_transcription_kwargs except Exception as e: - verbose_logger.warning("Error getting transcription kwargs %s", str(e)) + verbose_logger.debug("Error getting transcription kwargs %s", str(e)) return set() @staticmethod diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index cd49647464b..fbe3471663c 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -5,7 +5,10 @@ model_list: api_key: os.environ/AZURE_API_KEY api_base: http://0.0.0.0:8090 rpm: 3 - + - model_name: "gpt-4o-mini-openai" + litellm_params: + model: gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY litellm_settings: num_retries: 0 diff --git a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 415acdecee7..0a55991215a 100644 --- a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1,3 +1,4 @@ +import asyncio import datetime import json import os @@ -11,6 +12,10 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path +from unittest.mock import patch + +import litellm +from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.proxy_server import app, prisma_client @@ -400,3 +405,28 @@ async def test_ui_view_spend_logs_unauthorized(client): headers={"Authorization": "Bearer invalid-token"}, ) assert response.status_code == 401 or response.status_code == 403 + + +class TestSpendLogsPayload: + @pytest.mark.asyncio + async def test_spend_logs_payload_api_base_e2e(self): + litellm.callbacks = [_ProxyDBLogger(message_logging=False)] + # litellm._turn_on_debug() + + with patch.object( + litellm.proxy.proxy_server, "_set_spend_logs_payload" + ) as mock_client: + response = await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello, world!"}], + mock_response="Hello, world!", + ) + + assert response.choices[0].message.content == "Hello, world!" + + await asyncio.sleep(1) + + mock_client.assert_called_once() + + def test_spend_logs_payload_with_prompts_enabled(self): + pass From 75722c4b137508d78519663cf7d2856b46a713a1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 24 Mar 2025 14:45:20 -0700 Subject: [PATCH 2/5] test: add unit test --- .../test_litellm_logging.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 tests/litellm/litellm_core_utils/test_litellm_logging.py diff --git a/tests/litellm/litellm_core_utils/test_litellm_logging.py b/tests/litellm/litellm_core_utils/test_litellm_logging.py new file mode 100644 index 00000000000..eedc15dd362 --- /dev/null +++ b/tests/litellm/litellm_core_utils/test_litellm_logging.py @@ -0,0 +1,34 @@ +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import time + +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging + + +@pytest.fixture +def logging_obj(): + return LitellmLogging( + model="bedrock/claude-3-5-sonnet-20240620-v1:0", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="12345", + function_id="1245", + ) + + +def test_get_masked_api_base(logging_obj): + api_base = "https://api.openai.com/v1" + masked_api_base = logging_obj._get_masked_api_base(api_base) + assert masked_api_base == "https://api.openai.com/v1" + assert type(masked_api_base) == str From e1bad1befac865f5a26ca3f6d803c7d3bf8db284 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 24 Mar 2025 15:12:18 -0700 Subject: [PATCH 3/5] test: add e2e testing --- .../test_spend_management_endpoints.py | 56 ++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 0a55991215a..1943ea19d03 100644 --- a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -15,6 +15,7 @@ sys.path.insert( from unittest.mock import patch import litellm +from litellm.proxy._types import SpendLogsPayload from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.proxy_server import app, prisma_client @@ -415,11 +416,12 @@ class TestSpendLogsPayload: with patch.object( litellm.proxy.proxy_server, "_set_spend_logs_payload" - ) as mock_client: + ) as mock_client, patch.object(litellm.proxy.proxy_server, "prisma_client"): response = await litellm.acompletion( model="gpt-4o", messages=[{"role": "user", "content": "Hello, world!"}], mock_response="Hello, world!", + metadata={"user_api_key_end_user_id": "test_user_1"}, ) assert response.choices[0].message.content == "Hello, world!" @@ -428,5 +430,57 @@ class TestSpendLogsPayload: mock_client.assert_called_once() + kwargs = mock_client.call_args.kwargs + payload: SpendLogsPayload = kwargs["payload"] + expected_payload = SpendLogsPayload( + **{ + "request_id": "chatcmpl-34df56d5-4807-45c1-bb99-61e52586b802", + "call_type": "acompletion", + "api_key": "", + "cache_hit": "None", + "startTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 975883, tzinfo=datetime.timezone.utc + ), + "endTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 989132, tzinfo=datetime.timezone.utc + ), + "completionStartTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 989132, tzinfo=datetime.timezone.utc + ), + "model": "gpt-4o", + "user": "", + "team_id": "", + "metadata": '{"applied_guardrails": [], "batch_models": null, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "cache_key": "Cache OFF", + "spend": 0.00022500000000000002, + "total_tokens": 30, + "prompt_tokens": 10, + "completion_tokens": 20, + "request_tags": "[]", + "end_user": "test_user_1", + "api_base": "", + "model_group": "", + "model_id": "", + "requester_ip_address": None, + "custom_llm_provider": "openai", + "messages": "{}", + "response": "{}", + } + ) + + for key, value in expected_payload.items(): + if key in [ + "request_id", + "startTime", + "endTime", + "completionStartTime", + "endTime", + ]: + assert payload[key] is not None + else: + assert ( + payload[key] == value + ), f"Expected {key} to be {value}, but got {payload[key]}" + def test_spend_logs_payload_with_prompts_enabled(self): pass From 1dc15ef5bf27c888625d43cd01c166c935863403 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 24 Mar 2025 15:29:47 -0700 Subject: [PATCH 4/5] test(test_spend_management_endpoints.py): guarantee consistent spend logs --- .../test_spend_management_endpoints.py | 196 +++++++++++++++++- 1 file changed, 192 insertions(+), 4 deletions(-) diff --git a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 1943ea19d03..6623e322eed 100644 --- a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -12,12 +12,13 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import litellm from litellm.proxy._types import SpendLogsPayload from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.proxy_server import app, prisma_client +from litellm.router import Router @pytest.fixture @@ -410,7 +411,7 @@ async def test_ui_view_spend_logs_unauthorized(client): class TestSpendLogsPayload: @pytest.mark.asyncio - async def test_spend_logs_payload_api_base_e2e(self): + async def test_spend_logs_payload_e2e(self): litellm.callbacks = [_ProxyDBLogger(message_logging=False)] # litellm._turn_on_debug() @@ -482,5 +483,192 @@ class TestSpendLogsPayload: payload[key] == value ), f"Expected {key} to be {value}, but got {payload[key]}" - def test_spend_logs_payload_with_prompts_enabled(self): - pass + def mock_anthropic_response(*args, **kwargs): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "content": [{"text": "Hi! My name is Claude.", "type": "text"}], + "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", + "model": "claude-3-7-sonnet-20250219", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": None, + "type": "message", + "usage": {"input_tokens": 2095, "output_tokens": 503}, + } + return mock_response + + @pytest.mark.asyncio + async def test_spend_logs_payload_success_log_with_api_base(self): + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + litellm.callbacks = [_ProxyDBLogger(message_logging=False)] + # litellm._turn_on_debug() + + client = AsyncHTTPHandler() + + with patch.object( + litellm.proxy.proxy_server, "_set_spend_logs_payload" + ) as mock_client, patch.object( + litellm.proxy.proxy_server, "prisma_client" + ), patch.object( + client, "post", side_effect=self.mock_anthropic_response + ): + response = await litellm.acompletion( + model="claude-3-7-sonnet-20250219", + messages=[{"role": "user", "content": "Hello, world!"}], + metadata={"user_api_key_end_user_id": "test_user_1"}, + client=client, + ) + + assert response.choices[0].message.content == "Hi! My name is Claude." + + await asyncio.sleep(1) + + mock_client.assert_called_once() + + kwargs = mock_client.call_args.kwargs + payload: SpendLogsPayload = kwargs["payload"] + expected_payload = SpendLogsPayload( + **{ + "request_id": "chatcmpl-34df56d5-4807-45c1-bb99-61e52586b802", + "call_type": "acompletion", + "api_key": "", + "cache_hit": "None", + "startTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 975883, tzinfo=datetime.timezone.utc + ), + "endTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 989132, tzinfo=datetime.timezone.utc + ), + "completionStartTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 989132, tzinfo=datetime.timezone.utc + ), + "model": "claude-3-7-sonnet-20250219", + "user": "", + "team_id": "", + "metadata": '{"applied_guardrails": [], "batch_models": null, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "cache_key": "Cache OFF", + "spend": 0.01383, + "total_tokens": 2598, + "prompt_tokens": 2095, + "completion_tokens": 503, + "request_tags": "[]", + "end_user": "test_user_1", + "api_base": "https://api.anthropic.com/v1/messages", + "model_group": "", + "model_id": "", + "requester_ip_address": None, + "custom_llm_provider": "anthropic", + "messages": "{}", + "response": "{}", + } + ) + + for key, value in expected_payload.items(): + if key in [ + "request_id", + "startTime", + "endTime", + "completionStartTime", + "endTime", + ]: + assert payload[key] is not None + else: + assert ( + payload[key] == value + ), f"Expected {key} to be {value}, but got {payload[key]}" + + @pytest.mark.skip(reason="Test found bug, let's fix and repush.") + @pytest.mark.asyncio + async def test_spend_logs_payload_success_log_with_router(self): + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + litellm.callbacks = [_ProxyDBLogger(message_logging=False)] + # litellm._turn_on_debug() + + client = AsyncHTTPHandler() + + router = Router( + model_list=[ + { + "model_name": "my-anthropic-model-group", + "litellm_params": { + "model": "claude-3-7-sonnet-20250219", + }, + } + ] + ) + + with patch.object( + litellm.proxy.proxy_server, "_set_spend_logs_payload" + ) as mock_client, patch.object( + litellm.proxy.proxy_server, "prisma_client" + ), patch.object( + client, "post", side_effect=self.mock_anthropic_response + ): + response = await router.acompletion( + model="claude-3-7-sonnet-20250219", + messages=[{"role": "user", "content": "Hello, world!"}], + metadata={"user_api_key_end_user_id": "test_user_1"}, + client=client, + ) + + assert response.choices[0].message.content == "Hi! My name is Claude." + + await asyncio.sleep(1) + + mock_client.assert_called_once() + + kwargs = mock_client.call_args.kwargs + payload: SpendLogsPayload = kwargs["payload"] + expected_payload = SpendLogsPayload( + **{ + "request_id": "chatcmpl-34df56d5-4807-45c1-bb99-61e52586b802", + "call_type": "acompletion", + "api_key": "", + "cache_hit": "None", + "startTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 975883, tzinfo=datetime.timezone.utc + ), + "endTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 989132, tzinfo=datetime.timezone.utc + ), + "completionStartTime": datetime.datetime( + 2025, 3, 24, 22, 2, 42, 989132, tzinfo=datetime.timezone.utc + ), + "model": "claude-3-7-sonnet-20250219", + "user": "", + "team_id": "", + "metadata": '{"applied_guardrails": [], "batch_models": null, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "cache_key": "Cache OFF", + "spend": 0.01383, + "total_tokens": 2598, + "prompt_tokens": 2095, + "completion_tokens": 503, + "request_tags": "[]", + "end_user": "test_user_1", + "api_base": "https://api.anthropic.com/v1/messages", + "model_group": "my-anthropic-model-group", + "model_id": "", + "requester_ip_address": None, + "custom_llm_provider": "anthropic", + "messages": "{}", + "response": "{}", + } + ) + + for key, value in expected_payload.items(): + if key in [ + "request_id", + "startTime", + "endTime", + "completionStartTime", + "endTime", + ]: + assert payload[key] is not None + else: + assert ( + payload[key] == value + ), f"Expected {key} to be {value}, but got {payload[key]}" From 44e305648de5779a47789bb5b71e84b7c3867a70 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 24 Mar 2025 15:33:02 -0700 Subject: [PATCH 5/5] test(test_spend_management_endpoints.py): add unit testing for router + spend logs --- .../spend_tracking/test_spend_management_endpoints.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 6623e322eed..a5ee9ddf70d 100644 --- a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -580,7 +580,6 @@ class TestSpendLogsPayload: payload[key] == value ), f"Expected {key} to be {value}, but got {payload[key]}" - @pytest.mark.skip(reason="Test found bug, let's fix and repush.") @pytest.mark.asyncio async def test_spend_logs_payload_success_log_with_router(self): from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -597,6 +596,9 @@ class TestSpendLogsPayload: "litellm_params": { "model": "claude-3-7-sonnet-20250219", }, + "model_info": { + "id": "my-unique-model-id", + }, } ] ) @@ -609,7 +611,7 @@ class TestSpendLogsPayload: client, "post", side_effect=self.mock_anthropic_response ): response = await router.acompletion( - model="claude-3-7-sonnet-20250219", + model="my-anthropic-model-group", messages=[{"role": "user", "content": "Hello, world!"}], metadata={"user_api_key_end_user_id": "test_user_1"}, client=client, @@ -651,7 +653,7 @@ class TestSpendLogsPayload: "end_user": "test_user_1", "api_base": "https://api.anthropic.com/v1/messages", "model_group": "my-anthropic-model-group", - "model_id": "", + "model_id": "my-unique-model-id", "requester_ip_address": None, "custom_llm_provider": "anthropic", "messages": "{}",