mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge pull request #41472 from BerriAI/litellm_customer_budget_prometheus_metrics
feat(prometheus): add customer (end_user) budget gauges
This commit is contained in:
Yassin Kortam
GitHub
commit
4bf04b22b0
4 changed files with 710 additions and 13 deletions
|
|
@ -13,6 +13,7 @@ from datetime import datetime, timedelta
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
|
|
@ -36,6 +37,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
from litellm.litellm_core_utils.service_tier_utils import (
|
||||
get_service_tier_from_standard_logging_payload,
|
||||
)
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_DeletedVerificationToken,
|
||||
LiteLLM_TeamTable,
|
||||
|
|
@ -43,7 +45,9 @@ 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
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
|
@ -66,13 +70,26 @@ from litellm.types.utils import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
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
|
||||
|
||||
_IsNotNull = TypedDict("_IsNotNull", {"not": ReadOnly[None]})
|
||||
|
||||
|
||||
class _BudgetedCustomerFilter(TypedDict):
|
||||
budget_id: ReadOnly[_IsNotNull]
|
||||
|
||||
|
||||
_BudgetRowT: Final = TypeVar("_BudgetRowT")
|
||||
_TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
|
||||
|
||||
|
|
@ -116,8 +133,8 @@ def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrisma
|
|||
)
|
||||
|
||||
|
||||
class _OrgBudgetRow(Protocol):
|
||||
"""The budget columns joined onto an organization row."""
|
||||
class _JoinedBudgetRow(Protocol):
|
||||
"""The budget columns joined onto an organization or customer row."""
|
||||
|
||||
@property
|
||||
def max_budget(self) -> float | None: ...
|
||||
|
|
@ -126,6 +143,23 @@ class _OrgBudgetRow(Protocol):
|
|||
def budget_reset_at(self) -> datetime | None: ...
|
||||
|
||||
|
||||
class _CustomerBudgetRow(Protocol):
|
||||
"""The columns of a customer (end user) row that budget gauges read."""
|
||||
|
||||
@property
|
||||
def user_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def spend(self) -> float: ...
|
||||
|
||||
@property
|
||||
def litellm_budget_table(self) -> _JoinedBudgetRow | None: ...
|
||||
|
||||
|
||||
def _customer_budget_metrics_enabled() -> bool:
|
||||
return litellm.enable_end_user_cost_tracking_prometheus_only is True and not litellm.disable_end_user_cost_tracking
|
||||
|
||||
|
||||
class _ExcludedLabelMetric:
|
||||
"""Proxies a prometheus metric whose declared ``labelnames`` had globally
|
||||
excluded labels removed, dropping those labels from every ``labels(...)``
|
||||
|
|
@ -471,6 +505,24 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=self.get_labels_for_metric("litellm_user_budget_remaining_hours_metric"),
|
||||
)
|
||||
|
||||
self.litellm_remaining_customer_budget_metric = self._gauge_factory(
|
||||
"litellm_remaining_customer_budget_metric",
|
||||
"Remaining budget for customer (end user)",
|
||||
labelnames=self.get_labels_for_metric("litellm_remaining_customer_budget_metric"),
|
||||
)
|
||||
|
||||
self.litellm_customer_max_budget_metric = self._gauge_factory(
|
||||
"litellm_customer_max_budget_metric",
|
||||
"Maximum budget set for customer (end user)",
|
||||
labelnames=self.get_labels_for_metric("litellm_customer_max_budget_metric"),
|
||||
)
|
||||
|
||||
self.litellm_customer_budget_remaining_hours_metric = self._gauge_factory(
|
||||
"litellm_customer_budget_remaining_hours_metric",
|
||||
"Remaining hours for customer (end user) budget to be reset",
|
||||
labelnames=self.get_labels_for_metric("litellm_customer_budget_remaining_hours_metric"),
|
||||
)
|
||||
|
||||
########################################
|
||||
# LiteLLM Virtual API KEY metrics
|
||||
########################################
|
||||
|
|
@ -1334,7 +1386,7 @@ class PrometheusLogger(CustomLogger):
|
|||
self,
|
||||
metric: Any,
|
||||
metric_name: DEFINED_PROMETHEUS_METRICS,
|
||||
labels: dict[str, str | None],
|
||||
labels: Mapping[str, str | None],
|
||||
) -> None:
|
||||
"""
|
||||
Cap the cardinality of metrics that include the ``end_user`` label.
|
||||
|
|
@ -1501,6 +1553,7 @@ class PrometheusLogger(CustomLogger):
|
|||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
user_api_key_org_id=user_api_key_org_id,
|
||||
end_user_id=end_user_id,
|
||||
)
|
||||
|
||||
# set proxy virtual key rpm/tpm metrics
|
||||
|
|
@ -1930,12 +1983,14 @@ class PrometheusLogger(CustomLogger):
|
|||
response_cost: float,
|
||||
user_id: str | None = None,
|
||||
user_api_key_org_id: str | None = None,
|
||||
end_user_id: str | None = None,
|
||||
):
|
||||
if (
|
||||
isinstance(self.litellm_remaining_team_budget_metric, NoOpMetric)
|
||||
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 self._customer_budget_gauges_are_noop()
|
||||
):
|
||||
return
|
||||
|
||||
|
|
@ -1990,6 +2045,10 @@ class PrometheusLogger(CustomLogger):
|
|||
carried=OrgBudgetSnapshot.from_metadata(_metadata),
|
||||
org_alias=_org_alias if isinstance(_org_alias, str) else None,
|
||||
),
|
||||
self._set_customer_budget_metrics_after_api_request(
|
||||
end_user_id=end_user_id,
|
||||
response_cost=response_cost,
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
try:
|
||||
|
|
@ -2006,7 +2065,7 @@ class PrometheusLogger(CustomLogger):
|
|||
if isinstance(r, Exception):
|
||||
verbose_logger.debug(
|
||||
"[Non-Blocking] Prometheus: Budget metric lookup %s failed: %s",
|
||||
["key", "team", "user", "org"][i],
|
||||
("key", "team", "user", "org", "customer")[i],
|
||||
r,
|
||||
)
|
||||
|
||||
|
|
@ -3574,9 +3633,9 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
async def _initialize_budget_metrics(
|
||||
self,
|
||||
data_fetch_function: Callable[..., Awaitable[tuple[list[_BudgetRowT], int | None]]],
|
||||
set_metrics_function: Callable[[list[_BudgetRowT]], Awaitable[None]],
|
||||
data_type: Literal["teams", "keys", "users", "orgs"],
|
||||
data_fetch_function: Callable[..., Awaitable[tuple[Sequence[_BudgetRowT], int | None]]],
|
||||
set_metrics_function: Callable[[Sequence[_BudgetRowT]], Awaitable[None]],
|
||||
data_type: Literal["teams", "keys", "users", "orgs", "customers"],
|
||||
):
|
||||
"""
|
||||
Generic method to initialize budget metrics for teams or API keys.
|
||||
|
|
@ -3735,6 +3794,49 @@ class PrometheusLogger(CustomLogger):
|
|||
data_type="orgs",
|
||||
)
|
||||
|
||||
async def _initialize_customer_budget_metrics(self):
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_logger.debug("Prometheus: skipping customer metrics initialization, DB not initialized")
|
||||
return
|
||||
|
||||
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
|
||||
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}
|
||||
|
||||
async def fetch_customers(page_size: int, page: int) -> tuple[Sequence[_CustomerBudgetRow], int | None]:
|
||||
skip: Final = (page - 1) * page_size
|
||||
customers: Final = await customers_table.find_many(
|
||||
skip=skip,
|
||||
take=page_size,
|
||||
where=budgeted_customers,
|
||||
order=by_user_id,
|
||||
include=with_budget,
|
||||
)
|
||||
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=set_customer_metrics,
|
||||
data_type="customers",
|
||||
)
|
||||
|
||||
async def initialize_remaining_budget_metrics(self):
|
||||
"""
|
||||
Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies.
|
||||
|
|
@ -3765,11 +3867,12 @@ class PrometheusLogger(CustomLogger):
|
|||
"""
|
||||
Helper to initialize remaining budget metrics for all teams, API keys, and users.
|
||||
"""
|
||||
verbose_logger.debug("Emitting key, team, user, org budget metrics....")
|
||||
verbose_logger.debug("Emitting key, team, user, org, customer budget metrics....")
|
||||
await self._initialize_team_budget_metrics()
|
||||
await self._initialize_api_key_budget_metrics()
|
||||
await self._initialize_user_budget_metrics()
|
||||
await self._initialize_org_budget_metrics()
|
||||
await self._initialize_customer_budget_metrics()
|
||||
await self._initialize_user_and_team_count_metrics()
|
||||
|
||||
async def _initialize_user_and_team_count_metrics(self):
|
||||
|
|
@ -3805,27 +3908,27 @@ class PrometheusLogger(CustomLogger):
|
|||
verbose_logger.exception("Error initializing user/team count metrics: %s", e)
|
||||
|
||||
async def _set_key_list_budget_metrics(
|
||||
self, keys: list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]
|
||||
self, keys: Sequence[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]
|
||||
) -> None:
|
||||
"""Helper function to set budget metrics for a list of keys"""
|
||||
for key in keys:
|
||||
if isinstance(key, UserAPIKeyAuth):
|
||||
self._set_key_budget_metrics(key)
|
||||
|
||||
async def _set_team_list_budget_metrics(self, teams: list[LiteLLM_TeamTable]):
|
||||
async def _set_team_list_budget_metrics(self, teams: Sequence[LiteLLM_TeamTable]):
|
||||
"""Helper function to set budget metrics for a list of teams"""
|
||||
for team in teams:
|
||||
self._set_team_budget_metrics(team)
|
||||
|
||||
async def _set_user_list_budget_metrics(self, users: list[LiteLLM_UserTable]):
|
||||
async def _set_user_list_budget_metrics(self, users: Sequence[LiteLLM_UserTable]):
|
||||
"""Helper function to set budget metrics for a list of users"""
|
||||
for user in users:
|
||||
self._set_user_budget_metrics(user)
|
||||
|
||||
async def _set_org_list_budget_metrics(self, orgs: list):
|
||||
async def _set_org_list_budget_metrics(self, orgs: Sequence):
|
||||
"""Helper function to set budget metrics for a list of orgs"""
|
||||
for org in orgs:
|
||||
budget_table: _OrgBudgetRow | None = getattr(org, "litellm_budget_table", None)
|
||||
budget_table: _JoinedBudgetRow | None = getattr(org, "litellm_budget_table", None)
|
||||
self._set_org_budget_metrics(
|
||||
org_id=org.organization_id or "",
|
||||
org_alias=org.organization_alias or "",
|
||||
|
|
@ -3834,6 +3937,19 @@ class PrometheusLogger(CustomLogger):
|
|||
budget_reset_at=(getattr(budget_table, "budget_reset_at", None) if budget_table else None),
|
||||
)
|
||||
|
||||
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,
|
||||
max_budget=budget_table.max_budget if budget_table is not None else None,
|
||||
budget_reset_at=budget_table.budget_reset_at if budget_table is not None else None,
|
||||
)
|
||||
|
||||
async def _set_team_budget_metrics_after_api_request(
|
||||
self,
|
||||
user_api_team: str | None,
|
||||
|
|
@ -4083,6 +4199,98 @@ class PrometheusLogger(CustomLogger):
|
|||
self._get_remaining_hours_for_budget_reset(budget_reset_at=budget_reset_at)
|
||||
)
|
||||
|
||||
async def _set_customer_budget_metrics_after_api_request(
|
||||
self,
|
||||
end_user_id: str | None,
|
||||
response_cost: float,
|
||||
):
|
||||
if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled():
|
||||
return
|
||||
|
||||
if not end_user_id:
|
||||
return
|
||||
|
||||
from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
try:
|
||||
cached_customer: Final = await user_api_key_cache.async_get_cache(
|
||||
key=end_user_cache_key(end_user_id),
|
||||
model_type=LiteLLM_EndUserTable,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("[Non-Blocking] Prometheus: Error getting customer info: %s", e)
|
||||
return
|
||||
|
||||
if cached_customer is None:
|
||||
return
|
||||
|
||||
budget_table: Final = cached_customer.litellm_budget_table
|
||||
self._set_customer_budget_metrics(
|
||||
end_user_id=end_user_id,
|
||||
spend=cached_customer.spend + response_cost,
|
||||
max_budget=budget_table.max_budget if budget_table is not None else None,
|
||||
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,
|
||||
spend: float,
|
||||
max_budget: float | None,
|
||||
budget_reset_at: datetime | None,
|
||||
):
|
||||
_labels: Final[dict[str, str | None]] = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_customer_budget_metric"),
|
||||
enum_values=UserAPIKeyLabelValues(end_user=end_user_id),
|
||||
)
|
||||
if _labels.get(UserAPIKeyLabelNames.END_USER.value) is None:
|
||||
return
|
||||
|
||||
self.litellm_remaining_customer_budget_metric.labels(**_labels).set(
|
||||
self._safe_get_remaining_budget(
|
||||
max_budget=max_budget,
|
||||
spend=spend,
|
||||
)
|
||||
)
|
||||
self._track_end_user_metric_series(
|
||||
self.litellm_remaining_customer_budget_metric, "litellm_remaining_customer_budget_metric", _labels
|
||||
)
|
||||
|
||||
if max_budget is not None:
|
||||
self.litellm_customer_max_budget_metric.labels(**_labels).set(max_budget)
|
||||
self._track_end_user_metric_series(
|
||||
self.litellm_customer_max_budget_metric, "litellm_customer_max_budget_metric", _labels
|
||||
)
|
||||
|
||||
if budget_reset_at is not None:
|
||||
self.litellm_customer_budget_remaining_hours_metric.labels(**_labels).set(
|
||||
self._get_remaining_hours_for_budget_reset(budget_reset_at=budget_reset_at)
|
||||
)
|
||||
self._track_end_user_metric_series(
|
||||
self.litellm_customer_budget_remaining_hours_metric,
|
||||
"litellm_customer_budget_remaining_hours_metric",
|
||||
_labels,
|
||||
)
|
||||
|
||||
def _set_key_budget_metrics(self, user_api_key_dict: UserAPIKeyAuth):
|
||||
"""
|
||||
Set virtual key budget metrics
|
||||
|
|
|
|||
|
|
@ -262,6 +262,9 @@ DEFINED_PROMETHEUS_METRICS = Literal[
|
|||
"litellm_remaining_user_budget_metric",
|
||||
"litellm_user_max_budget_metric",
|
||||
"litellm_user_budget_remaining_hours_metric",
|
||||
"litellm_remaining_customer_budget_metric",
|
||||
"litellm_customer_max_budget_metric",
|
||||
"litellm_customer_budget_remaining_hours_metric",
|
||||
"litellm_deployment_state",
|
||||
"litellm_deployment_failure_responses",
|
||||
"litellm_deployment_total_requests",
|
||||
|
|
@ -733,6 +736,12 @@ class PrometheusMetricLabels:
|
|||
|
||||
litellm_user_budget_remaining_hours_metric = litellm_remaining_user_budget_metric
|
||||
|
||||
litellm_remaining_customer_budget_metric = (UserAPIKeyLabelNames.END_USER.value,)
|
||||
|
||||
litellm_customer_max_budget_metric = litellm_remaining_customer_budget_metric
|
||||
|
||||
litellm_customer_budget_remaining_hours_metric = litellm_remaining_customer_budget_metric
|
||||
|
||||
litellm_remaining_api_key_requests_for_model = [
|
||||
UserAPIKeyLabelNames.API_KEY_HASH.value,
|
||||
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from time import monotonic
|
||||
|
||||
import pytest
|
||||
|
|
@ -179,3 +180,50 @@ def test_prometheus_end_user_not_tracked_by_default():
|
|||
|
||||
prometheus_labels = prometheus_label_factory(labels, label_values)
|
||||
assert prometheus_labels["end_user"] is None
|
||||
|
||||
|
||||
def test_prometheus_customer_budget_series_are_capped_per_metric(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", True)
|
||||
monkeypatch.setattr(litellm, "prometheus_end_user_metrics_max_series_per_metric", 2)
|
||||
monkeypatch.setattr(litellm, "prometheus_end_user_metrics_ttl_seconds", None)
|
||||
logger = PrometheusLogger()
|
||||
|
||||
for index in range(5):
|
||||
logger._set_customer_budget_metrics(
|
||||
end_user_id=f"customer-{index}",
|
||||
spend=1.0,
|
||||
max_budget=10.0,
|
||||
budget_reset_at=None,
|
||||
)
|
||||
|
||||
assert set(logger.litellm_remaining_customer_budget_metric._metrics) == {("customer-3",), ("customer-4",)}
|
||||
assert set(logger.litellm_customer_max_budget_metric._metrics) == {("customer-3",), ("customer-4",)}
|
||||
|
||||
|
||||
def test_prometheus_customer_budget_series_expire_by_ttl(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", True)
|
||||
monkeypatch.setattr(litellm, "prometheus_end_user_metrics_max_series_per_metric", None)
|
||||
monkeypatch.setattr(litellm, "prometheus_end_user_metrics_ttl_seconds", 10.0)
|
||||
monkeypatch.setattr(litellm, "prometheus_end_user_metrics_cleanup_interval_seconds", 0.0)
|
||||
logger = PrometheusLogger()
|
||||
|
||||
current_time = [monotonic()]
|
||||
monkeypatch.setattr(bounded_prometheus_series_tracker.time, "monotonic", lambda: current_time[0])
|
||||
logger._set_customer_budget_metrics(
|
||||
end_user_id="customer-with-removed-budget",
|
||||
spend=1.0,
|
||||
max_budget=10.0,
|
||||
budget_reset_at=datetime.now(timezone.utc) + timedelta(hours=1),
|
||||
)
|
||||
|
||||
current_time[0] += 11.0
|
||||
logger._set_customer_budget_metrics(
|
||||
end_user_id="customer-still-budgeted",
|
||||
spend=1.0,
|
||||
max_budget=10.0,
|
||||
budget_reset_at=datetime.now(timezone.utc) + timedelta(hours=1),
|
||||
)
|
||||
|
||||
assert set(logger.litellm_remaining_customer_budget_metric._metrics) == {("customer-still-budgeted",)}
|
||||
assert set(logger.litellm_customer_max_budget_metric._metrics) == {("customer-still-budgeted",)}
|
||||
assert set(logger.litellm_customer_budget_remaining_hours_metric._metrics) == {("customer-still-budgeted",)}
|
||||
|
|
|
|||
|
|
@ -923,6 +923,438 @@ async def test_initialize_org_budget_metrics(prometheus_logger):
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def customer_metrics_enabled(monkeypatch):
|
||||
import litellm
|
||||
|
||||
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):
|
||||
return REGISTRY.get_sample_value(metric_name, {"end_user": end_user_id})
|
||||
|
||||
|
||||
def _mock_customer_row(user_id: str, spend: float, max_budget: float | None, budget_reset_at):
|
||||
budget_mock = MagicMock()
|
||||
budget_mock.max_budget = max_budget
|
||||
budget_mock.budget_reset_at = budget_reset_at
|
||||
row = MagicMock()
|
||||
row.user_id = user_id
|
||||
row.spend = spend
|
||||
row.litellm_budget_table = budget_mock
|
||||
return row
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spend, max_budget, expected_remaining",
|
||||
[(125.0, 500.0, 375.0), (500.0, 500.0, 0.0), (0.0, 500.0, 500.0)],
|
||||
)
|
||||
def test_set_customer_budget_metrics_emits_remaining_and_max_budget(
|
||||
prometheus_logger, customer_metrics_enabled, spend, max_budget, expected_remaining
|
||||
):
|
||||
prometheus_logger._set_customer_budget_metrics(
|
||||
end_user_id="cust-1",
|
||||
spend=spend,
|
||||
max_budget=max_budget,
|
||||
budget_reset_at=None,
|
||||
)
|
||||
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-1") == pytest.approx(
|
||||
expected_remaining
|
||||
)
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-1") == pytest.approx(max_budget)
|
||||
assert _customer_sample("litellm_customer_budget_remaining_hours_metric", "cust-1") is None
|
||||
|
||||
|
||||
def test_set_customer_budget_metrics_remaining_hours(prometheus_logger, customer_metrics_enabled):
|
||||
reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc)
|
||||
prometheus_logger._set_customer_budget_metrics(
|
||||
end_user_id="cust-1",
|
||||
spend=1.0,
|
||||
max_budget=10.0,
|
||||
budget_reset_at=reset_at,
|
||||
)
|
||||
|
||||
expected_hours = (reset_at - datetime.now(timezone.utc)).total_seconds() / 3600
|
||||
assert _customer_sample("litellm_customer_budget_remaining_hours_metric", "cust-1") == pytest.approx(
|
||||
expected_hours, abs=0.1
|
||||
)
|
||||
|
||||
|
||||
def test_set_customer_budget_metrics_not_emitted_when_end_user_tracking_off(prometheus_logger, monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", False)
|
||||
monkeypatch.setattr(litellm, "disable_end_user_cost_tracking", False)
|
||||
|
||||
prometheus_logger._set_customer_budget_metrics(
|
||||
end_user_id="cust-off",
|
||||
spend=1.0,
|
||||
max_budget=10.0,
|
||||
budget_reset_at=datetime(2099, 1, 1, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
assert prometheus_logger.litellm_remaining_customer_budget_metric._metrics == {}
|
||||
assert prometheus_logger.litellm_customer_max_budget_metric._metrics == {}
|
||||
assert prometheus_logger.litellm_customer_budget_remaining_hours_metric._metrics == {}
|
||||
|
||||
|
||||
def test_set_customer_budget_metrics_without_budget_only_emits_remaining(prometheus_logger, customer_metrics_enabled):
|
||||
prometheus_logger._set_customer_budget_metrics(
|
||||
end_user_id="cust-free",
|
||||
spend=3.0,
|
||||
max_budget=None,
|
||||
budget_reset_at=None,
|
||||
)
|
||||
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-free") == float("inf")
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-free") is None
|
||||
assert _customer_sample("litellm_customer_budget_remaining_hours_metric", "cust-free") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_remaining_budget_metrics_emits_customer_gauges_from_cached_end_user(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
):
|
||||
import sys
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
|
||||
end_user = LiteLLM_EndUserTable(
|
||||
user_id="cust-req",
|
||||
blocked=False,
|
||||
spend=300.0,
|
||||
budget_id="budget-1",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget-1", max_budget=1000.0),
|
||||
)
|
||||
get_end_user_object = AsyncMock()
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.prisma_client = None
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user)
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}),
|
||||
patch("litellm.proxy.auth.auth_checks.get_end_user_object", get_end_user_object), # test-quality-ok: [TQ008] assert the request path never reaches the DB-backed auth lookup
|
||||
):
|
||||
await prometheus_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=50.0,
|
||||
end_user_id="cust-req",
|
||||
)
|
||||
|
||||
get_end_user_object.assert_not_awaited()
|
||||
cache_read = mock_proxy_server.user_api_key_cache.async_get_cache
|
||||
cache_read.assert_awaited_once()
|
||||
assert cache_read.await_args.kwargs["key"] == "end_user_id:cust-req"
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-req") == pytest.approx(650.0)
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-req") == pytest.approx(1000.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_customer_budget_metrics_after_api_request_uses_cached_default_budget(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
):
|
||||
import sys
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
|
||||
end_user = LiteLLM_EndUserTable(
|
||||
user_id="cust-default",
|
||||
blocked=False,
|
||||
spend=0.5,
|
||||
budget_id=None,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="default-budget", max_budget=3.0),
|
||||
)
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user)
|
||||
|
||||
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-default",
|
||||
response_cost=0.5,
|
||||
)
|
||||
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-default") == pytest.approx(2.0)
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
):
|
||||
import sys
|
||||
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
|
||||
end_user = LiteLLM_EndUserTable(user_id="cust-no-budget", blocked=False, spend=2.0, budget_id=None)
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user)
|
||||
|
||||
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-no-budget",
|
||||
response_cost=1.0,
|
||||
)
|
||||
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-no-budget") == float("inf")
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-no-budget") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_customer_budget_metrics_after_api_request_skips_uncached_customer(
|
||||
prometheus_logger, customer_metrics_enabled
|
||||
):
|
||||
import sys
|
||||
|
||||
get_end_user_object = AsyncMock()
|
||||
mock_proxy_server = MagicMock()
|
||||
mock_proxy_server.prisma_client = MagicMock()
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
|
||||
with (
|
||||
patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}),
|
||||
patch("litellm.proxy.auth.auth_checks.get_end_user_object", get_end_user_object), # test-quality-ok: [TQ008] assert a cache miss does not fall back to the DB-backed auth lookup
|
||||
):
|
||||
await prometheus_logger._set_customer_budget_metrics_after_api_request(
|
||||
end_user_id="cust-uncached",
|
||||
response_cost=1.0,
|
||||
)
|
||||
|
||||
get_end_user_object.assert_not_awaited()
|
||||
mock_proxy_server.prisma_client.assert_not_called()
|
||||
assert prometheus_logger.litellm_remaining_customer_budget_metric._metrics == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_customer_budget_metrics_after_api_request_without_end_user_is_noop(prometheus_logger):
|
||||
import sys
|
||||
|
||||
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=None,
|
||||
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_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
|
||||
):
|
||||
import sys
|
||||
|
||||
reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc)
|
||||
rows = [
|
||||
_mock_customer_row("cust-a", 100.0, 500.0, None),
|
||||
_mock_customer_row("cust-b", 20.0, 50.0, reset_at),
|
||||
]
|
||||
find_many = AsyncMock(return_value=rows)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_endusertable.find_many = find_many
|
||||
mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=len(rows))
|
||||
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_args.kwargs["where"] == {"budget_id": {"not": None}}
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-a") == pytest.approx(400.0)
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-a") == pytest.approx(500.0)
|
||||
assert _customer_sample("litellm_customer_budget_remaining_hours_metric", "cust-a") is None
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-b") == pytest.approx(30.0)
|
||||
assert _customer_sample("litellm_customer_max_budget_metric", "cust-b") == pytest.approx(50.0)
|
||||
assert _customer_sample("litellm_customer_budget_remaining_hours_metric", "cust-b") > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"enable_prometheus_only, disable_end_user",
|
||||
[(False, False), (True, True)],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_customer_budget_metrics_skips_when_end_user_tracking_off(
|
||||
prometheus_logger, monkeypatch, enable_prometheus_only, disable_end_user
|
||||
):
|
||||
import sys
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", enable_prometheus_only)
|
||||
monkeypatch.setattr(litellm, "disable_end_user_cost_tracking", disable_end_user)
|
||||
|
||||
find_many = AsyncMock(return_value=[_mock_customer_row("cust-a", 100.0, 500.0, None)])
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_endusertable.find_many = find_many
|
||||
mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=1)
|
||||
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()
|
||||
|
||||
find_many.assert_not_awaited()
|
||||
assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-a") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialize_remaining_budget_metrics_includes_customers(prometheus_logger, customer_metrics_enabled):
|
||||
import sys
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(
|
||||
return_value=[_mock_customer_row("cust-startup", 5.0, 25.0, None)]
|
||||
)
|
||||
mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=1)
|
||||
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_remaining_budget_metrics()
|
||||
|
||||
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 = None
|
||||
mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user)
|
||||
|
||||
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
|
||||
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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue