tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-04 23:11:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add org budget metrics initialization at startup

Browse Source
This commit is contained in:
Josh
2026-03-23 19:33:57 -04:00
parent 7fcf99ffaf
commit 8a58281cbf
2 changed files with 112 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/prometheus.py
+60 -11
View File
@@ -2571,6 +2571,37 @@ class PrometheusLogger(CustomLogger):
data_type="users",
)
async def _initialize_org_budget_metrics(self):
"""
Initialize org budget metrics by reusing the generic pagination logic.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
verbose_logger.debug(
"Prometheus: skipping org metrics initialization, DB not initialized"
)
return
async def fetch_orgs(
page_size: int, page: int
) -> Tuple[list, Optional[int]]:
skip = (page - 1) * page_size
orgs = await prisma_client.db.litellm_organizationtable.find_many(
skip=skip,
take=page_size,
order={"created_at": "desc"},
include={"litellm_budget_table": True},
)
total_count = await prisma_client.db.litellm_organizationtable.count()
return orgs, total_count
await self._initialize_budget_metrics(
data_fetch_function=fetch_orgs,
set_metrics_function=self._set_org_list_budget_metrics,
data_type="orgs",
)
async def initialize_remaining_budget_metrics(self):
"""
Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies.
@@ -2605,10 +2636,11 @@ class PrometheusLogger(CustomLogger):
"""
Helper to initialize remaining budget metrics for all teams, API keys, and users.
"""
verbose_logger.debug("Emitting key, team, user budget metrics....")
verbose_logger.debug("Emitting key, team, user, org 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_user_and_team_count_metrics()
async def _initialize_user_and_team_count_metrics(self):
@@ -2664,6 +2696,20 @@ class PrometheusLogger(CustomLogger):
for user in users:
self._set_user_budget_metrics(user)
async def _set_org_list_budget_metrics(self, orgs: list):
"""Helper function to set budget metrics for a list of orgs"""
for org in orgs:
budget_table = getattr(org, "litellm_budget_table", None)
self._set_org_budget_metrics(
org_id=org.organization_id or "",
org_alias=org.organization_alias or "",
spend=org.spend or 0.0,
max_budget=budget_table.max_budget if budget_table else None,
budget_reset_at=getattr(budget_table, "budget_reset_at", None)
if budget_table
else None,
)
async def _set_team_budget_metrics_after_api_request(
self,
user_api_team: Optional[str],
@@ -2793,21 +2839,24 @@ class PrometheusLogger(CustomLogger):
"""
Set org budget metrics after an LLM API request
- Fetches org info from db
- Fetches org info via cache (get_org_object)
- Sets org budget metrics
"""
if not org_id:
return
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.auth.auth_checks import get_org_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
return
try:
org_row = await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": org_id},
include={"litellm_budget_table": True},
org_info = await get_org_object(
org_id=org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
include_budget_table=True,
)
except Exception as e:
verbose_logger.debug(
@@ -2815,12 +2864,12 @@ class PrometheusLogger(CustomLogger):
)
return
if org_row is None:
if org_info is None:
return
org_alias = org_row.organization_alias or ""
spend = org_row.spend or 0.0
budget_table = org_row.litellm_budget_table
org_alias = org_info.organization_alias or ""
_total_org_spend = (org_info.spend or 0.0) + response_cost
budget_table = org_info.litellm_budget_table
max_budget = budget_table.max_budget if budget_table else None
budget_reset_at = (
getattr(budget_table, "budget_reset_at", None) if budget_table else None
@@ -2829,7 +2878,7 @@ class PrometheusLogger(CustomLogger):
self._set_org_budget_metrics(
org_id=org_id,
org_alias=org_alias,
spend=spend,
spend=_total_org_spend,
max_budget=max_budget,
budget_reset_at=budget_reset_at,
)
tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
+52 -9
View File
@@ -616,7 +616,7 @@ def test_set_org_budget_metrics_remaining_hours(prometheus_logger):
@pytest.mark.asyncio
async def test_set_org_budget_metrics_after_api_request(prometheus_logger):
"""_set_org_budget_metrics_after_api_request fetches from DB and sets gauges."""
"""_set_org_budget_metrics_after_api_request uses cache helper and accounts for response_cost."""
import sys
prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
@@ -632,23 +632,29 @@ async def test_set_org_budget_metrics_after_api_request(prometheus_logger):
org_mock.organization_alias = "test-org"
org_mock.spend = 300.0
org_mock.litellm_budget_table = budget_mock
org_mock.model_dump.return_value = {}
mock_prisma = MagicMock()
mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock(
return_value=org_mock
)
mock_proxy_server = MagicMock()
mock_proxy_server.prisma_client = mock_prisma
mock_proxy_server.user_api_key_cache = MagicMock()
with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}):
with (
patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}),
patch(
"litellm.proxy.auth.auth_checks.get_org_object",
AsyncMock(return_value=org_mock),
),
):
await prometheus_logger._set_org_budget_metrics_after_api_request(
org_id="org-xyz",
response_cost=0.0,
response_cost=50.0,
)
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once()
# remaining budget should reflect spend + response_cost (300 + 50 = 350, remaining = 1000 - 350 = 650)
remaining_call = prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args
assert remaining_call is not None
assert remaining_call[0][0] == pytest.approx(650.0)
prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(
1000.0
)
@@ -670,3 +676,40 @@ async def test_set_org_budget_metrics_after_api_request_no_org_id(prometheus_log
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_not_called()
prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_not_called()
prometheus_logger.litellm_org_budget_remaining_hours_metric.labels().set.assert_not_called()
@pytest.mark.asyncio
async def test_initialize_org_budget_metrics(prometheus_logger):
"""_initialize_org_budget_metrics fetches all orgs and sets gauges for each."""
import sys
prometheus_logger.litellm_remaining_org_budget_metric = MagicMock()
prometheus_logger.litellm_org_max_budget_metric = MagicMock()
prometheus_logger.litellm_org_budget_remaining_hours_metric = MagicMock()
budget_mock = MagicMock()
budget_mock.max_budget = 500.0
budget_mock.budget_reset_at = None
org_mock = MagicMock()
org_mock.organization_id = "org-init"
org_mock.organization_alias = "init-org"
org_mock.spend = 100.0
org_mock.litellm_budget_table = budget_mock
mock_prisma = MagicMock()
mock_prisma.db.litellm_organizationtable.find_many = AsyncMock(
return_value=[org_mock]
)
mock_prisma.db.litellm_organizationtable.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_org_budget_metrics()
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once()
prometheus_logger.litellm_org_max_budget_metric.labels().set.assert_called_once_with(
500.0
)