scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge 97250556a4 into c2c2a623c0

Browse source
This commit is contained in:
Gyanu Mayank 2026-09-12 23:49:21 -07:00 • committed by GitHub
commit 853726b790
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 241 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

63
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -411,22 +411,28 @@ class _ProxyDBLogger(CustomLogger):
"that no poll task will settle"
)
return
await _release_budget_reservation(budget_reservation=budget_reservation)
# Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
# Use .get() for "stream" to avoid KeyError on health checks.
# WS session wrappers (_aresponses_websocket, _arealtime) also reach here with
# result=None; their per-turn costs are tracked on the inner aresponses/realtime calls.
if sl_object is None and (
skippable_non_model_call = sl_object is None and (
not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime")
):
)
if skippable_non_model_call:
await _release_budget_reservation(budget_reservation=budget_reservation)
verbose_proxy_logger.warning(
"Cost tracking - skipping, no standard_logging_object for call_type=%s",
kwargs.get("call_type", "unknown"),
)
return
if kwargs.get("stream") is not True or (
kwargs.get("stream") is True and "complete_streaming_response" in kwargs
):
if _is_completed_proxy_cost_call(kwargs):
# Releasing to $0 treats the call as free. Leaving the hold
# open is also wrong: the next priced request only
# reconciles its own reservation, so this one would keep
# blocking shared counters until TTL. Settle at the
# admission estimate instead. No spend-log row. There is
# no real cost to write.
await _settle_unpriced_success_reservation(budget_reservation=budget_reservation)
if sl_object is not None:
cost_tracking_failure_debug_info: dict | str = (
sl_object["response_cost_failure_debug_info"]
@ -438,6 +444,7 @@ class _ProxyDBLogger(CustomLogger):
raise Exception(
f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
)
await _release_budget_reservation(budget_reservation=budget_reservation)
except Exception as e:
error_msg = f"Error in tracking cost callback - {e}\n Traceback:{traceback.format_exc()}"
model = kwargs.get("model", "")
@ -709,6 +716,33 @@ async def _update_database_and_spend_counters(
return True
def _is_completed_proxy_cost_call(kwargs: dict) -> bool:
return kwargs.get("stream") is not True or (
kwargs.get("stream") is True
and ("complete_streaming_response" in kwargs or "async_complete_streaming_response" in kwargs)
)
async def _settle_unpriced_success_reservation(budget_reservation: dict | None) -> None:
reserved_cost = float(budget_reservation.get("reserved_cost") or 0.0) if budget_reservation else 0.0
try:
await _reconcile_budget_reservation(
budget_reservation=budget_reservation,
actual_cost=reserved_cost,
)
except Exception: # noqa: BLE001 # settle can fail on cache/redis; still raise cost-tracking after invalidating
verbose_proxy_logger.exception("Failed to settle budget reservation after unpriced successful call")
try:
await _invalidate_budget_reservation_counters(
budget_reservation=budget_reservation,
)
except Exception: # noqa: BLE001 # invalidate is best-effort so the outer cost-tracking error still surfaces
verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after settle failed")
finally:
if budget_reservation is not None:
budget_reservation["finalized"] = True
async def _release_budget_reservation(budget_reservation: dict | None) -> None:
if budget_reservation is None:
return
@ -722,6 +756,23 @@ async def _release_budget_reservation(budget_reservation: dict | None) -> None:
)
async def _reconcile_budget_reservation(
budget_reservation: dict | None, # mutable-ok: same reservation payload _release_budget_reservation takes
actual_cost: float,
) -> None:
if budget_reservation is None:
return
from litellm.proxy.spend_tracking.budget_reservation import (
reconcile_budget_reservation,
)
await reconcile_budget_reservation(
budget_reservation=budget_reservation,
actual_cost=actual_cost,
)
async def _invalidate_budget_reservation_counters(
budget_reservation: dict | None,
) -> None:

2
test-quality-budget.json
View file

@ -3,7 +3,7 @@
"limit": 733
},
"TQ002": {
"limit": 737
"limit": 736
},
"TQ003": {
"limit": 62

186
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -390,7 +390,13 @@ async def test_track_cost_callback_releases_budget_reservation_when_spend_tracki
@pytest.mark.asyncio
async def test_track_cost_callback_releases_budget_reservation_when_response_cost_missing():
async def test_track_cost_callback_settles_budget_reservation_when_response_cost_missing():
"""A successful unpriced model call must not be refunded to $0.
Settling at the admission estimate converts the hold into budget spend
without inventing a spend-log row. Health checks still release, covered
separately.
"""
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation)
@ -412,13 +418,21 @@ async def test_track_cost_callback_releases_budget_reservation_when_response_cos
}
with (
patch(
patch( # test-quality-ok: callback reads proxy_logging_obj from the module; no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj",
) as mock_proxy_logging,
patch(
patch( # test-quality-ok: assert the hold is settled, not refunded through release_budget_reservation
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
patch( # test-quality-ok: settle is a proxy-internal reservation call, not an HTTP boundary
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
new_callable=AsyncMock,
) as mock_reconcile_budget_reservation,
patch( # test-quality-ok: unpriced settle must not write a spend-log row
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database,
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
@ -429,9 +443,175 @@ async def test_track_cost_callback_releases_budget_reservation_when_response_cos
end_time=datetime.now(),
)
mock_release_budget_reservation.assert_not_awaited()
mock_reconcile_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
actual_cost=0.5,
)
mock_update_database.assert_not_called()
mock_proxy_logging.failed_tracking_alert.assert_called()
# The hold is converted, not dropped: later traffic still sees the reserved cost.
assert budget_reservation["reserved_cost"] == 0.5
assert budget_reservation.get("finalized") is not True
@pytest.mark.asyncio
async def test_track_cost_callback_settles_async_stream_when_response_cost_missing():
"""Async streams record completion on async_complete_streaming_response."""
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation)
kwargs = {
"model": "gpt-4",
"call_type": "acompletion",
"stream": True,
"async_complete_streaming_response": {"usage": {"total_tokens": 10}},
"litellm_params": {
"metadata": {
"user_api_key_auth": user_api_key_auth,
},
},
"standard_logging_object": {
"response_cost": None,
"response_cost_failure_debug_info": "missing custom price",
"request_tags": None,
},
}
with (
patch( # test-quality-ok: callback reads proxy_logging_obj from the module; no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj",
) as mock_proxy_logging,
patch( # test-quality-ok: async-complete streams must settle, not release
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
patch( # test-quality-ok: settle is a proxy-internal reservation call, not an HTTP boundary
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
new_callable=AsyncMock,
) as mock_reconcile_budget_reservation,
patch( # test-quality-ok: unpriced settle must not write a spend-log row
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database,
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
mock_release_budget_reservation.assert_not_awaited()
mock_reconcile_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
actual_cost=0.5,
)
mock_update_database.assert_not_called()
mock_proxy_logging.failed_tracking_alert.assert_called()
assert budget_reservation["reserved_cost"] == 0.5
assert budget_reservation.get("finalized") is not True
@pytest.mark.asyncio
async def test_track_cost_callback_invalidates_reservation_when_settle_fails():
"""A failed settle must not leave the hold pinning later traffic until TTL."""
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation)
kwargs = {
"model": "gpt-4",
"call_type": "acompletion",
"litellm_params": {
"metadata": {
"user_api_key_auth": user_api_key_auth,
},
},
"standard_logging_object": {
"response_cost": None,
"response_cost_failure_debug_info": "missing custom price",
"request_tags": None,
},
"stream": False,
}
with (
patch( # test-quality-ok: callback reads proxy_logging_obj from the module; no injection seam
"litellm.proxy.proxy_server.proxy_logging_obj",
) as mock_proxy_logging,
patch( # test-quality-ok: a failed settle must not refund through release_budget_reservation
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation,
patch( # test-quality-ok: force reconcile to fail so the invalidate path can be observed
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
new_callable=AsyncMock,
side_effect=RuntimeError("redis down"),
),
patch( # test-quality-ok: invalidate is the only way to unpin counters after settle fails
"litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters",
new_callable=AsyncMock,
) as mock_invalidate_budget_reservation_counters,
patch( # test-quality-ok: failed settle must not write a spend-log row
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database,
):
mock_proxy_logging.failed_tracking_alert = AsyncMock()
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
mock_release_budget_reservation.assert_not_awaited()
mock_invalidate_budget_reservation_counters.assert_awaited_once()
settled = mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"]
assert settled["reserved_cost"] == 0.5
assert settled["finalized"] is True
mock_update_database.assert_not_called()
mock_proxy_logging.failed_tracking_alert.assert_called()
@pytest.mark.asyncio
async def test_track_cost_callback_releases_budget_reservation_for_non_model_calls():
logger = _ProxyDBLogger()
budget_reservation = {"reserved_cost": 0.5, "entries": []}
user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation)
kwargs = {
"call_type": "health",
"litellm_params": {
"metadata": {
"user_api_key_auth": user_api_key_auth,
},
},
"stream": False,
}
with patch( # test-quality-ok: health checks have no cost row; release is the observable contract
"litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation",
new_callable=AsyncMock,
) as mock_release_budget_reservation:
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
mock_release_budget_reservation.assert_awaited_once_with(
budget_reservation=budget_reservation,
)
# Health checks refund the hold; they must not stamp it settled.
assert budget_reservation.get("finalized") is not True
assert budget_reservation["reserved_cost"] == 0.5
def test_get_budget_reservation_from_metadata_handles_dict_auth_object():