From 2821ba96670c5ee1f03ca5d585535a05d684a9a3 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:08:42 +0000 Subject: [PATCH] fix(prometheus): refresh default-budget customers and honor independent customer gauges Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/prometheus.py | 56 ++++++--- .../test_prometheus_user_team_metrics.py | 109 ++++++++++++++++++ 2 files changed, 152 insertions(+), 13 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6a728904ba8..4f67cb6b0e2 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -44,6 +44,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.table_repositories import EndUserRepository from litellm.repositories.team_repository import TeamRepository @@ -68,10 +69,15 @@ from litellm.types.utils import ( if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler - from prisma.types import LiteLLM_EndUserTableInclude, LiteLLM_EndUserTableOrderByInput + from prisma.types import ( + LiteLLM_BudgetTableWhereUniqueInput, + LiteLLM_EndUserTableInclude, + LiteLLM_EndUserTableOrderByInput, + ) from prometheus_client import Gauge from prometheus_client.metrics import MetricWrapperBase + from litellm.proxy.utils import PrismaClient from litellm.router import Router else: AsyncIOScheduler = Any @@ -1983,7 +1989,7 @@ class PrometheusLogger(CustomLogger): and isinstance(self.litellm_remaining_api_key_budget_metric, NoOpMetric) and isinstance(self.litellm_remaining_user_budget_metric, NoOpMetric) and isinstance(self.litellm_remaining_org_budget_metric, NoOpMetric) - and isinstance(self.litellm_remaining_customer_budget_metric, NoOpMetric) + and self._customer_budget_gauges_are_noop() ): return @@ -3794,15 +3800,17 @@ class PrometheusLogger(CustomLogger): verbose_logger.debug("Prometheus: skipping customer metrics initialization, DB not initialized") return - if isinstance(self.litellm_remaining_customer_budget_metric, NoOpMetric): + if self._customer_budget_gauges_are_noop(): return if not _customer_budget_metrics_enabled(): verbose_logger.debug("Prometheus: skipping customer metrics initialization, end_user tracking disabled") return + default_budget: Final = await self._get_default_customer_budget(prisma_client) customers_table: Final = EndUserRepository(prisma_client).table - budgeted_customers: Final[_BudgetedCustomerFilter] = {"budget_id": {"not": None}} + with_persisted_budget: Final[_BudgetedCustomerFilter] = {"budget_id": {"not": None}} + budgeted_customers: Final = None if default_budget is not None else with_persisted_budget by_user_id: Final[LiteLLM_EndUserTableOrderByInput] = {"user_id": "asc"} with_budget: Final[LiteLLM_EndUserTableInclude] = {"litellm_budget_table": True} @@ -3815,12 +3823,16 @@ class PrometheusLogger(CustomLogger): order=by_user_id, include=with_budget, ) - total_count: Final = await customers_table.count(where=budgeted_customers) + total_count: Final = await customers_table.count(where=budgeted_customers) if page == 1 else None return customers, total_count + async def set_customer_metrics(customers: Sequence[_CustomerBudgetRow]) -> None: + for customer in customers: + self._set_customer_budget_metrics_from_row(customer, default_budget=default_budget) + await self._initialize_budget_metrics( data_fetch_function=fetch_customers, - set_metrics_function=self._set_customer_list_budget_metrics, + set_metrics_function=set_customer_metrics, data_type="customers", ) @@ -3924,12 +3936,12 @@ class PrometheusLogger(CustomLogger): budget_reset_at=(getattr(budget_table, "budget_reset_at", None) if budget_table else None), ) - async def _set_customer_list_budget_metrics(self, customers: Sequence[_CustomerBudgetRow]): - for customer in customers: - self._set_customer_budget_metrics_from_row(customer) - - def _set_customer_budget_metrics_from_row(self, customer: _CustomerBudgetRow): - budget_table: Final = customer.litellm_budget_table + def _set_customer_budget_metrics_from_row( + self, customer: _CustomerBudgetRow, default_budget: _JoinedBudgetRow | None + ): + budget_table: Final = ( + customer.litellm_budget_table if customer.litellm_budget_table is not None else default_budget + ) self._set_customer_budget_metrics( end_user_id=customer.user_id, spend=customer.spend, @@ -4191,7 +4203,7 @@ class PrometheusLogger(CustomLogger): end_user_id: str | None, response_cost: float, ): - if isinstance(self.litellm_remaining_customer_budget_metric, NoOpMetric): + if self._customer_budget_gauges_are_noop(): return if not end_user_id: @@ -4224,6 +4236,24 @@ class PrometheusLogger(CustomLogger): budget_reset_at=None, ) + async def _get_default_customer_budget(self, prisma_client: PrismaClient) -> _JoinedBudgetRow | None: + default_budget_id: Final = litellm.max_end_user_budget_id + if default_budget_id is None: + return None + default_budget_key: Final[LiteLLM_BudgetTableWhereUniqueInput] = {"budget_id": default_budget_id} + try: + return await BudgetRepository(prisma_client).table.find_unique(where=default_budget_key) + except Exception as e: + verbose_logger.debug("[Non-Blocking] Prometheus: Error getting default customer budget: %s", e) + return None + + def _customer_budget_gauges_are_noop(self) -> bool: + return ( + isinstance(self.litellm_remaining_customer_budget_metric, NoOpMetric) + and isinstance(self.litellm_customer_max_budget_metric, NoOpMetric) + and isinstance(self.litellm_customer_budget_remaining_hours_metric, NoOpMetric) + ) + def _set_customer_budget_metrics( self, end_user_id: str, diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index dee710ba1b1..077ea305053 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -929,6 +929,7 @@ def customer_metrics_enabled(monkeypatch): monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", True) monkeypatch.setattr(litellm, "disable_end_user_cost_tracking", False) + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) def _customer_sample(metric_name: str, end_user_id: str): @@ -1152,6 +1153,114 @@ async def test_initialize_remaining_budget_metrics_includes_customers(prometheus assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-startup") == pytest.approx(20.0) +@pytest.mark.asyncio +async def test_initialize_customer_budget_metrics_counts_once_across_pages(prometheus_logger, customer_metrics_enabled): + import sys + + pages = [ + [_mock_customer_row(f"cust-{i}", 1.0, 10.0, None) for i in range(50)], + [_mock_customer_row(f"cust-{i}", 1.0, 10.0, None) for i in range(50, 100)], + [_mock_customer_row("cust-100", 1.0, 10.0, None)], + ] + find_many = AsyncMock(side_effect=pages) + count = AsyncMock(return_value=101) + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = find_many + mock_prisma.db.litellm_endusertable.count = count + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = mock_prisma + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._initialize_customer_budget_metrics() + + assert find_many.await_count == 3 + count.assert_awaited_once() + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-100") == pytest.approx(9.0) + + +@pytest.mark.asyncio +async def test_initialize_customer_budget_metrics_applies_default_budget_to_unbudgeted_customers( + prometheus_logger, customer_metrics_enabled, monkeypatch +): + import sys + + import litellm + + monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-customer-budget") + reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc) + default_budget = MagicMock() + default_budget.max_budget = 10.0 + default_budget.budget_reset_at = reset_at + explicit_row = _mock_customer_row("cust-explicit", 5.0, 100.0, None) + default_row = _mock_customer_row("cust-default", 2.0, None, None) + default_row.litellm_budget_table = None + find_many = AsyncMock(return_value=[explicit_row, default_row]) + find_unique = AsyncMock(return_value=default_budget) + mock_prisma = MagicMock() + mock_prisma.db.litellm_endusertable.find_many = find_many + mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=2) + mock_prisma.db.litellm_budgettable.find_unique = find_unique + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = mock_prisma + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._initialize_customer_budget_metrics() + + assert find_unique.await_args.kwargs["where"] == {"budget_id": "default-customer-budget"} + assert find_many.await_args.kwargs["where"] is None + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-explicit") == pytest.approx(95.0) + assert _customer_sample("litellm_customer_max_budget_metric", "cust-explicit") == pytest.approx(100.0) + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-default") == pytest.approx(8.0) + assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(10.0) + assert _customer_sample("litellm_customer_budget_remaining_hours_metric", "cust-default") > 0 + + +@pytest.mark.asyncio +async def test_customer_max_budget_gauge_emitted_when_only_it_is_configured(customer_metrics_enabled, monkeypatch): + import sys + + import litellm + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.types.integrations.prometheus import NoOpMetric + + monkeypatch.setattr( + litellm, + "prometheus_metrics_config", + [{"group": "customer-max-only", "metrics": ["litellm_customer_max_budget_metric"]}], + ) + logger = PrometheusLogger() + assert isinstance(logger.litellm_remaining_customer_budget_metric, NoOpMetric) + assert not isinstance(logger.litellm_customer_max_budget_metric, NoOpMetric) + + end_user = LiteLLM_EndUserTable( + user_id="cust-max-only", + blocked=False, + spend=1.0, + budget_id="budget-1", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget-1", max_budget=40.0), + ) + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = MagicMock() + mock_proxy_server.user_api_key_cache = MagicMock() + + with ( + patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}), + patch("litellm.proxy.auth.auth_checks.get_end_user_object", AsyncMock(return_value=end_user)), # test-quality-ok: [TQ008] logger resolves customers through the proxy auth lookup, no injection seam + ): + await logger._increment_remaining_budget_metrics( + user_api_team=None, + user_api_team_alias=None, + user_api_key=None, + user_api_key_alias=None, + litellm_params={"metadata": {}}, + response_cost=1.0, + end_user_id="cust-max-only", + ) + + assert _customer_sample("litellm_customer_max_budget_metric", "cust-max-only") == pytest.approx(40.0) + + def test_default_latency_buckets(prometheus_logger): """PrometheusLogger uses the new reduced default latency buckets.""" from litellm.types.integrations.prometheus import LATENCY_BUCKETS