From 31752c7b78f1690aa8eb4da1cbcfe0135cfb65dc Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Thu, 19 Feb 2026 11:11:52 +0530 Subject: [PATCH] fix: mypy issues from auth_checks --- .../integrations/SlackAlerting/budget_alert_types.py | 10 ++++++++++ litellm/integrations/SlackAlerting/slack_alerting.py | 9 +++++++-- litellm/proxy/utils.py | 11 +++++++---- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index 205c5c89e3..ea80b25854 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -74,6 +74,14 @@ class ProjectedLimitExceededAlert(BaseBudgetAlertType): return user_info.token or "default_id" +class ProjectBudgetAlert(BaseBudgetAlertType): + def get_event_message(self) -> str: + return "Project Budget: " + + def get_id(self, user_info: CallInfo) -> str: + return user_info.token or "default_id" + + def get_budget_alert_type( type: Literal[ "token_budget", @@ -84,6 +92,7 @@ def get_budget_alert_type( "organization_budget", "proxy_budget", "projected_limit_exceeded", + "project_budget", ], ) -> BaseBudgetAlertType: """Factory function to get the appropriate budget alert type class""" @@ -97,6 +106,7 @@ def get_budget_alert_type( "organization_budget": OrganizationBudgetAlert(), "token_budget": TokenBudgetAlert(), "projected_limit_exceeded": ProjectedLimitExceededAlert(), + "project_budget": ProjectBudgetAlert(), } if type in alert_types: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 8fb3e132de..a525856db8 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -538,6 +538,7 @@ class SlackAlerting(CustomBatchLogger): "organization_budget", "proxy_budget", "projected_limit_exceeded", + "project_budget", ], user_info: CallInfo, ): @@ -1378,9 +1379,13 @@ Model Info: """ if self.alerting is None: return - + # Start periodic flush if not already started - if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0: + if ( + not self.periodic_started + and self.alerting is not None + and len(self.alerting) > 0 + ): asyncio.create_task(self.periodic_flush()) self.periodic_started = True diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 94471f0e32..1a1764324a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1238,7 +1238,8 @@ class ProxyLogging: if result.terminal_action == "modify_response": raise ModifyResponseException( - message=result.modify_response_message or "Response modified by pipeline", + message=result.modify_response_message + or "Response modified by pipeline", model=data.get("model", "unknown"), request_data=data, guardrail_name=f"pipeline:{policy_name}", @@ -1321,7 +1322,6 @@ class ProxyLogging: metadata = data.get("metadata", data.get("litellm_metadata", {})) or {} pipeline_managed: set = metadata.get("_pipeline_managed_guardrails", set()) - for callback in litellm.callbacks: start_time = time.time() _callback = None @@ -1337,7 +1337,10 @@ class ProxyLogging: and data is not None ): # Skip guardrails managed by a pipeline - if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: + if ( + _callback.guardrail_name + and _callback.guardrail_name in pipeline_managed + ): continue result = await self._process_guardrail_callback( @@ -1491,6 +1494,7 @@ class ProxyLogging: "organization_budget", "proxy_budget", "projected_limit_exceeded", + "project_budget", ], user_info: CallInfo, ): @@ -1885,7 +1889,6 @@ class ProxyLogging: from litellm.types.guardrails import GuardrailEventHooks - guardrail_callbacks: List[CustomGuardrail] = [] other_callbacks: List[CustomLogger] = [] try: