From 4ee44302c2dfee92589eec7ed5c2689d43e21064 Mon Sep 17 00:00:00 2001 From: spawrks Date: Sun, 13 Sep 2026 09:24:15 -0500 Subject: [PATCH] fix(proxy): redact response metadata in spend logs --- .../spend_tracking/spend_tracking_utils.py | 30 +++-- .../test_proxy_config_unit_test.py | 8 +- .../test_spend_tracking_utils.py | 116 ++++++++++-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 4 files changed, 95 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 4a3cddfdc40..f72d73316ee 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -199,14 +199,19 @@ def _get_spend_logs_metadata( _already_redacted: Final = ( isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == _raw_key ) + store_responses: Final = _should_store_responses_in_spend_logs() clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["batch_successful_requests"] = batch_successful_requests clean_metadata["batch_failed_requests"] = batch_failed_requests - clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata + clean_metadata["mcp_tool_call_metadata"] = _get_mcp_tool_call_metadata_for_spend_logs_payload( + mcp_tool_call_metadata, + store_responses=store_responses, + ) clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload( - vector_store_request_metadata + vector_store_request_metadata, + store_responses=store_responses, ) clean_metadata["guardrail_information"] = _sanitize_guardrail_information_for_spend_logs(guardrail_information) clean_metadata["usage_object"] = usage_object @@ -1322,16 +1327,27 @@ def _get_proxy_server_request_for_spend_logs_payload( return "{}" +def _get_mcp_tool_call_metadata_for_spend_logs_payload( + mcp_tool_call_metadata: StandardLoggingMCPToolCall | None, + *, + store_responses: bool, +) -> StandardLoggingMCPToolCall | None: + if mcp_tool_call_metadata is None or store_responses: + return mcp_tool_call_metadata + return cast( + StandardLoggingMCPToolCall, + {key: value for key, value in mcp_tool_call_metadata.items() if key != "result"}, + ) + + def _get_vector_store_request_for_spend_logs_payload( vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None, + *, + store_responses: bool, ) -> list[StandardLoggingVectorStoreRequest] | None: - """ - If user does not want to store prompts and responses, then remove the content from the vector store request metadata - """ - if should_store_prompts_and_responses_in_spend_logs(): + if store_responses: return vector_store_request_metadata - # if user does not want to store prompts and responses, then remove the content from the vector store request metadata if vector_store_request_metadata is None: return None for vector_store_request in vector_store_request_metadata: diff --git a/tests/proxy_unit_tests/test_proxy_config_unit_test.py b/tests/proxy_unit_tests/test_proxy_config_unit_test.py index 909acfdceb7..ce05b769ad6 100644 --- a/tests/proxy_unit_tests/test_proxy_config_unit_test.py +++ b/tests/proxy_unit_tests/test_proxy_config_unit_test.py @@ -359,7 +359,9 @@ class TestYamlStorePromptsDbOverride: proxy_config = self._make_proxy_config_with_yaml_keys({"store_responses_in_spend_logs"}) test_general_settings = {"store_responses_in_spend_logs": False} - with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings): + with mock.patch( # test-quality-ok: verifies synchronization with the process-wide settings store + "litellm.proxy.proxy_server.general_settings", test_general_settings + ): await proxy_config._update_general_settings( db_general_settings={"store_responses_in_spend_logs": True}, ) @@ -384,7 +386,9 @@ class TestYamlStorePromptsDbOverride: proxy_config = self._make_proxy_config_with_yaml_keys({"master_key"}) test_general_settings = {"master_key": "sk-test"} - with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings): + with mock.patch( # test-quality-ok: verifies synchronization with the process-wide settings store + "litellm.proxy.proxy_server.general_settings", test_general_settings + ): await proxy_config._update_general_settings( db_general_settings={"store_responses_in_spend_logs": db_value}, ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 67ffa5647dc..8beb4871354 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -27,6 +27,7 @@ from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_messages_for_spend_logs_payload, + _get_mcp_tool_call_metadata_for_spend_logs_payload, _get_proxy_server_request_for_spend_logs_payload, _get_request_duration_ms, _get_response_for_spend_logs_payload, @@ -47,9 +48,11 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, + StandardLoggingMCPToolCall, StandardLoggingMetadata, StandardLoggingModelInformation, StandardLoggingPayload, + StandardLoggingVectorStoreRequest, ) @@ -80,10 +83,16 @@ def test_classifier_audit_spend_storage_obeys_privacy_and_truncation(monkeypatch "classifier_input": {"system": "rubric" * 1000, "messages": [{"role": "user", "content": "ask"}]}, "originating_request_masked": {"input": "source-only", "api_key": "REDACTED"}, } - stored: Final = json.loads(_get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params={"proxy_server_request": {"body": {"model": "classifier"}}}, - kwargs={"standard_logging_object": audit, "standard_callback_dynamic_params": {"turn_off_message_logging": redact}}, - )) + stored: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, + litellm_params={"proxy_server_request": {"body": {"model": "classifier"}}}, + kwargs={ + "standard_logging_object": audit, + "standard_callback_dynamic_params": {"turn_off_message_logging": redact}, + }, + ) + ) if not store_prompts or redact: assert "classifier_input" not in stored assert "originating_request_masked" not in stored @@ -209,9 +218,7 @@ def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") - cost_session: Final = _get_batch_trace_session_id( - call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" - ) + cost_session: Final = _get_batch_trace_session_id(call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost") assert create_session == cost_session == "batch-uid-1" @@ -657,52 +664,61 @@ def test_sanitize_request_body_for_spend_logs_payload_circular_reference(): assert sanitized == {"b": {"a": {}}} # Should return empty dict for circular reference -@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") -def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( - mock_should_store, +@pytest.mark.parametrize( + ("store_responses", "expected_text"), + [ + (False, REDACTED_BY_LITELM_STRING), + (True, "sensitive information"), + ], +) +def test_get_vector_store_request_for_spend_logs_payload_uses_response_setting( + store_responses: bool, expected_text: str ): - # When should_store_prompts_and_responses_in_spend_logs returns True - mock_should_store.return_value = True - - # Sample vector store request metadata - vector_store_request = [ + vector_store_request: Final[list[StandardLoggingVectorStoreRequest]] = [ {"vector_store_search_response": {"data": [{"content": [{"text": "sensitive information", "type": "text"}]}]}} ] - # When store_prompts is True, the original data should be returned unchanged - result = _get_vector_store_request_for_spend_logs_payload(vector_store_request) - assert result == vector_store_request - assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == "sensitive information" + result: Final = _get_vector_store_request_for_spend_logs_payload( + vector_store_request, + store_responses=store_responses, + ) - -@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") -def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( - mock_should_store, -): - # When should_store_prompts_and_responses_in_spend_logs returns False - mock_should_store.return_value = False - - # Sample vector store request metadata - vector_store_request = [ - {"vector_store_search_response": {"data": [{"content": [{"text": "sensitive information", "type": "text"}]}]}} - ] - - # When store_prompts is False, text should be redacted - result = _get_vector_store_request_for_spend_logs_payload(vector_store_request) assert result is not None - assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == REDACTED_BY_LITELM_STRING - # Ensure other fields are unchanged + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == expected_text assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] == "text" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") -def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_store): - # When input is None - mock_should_store.return_value = False - result = _get_vector_store_request_for_spend_logs_payload(None) +def test_get_vector_store_request_for_spend_logs_payload_null_input(): + result = _get_vector_store_request_for_spend_logs_payload(None, store_responses=False) assert result is None +@pytest.mark.parametrize( + ("store_responses", "expected_result"), + [ + (False, None), + (True, {"content": "sensitive response"}), + ], +) +def test_get_spend_logs_metadata_uses_response_setting_for_mcp_result( + store_responses: bool, expected_result: Mapping[str, str] | None +): + mcp_metadata: Final[StandardLoggingMCPToolCall] = { + "name": "search", + "arguments": {"query": "request content"}, + "result": {"content": "sensitive response"}, + } + + stored_mcp_metadata: Final = _get_mcp_tool_call_metadata_for_spend_logs_payload( + mcp_metadata, + store_responses=store_responses, + ) + assert stored_mcp_metadata is not None + assert stored_mcp_metadata["name"] == "search" + assert stored_mcp_metadata["arguments"] == {"query": "request content"} + assert stored_mcp_metadata.get("result") == expected_result + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store): """ @@ -1741,7 +1757,9 @@ def test_should_store_responses_in_spend_logs( settings, expected, ): - with patch("litellm.proxy.proxy_server.general_settings", settings): + with patch( # test-quality-ok: isolates the process-wide settings source for retention resolution + "litellm.proxy.proxy_server.general_settings", settings + ): assert _should_store_responses_in_spend_logs() is expected @@ -1765,7 +1783,9 @@ def test_spend_logs_can_store_request_without_response(): {"response": {"role": "assistant", "content": "Hi there!"}}, ) - with patch("litellm.proxy.proxy_server.general_settings", settings): + with patch( # test-quality-ok: isolates the process-wide settings source for the retention integration check + "litellm.proxy.proxy_server.general_settings", settings + ): request_result = _get_proxy_server_request_for_spend_logs_payload( metadata={}, litellm_params=kwargs["litellm_params"], @@ -4529,7 +4549,7 @@ ANTHROPIC_MESSAGES_SSE_CHUNKS: Final = ( 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n', 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' '"usage":{"output_tokens":4}}\n\n', - "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', ) @@ -4567,9 +4587,7 @@ def test_spend_log_request_id_is_the_message_id_a_non_streaming_messages_caller_ """ logging_obj = _anthropic_messages_logging_obj(stream=False) - logged_response = logging_obj._handle_anthropic_messages_response_logging( - result=ANTHROPIC_MESSAGES_RESPONSE - ) + logged_response = logging_obj._handle_anthropic_messages_response_logging(result=ANTHROPIC_MESSAGES_RESPONSE) assert logged_response.id == "msg_01Lit6806NonStreaming" assert ( @@ -4645,9 +4663,7 @@ def test_spend_log_request_id_still_falls_back_to_litellm_call_id_without_a_prov end_time=datetime.datetime.now(timezone.utc), logging_obj=logging_obj, ) - assert logging_obj.model_call_details["complete_streaming_response"].id == ( - "6806cafe-0000-4000-8000-000000000001" - ) + assert logging_obj.model_call_details["complete_streaming_response"].id == ("6806cafe-0000-4000-8000-000000000001") def test_spend_log_request_id_for_chat_completions_is_untouched(): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0ff8817c4e7..0e020ef170c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)