scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(prometheus): fall back to DB for max_budget when metadata value is None

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-18 10:38:23 -07:00
parent e2eaa8d711
commit 94fbcac6f4
No known key found for this signature in database
1 changed files with 42 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
litellm/integrations/prometheus.py
View file

@ -51,6 +51,7 @@ if TYPE_CHECKING:
else:
AsyncIOScheduler = Any
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger):
amount: float = 1.0,
) -> None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name=metric_name
),
supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name),
enum_values=enum_values,
label_context=label_context,
)
@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger):
user_api_key = hash_token(user_api_key)
label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request.
label_context = PrometheusLabelFactoryContext(
enum_values
) # amortized per request.
# increment total LLM requests and spend metric
self._increment_top_level_request_and_spend_metrics(
@ -3001,6 +3002,22 @@ class PrometheusLogger(CustomLogger):
"""
Assemble a LiteLLM_TeamTable object
"""
if max_budget is None:
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
db_team = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
max_budget = db_team.max_budget
if budget_reset_at is None:
budget_reset_at = db_team.budget_reset_at
except Exception:
pass
_total_team_spend = (spend or 0) + response_cost
team_object = LiteLLM_TeamTable(
team_id=team_id,
@ -3290,6 +3307,24 @@ class PrometheusLogger(CustomLogger):
"""
Assemble a LiteLLM_UserTable object
"""
if max_budget is None:
from litellm.proxy.auth.auth_checks import get_user_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
db_user = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
if db_user is not None:
max_budget = db_user.max_budget
if budget_reset_at is None:
budget_reset_at = db_user.budget_reset_at
except Exception:
pass
_total_user_spend = (spend or 0) + response_cost
user_object = LiteLLM_UserTable(
user_id=user_id,
@ -3455,7 +3490,9 @@ def _prometheus_labels_from_context(
}
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user()
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = (
ctx.get_resolved_end_user()
)
for sk, val in ctx._custom_by_sanitized_key.items():
if sk in supported_enum_labels: