scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 96

fix(proxy): keep the raw client model out of the stored request body when a spend row is placeholdered

Browse source 
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.
This commit is contained in:
mateo-berri 2026-09-19 03:20:31 -07:00
parent df3a37857c
commit a49fbc6272
4 changed files with 92 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -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:
"""

9
tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py
View file

@ -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"]

9
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -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: ")

58
tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -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"
),
),
],
)