mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(spend): keep a failed follow-up's billing state off the request-merged call details
This commit is contained in:
mateo-berri
parent
8a7fbe9e47
commit
5ab5cd7181
2 changed files with 43 additions and 36 deletions
|
|
@ -591,8 +591,6 @@ _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
|
|||
_SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object])
|
||||
_GOOGLE_USAGE_FIELDS: Final = frozenset(UsageMetadata.__annotations__)
|
||||
_REDACTED_RESULT: Final = TypeAdapter(object)
|
||||
_FAILED_ATTEMPTS_USAGE: Final = TypeAdapter(tuple[ServedRound, ...])
|
||||
_ATTEMPT_IDS: Final = TypeAdapter(tuple[str, ...])
|
||||
_LoggedResponse: Final = TypeVar("_LoggedResponse")
|
||||
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
|
||||
|
||||
|
|
@ -941,6 +939,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self._sync_success_scheduled: bool = False
|
||||
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self._stream_failure_claimed: bool = False
|
||||
self._stash_owed_on_success: bool = False
|
||||
self._stash_billed_by_continuation: bool = False
|
||||
self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = ()
|
||||
self._failed_attempt_ids_owed_on_success: tuple[str, ...] = ()
|
||||
self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = ()
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
||||
def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
|
||||
|
|
@ -2748,8 +2751,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
own usage while the request's failure log still bills every attempt."""
|
||||
usage: Final = self._model_call_details_view.get("combined_usage_object")
|
||||
response_cost: Final = self._model_call_details_view.get("response_cost")
|
||||
billed_by_continuation: Final = self.model_call_details.pop("stash_billed_by_continuation", False) is True
|
||||
owed_on_success: Final = self.model_call_details.pop("stash_owed_on_success", False) is True
|
||||
billed_by_continuation: Final = self._stash_billed_by_continuation
|
||||
owed_on_success: Final = self._stash_owed_on_success
|
||||
self._stash_billed_by_continuation = False
|
||||
self._stash_owed_on_success = False
|
||||
captured_attempt_ids: Final = (
|
||||
self.capture_failed_attempt_usage(usage)
|
||||
if isinstance(usage, Usage) and type(self).usage_capture_callback is not None
|
||||
|
|
@ -2770,18 +2775,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
if not owed_on_success or billed_by_continuation:
|
||||
return
|
||||
self.model_call_details["failed_attempts_owed_on_success"] = (
|
||||
*self.model_call_details.get("failed_attempts_owed_on_success", ()),
|
||||
attempt,
|
||||
)
|
||||
if captured_attempt_ids:
|
||||
self.model_call_details["failed_attempt_ids_owed_on_success"] = (
|
||||
*_ATTEMPT_IDS.validate_python(self.model_call_details.get("failed_attempt_ids_owed_on_success", ())),
|
||||
*captured_attempt_ids,
|
||||
)
|
||||
self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt)
|
||||
self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids)
|
||||
|
||||
def mark_stash_billed_by_continuation(self) -> None:
|
||||
self.model_call_details["stash_billed_by_continuation"] = True
|
||||
self._stash_billed_by_continuation = True
|
||||
|
||||
def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]:
|
||||
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY
|
||||
|
|
@ -2864,25 +2862,15 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
own_cost: Final = self.model_call_details.get("response_cost")
|
||||
if own_usage is None or not isinstance(own_cost, (int, float)):
|
||||
return
|
||||
failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(
|
||||
(
|
||||
*self.model_call_details.get("failed_attempts_billed_on_success", ()),
|
||||
*self.model_call_details.pop("failed_attempts_owed_on_success", ()),
|
||||
)
|
||||
)
|
||||
failed_attempts: Final = self._failed_attempts_owed_on_success
|
||||
if not failed_attempts:
|
||||
return
|
||||
failed_attempts_cost: Final = sum(cost for _, cost in failed_attempts)
|
||||
self.model_call_details.pop("failed_attempts_usage", None)
|
||||
self.failed_attempts_billed_on_success = failed_attempts
|
||||
self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts
|
||||
included_attempt_ids: Final = _ATTEMPT_IDS.validate_python(
|
||||
self.model_call_details.pop("failed_attempt_ids_owed_on_success", ())
|
||||
)
|
||||
if included_attempt_ids:
|
||||
self.model_call_details["_accounting_included_attempt_ids"] = (
|
||||
*_ATTEMPT_IDS.validate_python(self.model_call_details.get("_accounting_included_attempt_ids", ())),
|
||||
*included_attempt_ids,
|
||||
)
|
||||
if self._failed_attempt_ids_owed_on_success:
|
||||
self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success
|
||||
self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects(
|
||||
usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)]
|
||||
)
|
||||
|
|
@ -2906,7 +2894,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
def record_served_rounds_for_failure(self, served_rounds: tuple[ServedRound, ...]) -> None:
|
||||
if not served_rounds:
|
||||
return
|
||||
self.model_call_details["stash_owed_on_success"] = True
|
||||
self._stash_owed_on_success = True
|
||||
self.record_partial_usage_for_failure(
|
||||
served_rounds[0][0]
|
||||
if len(served_rounds) == 1
|
||||
|
|
@ -7892,11 +7880,8 @@ def _with_failed_attempts_cost(breakdown: CostBreakdown, own_cost: float, failed
|
|||
)
|
||||
|
||||
|
||||
def _cached_cost_breakdown(served: CostBreakdown | None, model_call_details: Mapping[str, object]) -> CostBreakdown:
|
||||
def _cached_cost_breakdown(served: CostBreakdown | None, failed_attempts: tuple[ServedRound, ...]) -> CostBreakdown:
|
||||
cached: Final = _cache_hit_cost_breakdown(served)
|
||||
failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(
|
||||
model_call_details.get("failed_attempts_billed_on_success", ())
|
||||
)
|
||||
if not failed_attempts:
|
||||
return cached
|
||||
return _with_failed_attempts_cost(cached, 0.0, sum(cost for _, cost in failed_attempts))
|
||||
|
|
@ -8095,7 +8080,7 @@ def get_standard_logging_object_payload(
|
|||
if finalized_cost is not None
|
||||
else cost_breakdown_with_guardrail(
|
||||
(
|
||||
_cached_cost_breakdown(uncached_price.cost_breakdown, kwargs)
|
||||
_cached_cost_breakdown(uncached_price.cost_breakdown, logging_obj.failed_attempts_billed_on_success)
|
||||
if uncached_price is not None
|
||||
else logging_obj.cost_breakdown
|
||||
),
|
||||
|
|
|
|||
|
|
@ -4418,7 +4418,9 @@ def test_get_logging_payload_uses_recovered_combined_usage_on_failure():
|
|||
assert payload["total_tokens"] == 31
|
||||
|
||||
|
||||
def _fallback_after_a_failed_follow_up(cache_hit: bool):
|
||||
def _fallback_after_a_failed_follow_up(
|
||||
cache_hit: bool, request_params: Mapping[str, object] = MappingProxyType({})
|
||||
):
|
||||
from litellm.litellm_core_utils.hidden_params import set_hidden_param
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
|
@ -4437,7 +4439,7 @@ def _fallback_after_a_failed_follow_up(cache_hit: bool):
|
|||
model="fallback-model",
|
||||
custom_llm_provider="openai",
|
||||
litellm_params={"metadata": {"user_api_key": "sk-test"}},
|
||||
optional_params={},
|
||||
optional_params=dict(request_params),
|
||||
)
|
||||
fallback: Final = ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4))
|
||||
set_hidden_param(fallback, "response_cost", 0.1)
|
||||
|
|
@ -4461,6 +4463,26 @@ def test_success_row_bills_a_failed_follow_ups_cost_and_tokens(cache_hit: bool,
|
|||
assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (3 + 5, 1 + 2, 4 + 7)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("cache_hit", "own_cost"), [(False, 0.1), (True, 0.0)], ids=["fresh", "cached"])
|
||||
def test_request_params_cannot_change_what_a_failed_follow_up_bills(cache_hit: bool, own_cost: float):
|
||||
forged_attempt: Final = [{"prompt_tokens": 1000, "completion_tokens": 1000, "total_tokens": 2000}, -100.0]
|
||||
logger, fallback, _, now = _fallback_after_a_failed_follow_up(
|
||||
cache_hit,
|
||||
{
|
||||
"stash_billed_by_continuation": True,
|
||||
"failed_attempts_owed_on_success": [forged_attempt],
|
||||
"failed_attempts_billed_on_success": [forged_attempt],
|
||||
},
|
||||
)
|
||||
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs=logger.model_call_details, response_obj=fallback, start_time=now, end_time=now
|
||||
)
|
||||
|
||||
assert payload["spend"] == pytest.approx(own_cost + 0.4)
|
||||
assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (8, 3, 11)
|
||||
|
||||
|
||||
def test_cached_fallback_cost_breakdown_adds_up_to_the_billed_failed_follow_up():
|
||||
_, _, prepared, _ = _fallback_after_a_failed_follow_up(cache_hit=True)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue