[Fix] Unify cost calc in success_handler dict and typed branches (#26629)

* Unify cost calc in success_handler dict and typed branches

* Trim verbose comments and docstrings

---------

Co-authored-by: Michael Riad Zaky <michaelr@Mac.localdomain>
Co-authored-by: Michael Riad Zaky <michaelr@Michaels-MacBook-Air.local>
This commit is contained in:
Michael-RZ-Berri 2026-04-28 17:05:36 -07:00 committed by GitHub
parent 21ed38971d
commit 0520d5ce11
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 142 additions and 13 deletions

View file

@ -1467,6 +1467,8 @@ class Logging(LiteLLMLoggingBaseClass):
LiteLLMRealtimeStreamLoggingObject,
OpenAIModerationResponse,
"SearchResponse",
dict,
list,
],
cache_hit: Optional[bool] = None,
litellm_model_name: Optional[str] = None,
@ -1744,6 +1746,7 @@ class Logging(LiteLLMLoggingBaseClass):
start_time,
end_time,
):
"""Resolve hidden params, compute response cost, and emit the standard logging payload."""
hidden_params = getattr(logging_result, "_hidden_params", {})
if hidden_params:
if self.model_call_details.get("litellm_params") is not None:
@ -1877,24 +1880,12 @@ class Logging(LiteLLMLoggingBaseClass):
):
if self._is_recognized_call_type_for_logging(
logging_result=logging_result
):
) or isinstance(logging_result, (dict, list)):
self._process_hidden_params_and_response_cost(
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
)
elif isinstance(result, dict) or isinstance(result, list):
self.model_call_details["standard_logging_object"] = (
self._build_standard_logging_payload(
result, start_time, end_time
)
)
if (
standard_logging_payload := self.model_call_details.get(
"standard_logging_object"
)
) is not None:
emit_standard_logging_payload(standard_logging_payload)
elif standard_logging_object is not None:
self.model_call_details["standard_logging_object"] = (
standard_logging_object

View file

@ -2534,3 +2534,141 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload is not None
assert payload["litellm_call_id"] == call_id
def _make_dict_logging_obj():
"""Build a Logging instance configured for a non-streaming dict result."""
obj = LitellmLogging(
model="claude-haiku-4-5@20251001",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
litellm_call_id="test-call-id",
start_time=time.time(),
function_id="test-fn",
)
obj.model_call_details = {
"model": "claude-haiku-4-5@20251001",
"custom_llm_provider": "vertex_ai",
"litellm_params": {"metadata": {}},
"response_cost": None,
}
return obj
def test_success_handler_computes_cost_for_dict_response():
"""Non-streaming dict responses run through the cost calculator."""
logging_obj = _make_dict_logging_obj()
expected_cost = 0.42
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=expected_cost,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=False,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result={"id": "msg_1"},
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_called_once()
assert logging_obj.model_call_details["response_cost"] == expected_cost
def test_success_handler_preserves_precomputed_cost_for_dict_response():
"""Precomputed response_cost on model_call_details must not be overwritten."""
logging_obj = _make_dict_logging_obj()
precomputed_cost = 1.23
logging_obj.model_call_details["response_cost"] = precomputed_cost
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=9.99,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": precomputed_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=False,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result={"id": "msg_2"},
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_not_called()
assert logging_obj.model_call_details["response_cost"] == precomputed_cost
def test_success_handler_unified_helper_runs_for_typed_results():
"""Recognized typed responses still flow through the unified helper."""
logging_obj = _make_dict_logging_obj()
expected_cost = 0.10
typed_result = MagicMock()
typed_result._hidden_params = {}
with (
patch.object(
logging_obj,
"_response_cost_calculator",
return_value=expected_cost,
) as mock_calc,
patch.object(
logging_obj,
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
patch(
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
return_value=True,
),
patch.object(
logging_obj,
"_transform_usage_objects",
side_effect=lambda result: result,
),
):
logging_obj.success_handler(
result=typed_result,
start_time=time.time(),
end_time=time.time(),
)
mock_calc.assert_called_once()
assert logging_obj.model_call_details["response_cost"] == expected_cost