mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
[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:
parent
21ed38971d
commit
0520d5ce11
2 changed files with 142 additions and 13 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue