From 0520d5ce117a51994a862b8df6384fa6b1a52d74 Mon Sep 17 00:00:00 2001 From: Michael-RZ-Berri Date: Tue, 28 Apr 2026 17:05:36 -0700 Subject: [PATCH] [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 Co-authored-by: Michael Riad Zaky --- litellm/litellm_core_utils/litellm_logging.py | 17 +-- .../test_litellm_logging.py | 138 ++++++++++++++++++ 2 files changed, 142 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fb103afea04..829c1c9ca07 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 3348118a020..1764d9c609f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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