scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

fix(prometheus): refresh default-budget customers and honor independent customer gauges

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-16 19:08:42 +00:00
parent 8c89cff0e0
commit 2821ba9667
2 changed files with 152 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/integrations/prometheus.py
View file

@ -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,

109
tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
View file

@ -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