mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(prometheus): skip customer budget cache read when end_user tracking is off
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
e7d537442d
commit
81d57cbad7
2 changed files with 24 additions and 1 deletions
|
|
@ -4204,7 +4204,7 @@ class PrometheusLogger(CustomLogger):
|
|||
end_user_id: str | None,
|
||||
response_cost: float,
|
||||
):
|
||||
if self._customer_budget_gauges_are_noop():
|
||||
if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled():
|
||||
return
|
||||
|
||||
if not end_user_id:
|
||||
|
|
|
|||
|
|
@ -1150,6 +1150,29 @@ async def test_set_customer_budget_metrics_after_api_request_without_end_user_is
|
|||
assert prometheus_logger.litellm_remaining_customer_budget_metric._metrics == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_customer_budget_metrics_after_api_request_skips_cache_when_end_user_tracking_off(
|
||||
prometheus_logger, monkeypatch
|
||||
):
|
||||
import sys
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", False)
|
||||
monkeypatch.setattr(litellm, "disable_end_user_cost_tracking", False)
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock()
|
||||
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
|
||||
await prometheus_logger._set_customer_budget_metrics_after_api_request(
|
||||
end_user_id="cust-off",
|
||||
response_cost=1.0,
|
||||
)
|
||||
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache.assert_not_awaited()
|
||||
assert prometheus_logger.litellm_remaining_customer_budget_metric._metrics == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_customer_budget_metrics_emits_gauges_for_budgeted_customers(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue