scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 128

fix(prometheus): read the committed org spend on failed requests without the auth object

Browse source
This commit is contained in:
mateo-berri 2026-10-30 10:47:51 -07:00
parent aec15b56b6
commit 0ce16cd68d
2 changed files with 24 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/integrations/prometheus.py
View file

@ -377,6 +377,12 @@ def _customer_spend_after_request(cached_spend: float, committed_spend: float |
return max(committed_spend + response_cost, cached_spend)
def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys:
if org is None or org.max_budget is None:
return EntitySpendCounterKeys()
return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None)
async def _committed_spend_now(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend:
if not counter_keys.tracked():
return CommittedBudgetSpend()
@ -2862,16 +2868,16 @@ class PrometheusLogger(CustomLogger):
get_litellm_metadata_from_kwargs(kwargs)
)
failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias")
failure_at_auth: Final = committed_budget_spend_from_metadata(failure_metadata)
failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata)
failure_committed: Final = _committed_spend_before_request(
at_auth=failure_at_auth,
now=await _committed_spend_now(EntitySpendCounterKeys(org=failure_at_auth.counter_keys.org)),
at_auth=committed_budget_spend_from_metadata(failure_metadata),
now=await _committed_spend_now(_budgeted_org_counter(user_api_key_org_id, failure_org)),
response_cost=0,
)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
response_cost=0,
carried=OrgBudgetSnapshot.from_metadata(failure_metadata),
carried=failure_org,
org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None,
committed_spend=failure_committed.org,
)

14
tests/unit/integrations/test_prometheus_carried_budget_state.py
View file

@ -434,6 +434,20 @@ async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_commit
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0))
@pytest.mark.asyncio
async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend(
prometheus_logger, getters, spend_counter_cache
):
_stub_gauges(prometheus_logger)
_store_committed(spend_counter_cache, key=3.0, team=30.0, user=8.0, org=65.0)
metadata = _request_metadata(_token_committed_at_auth())
del metadata["user_api_key_auth"]
await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT)
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0))
@pytest.mark.asyncio
async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters):
_stub_gauges(prometheus_logger)