From a49fbc6272a5ba8dd7b90918ba1ebd0ddfc6ffb1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:20:31 -0700 Subject: [PATCH] fix(proxy): keep the raw client model out of the stored request body when a spend row is placeholdered With store_prompts_in_spend_logs on, the persisted request body kept the client's model string even when the row's model, model_group, and error text had been replaced by the unknown-model placeholder. The body's model now takes the same placeholder on those rows. Also annotates the new test locals with Final and wraps the four test lines that ran past 120 characters. --- .../spend_tracking/spend_tracking_utils.py | 28 ++++++++- .../test_files_common_utils.py | 9 ++- .../test_pass_through_endpoints.py | 9 +-- .../test_spend_tracking_utils.py | 58 ++++++++++++++++++- 4 files changed, 92 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 72af80bb3eb..055e128e0c4 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -756,7 +756,11 @@ def get_logging_payload( ), response=_get_response_for_spend_logs_payload(payload=standard_logging_payload, kwargs=kwargs), proxy_server_request=_get_proxy_server_request_for_spend_logs_payload( - metadata=metadata, litellm_params=litellm_params, kwargs=kwargs + metadata=metadata, + litellm_params=( + _placeholder_stored_request_body_model(litellm_params) if model_is_placeholdered else litellm_params + ), + kwargs=kwargs, ), session_id=_get_session_id_for_spend_log( kwargs=kwargs, @@ -1416,9 +1420,29 @@ def _convert_mapping_to_json_serializable(obj: Mapping[str, object]) -> dict[str return dict(obj) +def _placeholder_stored_request_body_model(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + proxy_server_request: Final = litellm_params.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return litellm_params + request_body: Final = proxy_server_request.get("body") + if not isinstance(request_body, Mapping) or "model" not in request_body: + return litellm_params + return MappingProxyType( + { + **litellm_params, + "proxy_server_request": MappingProxyType( + { + **proxy_server_request, + "body": MappingProxyType({**request_body, "model": UNKNOWN_MODEL_SPEND_LOG_MODEL}), + } + ), + } + ) + + def _get_proxy_server_request_for_spend_logs_payload( metadata: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], kwargs: dict | None = None, ) -> str: """ diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 77cd1358606..ef8af7bdbd3 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -1,4 +1,5 @@ from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -13,15 +14,17 @@ from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.utils import LiteLLMBatch -_RAW_MODEL_WITH_PROMPT = "opus-4.6 Please summarize my medical records\nPatient has diabetes" +_RAW_MODEL_WITH_PROMPT: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes" def test_get_credentials_for_model_rejects_an_unknown_model_without_persisting_the_raw_model(): - llm_router = MagicMock() + llm_router: Final = MagicMock() llm_router.get_deployment_credentials_with_provider.return_value = None with pytest.raises(ProxyModelNotFoundError) as raised: - get_credentials_for_model(llm_router=llm_router, model_id=_RAW_MODEL_WITH_PROMPT, operation_context="file upload") + get_credentials_for_model( + llm_router=llm_router, model_id=_RAW_MODEL_WITH_PROMPT, operation_context="file upload" + ) assert (raised.value.status_code, handle_exception_on_proxy(raised.value).code) == (400, "400") assert _RAW_MODEL_WITH_PROMPT in raised.value.detail["error"] diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 81911665b62..fb89e3a6973 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7,6 +7,7 @@ from collections.abc import Callable from contextlib import ExitStack, contextmanager from io import BytesIO from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -6448,8 +6449,8 @@ async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_err async def test_chat_completion_pass_through_endpoint_keeps_the_raw_model_out_of_the_spend_log_error( monkeypatch: pytest.MonkeyPatch, ): - raw_model = "opus-4.6 Please summarize my medical records\nPatient has diabetes" - proxy_logging = MagicMock() + raw_model: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes" + proxy_logging: Final = MagicMock() proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) proxy_logging.post_call_failure_hook = AsyncMock() @@ -6462,7 +6463,7 @@ async def test_chat_completion_pass_through_endpoint_keeps_the_raw_model_out_of_ monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) - request = MagicMock(spec=Request) + request: Final = MagicMock(spec=Request) request.body = AsyncMock( return_value=json.dumps({"model": raw_model, "messages": [{"role": "user", "content": "hi"}]}).encode() ) @@ -6475,7 +6476,7 @@ async def test_chat_completion_pass_through_endpoint_keeps_the_raw_model_out_of_ user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) - logged_exception = proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"] + logged_exception: Final = proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"] assert isinstance(logged_exception, ProxyModelNotFoundError) assert logged_exception.retryable_with_model_read_through is False assert logged_exception.spend_log_error_message.startswith("completion: ") 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 fbe8b10363f..50f4d2dcf5a 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 @@ -1107,6 +1107,50 @@ def test_get_logging_payload_keeps_a_whitespace_model_name_on_success_or_a_route assert payload["model"] == _RAW_MODEL_WITH_PROMPT +@pytest.mark.parametrize("redact_messages", [False, True]) +@pytest.mark.parametrize( + ("metadata", "expected_stored_model"), + [ + ({"user_api_key": "sk-test", "status": "failure"}, UNKNOWN_MODEL_SPEND_LOG_MODEL), + ( + {"user_api_key": "sk-test", "status": "failure", "model_info": {"id": "routed-deployment"}}, + _RAW_MODEL_WITH_PROMPT, + ), + ], +) +def test_get_logging_payload_placeholders_the_stored_request_body_model_only_when_the_row_is_placeholdered( + monkeypatch: pytest.MonkeyPatch, + metadata: dict[str, object], + expected_stored_model: str, + redact_messages: bool, +): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"store_prompts_in_spend_logs": True}) + kwargs: Final = { + "model": _RAW_MODEL_WITH_PROMPT, + "call_type": "amoderation", + "standard_callback_dynamic_params": {"turn_off_message_logging": redact_messages}, + "litellm_params": { + "metadata": metadata, + "proxy_server_request": { + "url": "http://localhost:4000/v1/moderations", + "body": {"input": "hi", "model": _RAW_MODEL_WITH_PROMPT}, + }, + }, + } + + payload: Final = get_logging_payload( + kwargs=kwargs, + response_obj=ValueError("Invalid value for 'model'"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + stored_request_body: Final = json.loads(payload["proxy_server_request"]) + assert stored_request_body["model"] == expected_stored_model + + _WHITESPACE_MODEL_GROUP: Final = "Broken GPT Mini" _WHITESPACE_MODEL_GROUP_ALIAS: Final = "Broken GPT Alias" _COOLDOWN_ERROR_MESSAGE: Final = ( @@ -1223,7 +1267,9 @@ def test_get_logging_payload_persists_no_raw_model_for_a_prompt_shaped_moderatio error_information: Final = _sanitize_error_information_for_spend_logs( StandardLoggingPayloadSetup.get_error_information( original_exception=provider_rejection, - traceback_str=f"Traceback (most recent call last):\n ...\nlitellm.exceptions.BadRequestError: {provider_rejection}", + traceback_str=( + f"Traceback (most recent call last):\n ...\nlitellm.exceptions.BadRequestError: {provider_rejection}" + ), ), original_exception=provider_rejection, ) @@ -1272,8 +1318,14 @@ _TRUNCATION_MARKER_TEXT: Final = ( f"OpenAIException - {{'message': '{UNKNOWN_MODEL_SPEND_LOG_MODEL}'}}", ), ( - f"Invalid model {_RAW_MODEL_WITH_PROMPT[:20]}{_TRUNCATION_MARKER_TEXT}{_RAW_MODEL_WITH_PROMPT[30:]} rejected", - f"Invalid model {UNKNOWN_MODEL_SPEND_LOG_MODEL}{_TRUNCATION_MARKER_TEXT}{UNKNOWN_MODEL_SPEND_LOG_MODEL} rejected", + ( + f"Invalid model {_RAW_MODEL_WITH_PROMPT[:20]}{_TRUNCATION_MARKER_TEXT}" + f"{_RAW_MODEL_WITH_PROMPT[30:]} rejected" + ), + ( + f"Invalid model {UNKNOWN_MODEL_SPEND_LOG_MODEL}{_TRUNCATION_MARKER_TEXT}" + f"{UNKNOWN_MODEL_SPEND_LOG_MODEL} rejected" + ), ), ], )