mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-22 05:24:40 +00:00
Unify cost calc in success_handler dict and typed branches
This commit is contained in:
Michael Riad Zaky
Michael Riad Zaky
parent
600d7b4a20
commit
ba1ddc6d6c
2 changed files with 166 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,
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue