From ba1ddc6d6cb1441b34bdd6bd09e0dbf7ffd9b236 Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Mon, 27 Apr 2026 11:08:10 -0700 Subject: [PATCH] Unify cost calc in success_handler dict and typed branches --- litellm/litellm_core_utils/litellm_logging.py | 31 ++-- .../test_litellm_logging.py | 148 ++++++++++++++++++ 2 files changed, 166 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fb103afea04..1a885e31617 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, @@ -1476,6 +1478,10 @@ class Logging(LiteLLMLoggingBaseClass): Calculate response cost using result + logging object variables. used for consistent cost calculation across response headers + logging integrations. + + `dict`/`list` results (e.g. non-streaming vertex rawPredict) are + accepted - the underlying calculator duck-types over the response + shape via `usage` extraction. """ if cache_hit is None: @@ -1744,6 +1750,13 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, ): + """ + Compute response_cost, build the standard logging payload, and emit + it. Handles both typed responses (ModelResponse, etc.) and raw dict + / list results (e.g. non-streaming vertex rawPredict, anthropic + passthrough). Without this, the dict path skipped cost calculation + and spend was logged as $0.00. + """ hidden_params = getattr(logging_result, "_hidden_params", {}) if hidden_params: if self.model_call_details.get("litellm_params") is not None: @@ -1875,26 +1888,18 @@ class Logging(LiteLLMLoggingBaseClass): and result is not None and self.stream is not True ): + # Single path for both typed responses (ModelResponse, etc.) + # and raw dict/list results. Previously the dict branch + # skipped cost calculation, which caused non-streaming + # rawPredict spend to log as $0.00. 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..ad9c25e0af5 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,151 @@ 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 raw dict responses (e.g. vertex rawPredict via + /v1/messages from Claude Code) must run through the cost calculator. + Pre-fix this branch skipped cost calc and spend was logged as $0.00. + """ + 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(): + """ + If a pass-through handler has already computed and stored response_cost + on model_call_details, the unified path must NOT overwrite it. + """ + 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 same helper. Sanity + check that unifying the two branches did not regress the typed path. + """ + 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