fix(proxy): only remap prompt-shaped model strings on failed spend rows

This commit is contained in:
mateo-berri 2026-09-10 15:36:32 -07:00
parent 368de5df68
commit 098bfe0bbc
2 changed files with 26 additions and 11 deletions

View file

@ -443,9 +443,12 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
resolved_model: Final = (
standard_logging_payload.get("model") if standard_logging_payload is not None else None
) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
failed_with_prompt_shaped_model: Final = _get_status_for_spend_log(
metadata=metadata
) == "failure" and not _looks_like_model_name(resolved_model)
model_name: Final = (
UNKNOWN_MODEL_SPEND_LOG_MODEL
if rejected_as_unknown_model or not _looks_like_model_name(resolved_model)
if rejected_as_unknown_model or failed_with_prompt_shaped_model
else resolved_model
)
litellm_call_id: Final = cast(

View file

@ -141,9 +141,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"
@ -976,6 +974,24 @@ def test_get_logging_payload_replaces_rejected_or_prompt_shaped_models_with_the_
assert payload["model"] == expected_model
def test_get_logging_payload_keeps_a_whitespace_model_name_on_success():
kwargs: Final = {
"model": _RAW_MODEL_WITH_PROMPT,
"messages": [{"role": "user", "content": "hi"}],
"call_type": "acompletion",
"litellm_params": {"metadata": {"user_api_key": "sk-test"}},
}
payload: Final = get_logging_payload(
kwargs=kwargs,
response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[]),
start_time=datetime.datetime.now(timezone.utc),
end_time=datetime.datetime.now(timezone.utc),
)
assert payload["model"] == _RAW_MODEL_WITH_PROMPT
@patch("litellm.proxy.proxy_server.master_key", None)
@patch("litellm.proxy.proxy_server.general_settings", {})
def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_none():
@ -4306,7 +4322,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',
)
@ -4344,9 +4360,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 (
@ -4422,9 +4436,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():