scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

Merge 11a0012d88 into ecdc0f9ab9

Browse source
This commit is contained in:
amarrtech 2026-10-24 03:31:54 +00:00 • committed by GitHub
commit e928c0e360
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 36 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/integrations/prometheus.py
View file

@ -3995,7 +3995,6 @@ class PrometheusLogger(CustomLogger):
"""
Initialize API key budget metrics by reusing the generic pagination logic.
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.proxy.management_endpoints.key_management_endpoints import (
list_key_helper,
)
@ -4019,7 +4018,6 @@ class PrometheusLogger(CustomLogger):
team_id=None,
key_alias=None,
key_hash=None,
exclude_team_id=UI_SESSION_TOKEN_TEAM_ID,
return_full_object=True,
organization_id=None,
)

36
tests/unit/integrations/test_prometheus_user_team_metrics.py
View file

@ -3,11 +3,13 @@ Unit tests for Prometheus user and team count metrics
"""
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from prometheus_client import REGISTRY
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
from litellm.integrations.prometheus import PrometheusLogger
@ -170,6 +172,40 @@ class TestPrometheusUserTeamCountMetrics:
prometheus_logger._initialize_api_key_budget_metrics.assert_called_once()
prometheus_logger._initialize_user_and_team_count_metrics.assert_called_once()
@pytest.mark.asyncio
async def test_initialize_api_key_budget_metrics_keeps_keys_without_team(self, prometheus_logger):
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
unscoped_key: Final = UserAPIKeyAuth(
token="hashed-unscoped-key",
key_alias="unscoped-key",
team_id=None,
spend=25.0,
max_budget=100.0,
)
mock_find_many: Final = AsyncMock(return_value=[unscoped_key])
mock_prisma: Final = MagicMock()
mock_prisma.db.litellm_verificationtoken.find_many = mock_find_many
mock_prisma.db.litellm_verificationtoken.count = AsyncMock(return_value=1)
with patch.object(proxy_server, "prisma_client", mock_prisma):
await prometheus_logger._initialize_api_key_budget_metrics()
expected_where: Final = {
"OR": [
{"team_id": None},
{"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
]
}
assert mock_find_many.await_args.kwargs["where"] == expected_where
labels: Final = {
"hashed_api_key": "hashed-unscoped-key",
"api_key_alias": "unscoped-key",
}
assert REGISTRY.get_sample_value("litellm_remaining_api_key_budget_metric", labels) == 75.0
assert REGISTRY.get_sample_value("litellm_api_key_max_budget_metric", labels) == 100.0
def test_active_users_metric_initialized(self, prometheus_logger):
"""litellm_active_users gauge must exist alongside litellm_total_users."""
assert hasattr(prometheus_logger, "litellm_active_users_metric")