scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 51

avoid user lookup for Prometheus email labels

Browse source 
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin 2026-05-08 22:37:33 +00:00
parent 321fe7464a
commit dfbd57a829
No known key found for this signature in database
2 changed files with 20 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/integrations/prometheus.py
View file

@ -3503,6 +3503,17 @@ class PrometheusLogger(CustomLogger):
- Set user budget metrics
"""
if user_id:
if user_spend is not None and user_max_budget is not None:
self._set_user_budget_metrics(
LiteLLM_UserTable(
user_id=user_id,
user_email=user_email,
spend=user_spend + response_cost,
max_budget=user_max_budget,
)
)
return
user_object = await self._assemble_user_object(
user_id=user_id,
user_email=user_email,
@ -3556,9 +3567,6 @@ class PrometheusLogger(CustomLogger):
if user_info:
user_object.budget_reset_at = user_info.budget_reset_at
user_info_email = getattr(user_info, "user_email", None)
if user_object.user_email is None and isinstance(user_info_email, str):
user_object.user_email = user_info_email
if user_object.max_budget is None and user_info.max_budget is not None:
user_object.max_budget = user_info.max_budget

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

@ -653,32 +653,30 @@ def test_user_budget_metrics_emit_user_email_label(prometheus_logger):
)
async def test_set_user_budget_metrics_after_api_request_uses_db_user_email(
async def test_set_user_budget_metrics_after_api_request_uses_metadata_without_lookup(
prometheus_logger,
):
"""The request path should populate user_email from user lookup when needed."""
"""Complete request metadata should emit budget labels without user lookup."""
prometheus_logger.litellm_remaining_user_budget_metric = MagicMock()
prometheus_logger.litellm_user_max_budget_metric = MagicMock()
prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock()
db_user = MagicMock()
db_user.max_budget = 500.0
db_user.budget_reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc)
db_user.user_email = "db-user@example.com"
with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user:
mock_get_user.return_value = db_user
await prometheus_logger._set_user_budget_metrics_after_api_request(
user_id="user-abc-123",
user_email=None,
user_email="metadata-user@example.com",
user_spend=120.0,
user_max_budget=None,
user_max_budget=500.0,
response_cost=0.5,
)
mock_get_user.assert_not_called()
prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with(
user="user-abc-123",
user_email="db-user@example.com",
user_email="metadata-user@example.com",
)
prometheus_logger.litellm_remaining_user_budget_metric.labels().set.assert_called_once_with(
379.5
)