From 287c4ae6f9d9f88776aff190e9ed932e65993f9b Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 31 Mar 2026 14:26:46 -0700 Subject: [PATCH] fix(prometheus): eliminate DB calls in prometheus budget metrics path budget_reset_at (key/team/user) is now propagated through request metadata at auth time instead of being fetched from DB on every request in the prometheus logging path. - add team_budget_reset_at + user_budget_reset_at fields to UserAPIKeyAuth - populate them from team/user objects during auth (cached in 60s DualCache) - write all three budget_reset_at values to request metadata in pre_call_utils - remove get_key_object / get_team_object / get_user_object DB calls from _assemble_key_object, _assemble_team_object, _assemble_user_object in prometheus.py - add _parse_budget_reset_at helper to parse ISO strings from metadata eliminates up to 3 DB queries per request in the prometheus logging worker --- litellm/integrations/prometheus.py | 149 ++++++++++-------------- litellm/proxy/_types.py | 98 ++++++++-------- litellm/proxy/auth/user_api_key_auth.py | 16 ++- litellm/proxy/litellm_pre_call_utils.py | 40 ++++--- litellm/types/utils.py | 2 + 5 files changed, 147 insertions(+), 158 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index b3bf792e93b..6d1e57e3553 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1097,9 +1097,11 @@ class PrometheusLogger(CustomLogger): ), client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), user_agent=standard_logging_payload["metadata"].get("user_agent"), - stream=str(standard_logging_payload.get("stream")) - if litellm.prometheus_emit_stream_label - else None, + stream=( + str(standard_logging_payload.get("stream")) + if litellm.prometheus_emit_stream_label + else None + ), ) if ( @@ -1317,6 +1319,16 @@ class PrometheusLogger(CustomLogger): _user_spend = _metadata.get("user_api_key_user_spend", None) _user_max_budget = _metadata.get("user_api_key_user_max_budget", None) + _key_budget_reset_at = self._parse_budget_reset_at( + _metadata.get("user_api_key_budget_reset_at") + ) + _team_budget_reset_at = self._parse_budget_reset_at( + _metadata.get("user_api_key_team_budget_reset_at") + ) + _user_budget_reset_at = self._parse_budget_reset_at( + _metadata.get("user_api_key_user_budget_reset_at") + ) + results = await asyncio.gather( self._set_api_key_budget_metrics_after_api_request( user_api_key=user_api_key, @@ -1324,6 +1336,7 @@ class PrometheusLogger(CustomLogger): response_cost=response_cost, key_max_budget=_api_key_max_budget, key_spend=_api_key_spend, + budget_reset_at=_key_budget_reset_at, ), self._set_team_budget_metrics_after_api_request( user_api_team=user_api_team, @@ -1331,12 +1344,14 @@ class PrometheusLogger(CustomLogger): team_spend=_team_spend, team_max_budget=_team_max_budget, response_cost=response_cost, + budget_reset_at=_team_budget_reset_at, ), self._set_user_budget_metrics_after_api_request( user_id=user_id, user_spend=_user_spend, user_max_budget=_user_max_budget, response_cost=response_cost, + budget_reset_at=_user_budget_reset_at, ), self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, @@ -1767,9 +1782,11 @@ class PrometheusLogger(CustomLogger): client_ip=_metadata.get("requester_ip_address"), user_agent=_metadata.get("user_agent"), model_id=model_id, - stream=str(request_data.get("stream")) - if litellm.prometheus_emit_stream_label - else None, + stream=( + str(request_data.get("stream")) + if litellm.prometheus_emit_stream_label + else None + ), ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( @@ -2093,9 +2110,9 @@ class PrometheusLogger(CustomLogger): ): try: verbose_logger.debug("setting remaining tokens requests metric") - standard_logging_payload: Optional[ - StandardLoggingPayload - ] = request_kwargs.get("standard_logging_object") + standard_logging_payload: Optional[StandardLoggingPayload] = ( + request_kwargs.get("standard_logging_object") + ) if standard_logging_payload is None: return @@ -2638,7 +2655,7 @@ class PrometheusLogger(CustomLogger): self, data_fetch_function: Callable[..., Awaitable[Tuple[List[Any], Optional[int]]]], set_metrics_function: Callable[[List[Any]], Awaitable[None]], - data_type: Literal["teams", "keys", "users"], + data_type: Literal["teams", "keys", "users", "orgs"], ): """ Generic method to initialize budget metrics for teams or API keys. @@ -2728,9 +2745,7 @@ class PrometheusLogger(CustomLogger): ) return - async def fetch_keys( - page_size: int, page: int - ) -> Tuple[ + async def fetch_keys(page_size: int, page: int) -> Tuple[ List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]], Optional[int], ]: @@ -2921,9 +2936,11 @@ class PrometheusLogger(CustomLogger): 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, + budget_reset_at=( + getattr(budget_table, "budget_reset_at", None) + if budget_table + else None + ), ) async def _set_team_budget_metrics_after_api_request( @@ -2933,6 +2950,7 @@ class PrometheusLogger(CustomLogger): team_spend: Optional[float], team_max_budget: Optional[float], response_cost: float, + budget_reset_at: Optional[datetime] = None, ): """ Set team budget metrics after an LLM API request @@ -2948,6 +2966,7 @@ class PrometheusLogger(CustomLogger): spend=team_spend, max_budget=team_max_budget, response_cost=response_cost, + budget_reset_at=budget_reset_at, ) self._set_team_budget_metrics(team_object) @@ -2959,17 +2978,11 @@ class PrometheusLogger(CustomLogger): spend: Optional[float], max_budget: Optional[float], response_cost: float, + budget_reset_at: Optional[datetime] = None, ) -> LiteLLM_TeamTable: """ Assemble a LiteLLM_TeamTable object - - for fields not available in metadata, we fetch from db - Fields not available in metadata: - - `budget_reset_at` """ - from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - _total_team_spend = (spend or 0) + response_cost team_object = LiteLLM_TeamTable( team_id=team_id, @@ -2977,23 +2990,7 @@ class PrometheusLogger(CustomLogger): spend=_total_team_spend, max_budget=max_budget, ) - try: - team_info = await get_team_object( - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting team info: {str(e)}" - ) - return team_object - - if team_info: - team_object.budget_reset_at = team_info.budget_reset_at - if team_object.max_budget is None and team_info.max_budget is not None: - team_object.max_budget = team_info.max_budget - + team_object.budget_reset_at = budget_reset_at return team_object def _set_team_budget_metrics( @@ -3204,6 +3201,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, key_max_budget: Optional[float], key_spend: Optional[float], + budget_reset_at: Optional[datetime] = None, ): if user_api_key: user_api_key_dict = await self._assemble_key_object( @@ -3212,6 +3210,7 @@ class PrometheusLogger(CustomLogger): key_max_budget=key_max_budget, key_spend=key_spend, response_cost=response_cost, + budget_reset_at=budget_reset_at, ) self._set_key_budget_metrics(user_api_key_dict) @@ -3222,13 +3221,11 @@ class PrometheusLogger(CustomLogger): key_max_budget: Optional[float], key_spend: Optional[float], response_cost: float, + budget_reset_at: Optional[datetime] = None, ) -> UserAPIKeyAuth: """ Assemble a UserAPIKeyAuth object """ - from litellm.proxy.auth.auth_checks import get_key_object - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - _total_key_spend = (key_spend or 0) + response_cost user_api_key_dict = UserAPIKeyAuth( token=user_api_key, @@ -3236,20 +3233,7 @@ class PrometheusLogger(CustomLogger): max_budget=key_max_budget, spend=_total_key_spend, ) - try: - if user_api_key_dict.token: - key_object = await get_key_object( - hashed_token=user_api_key_dict.token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - if key_object: - user_api_key_dict.budget_reset_at = key_object.budget_reset_at - except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting key info: {str(e)}" - ) - + user_api_key_dict.budget_reset_at = budget_reset_at return user_api_key_dict async def _set_user_budget_metrics_after_api_request( @@ -3258,12 +3242,12 @@ class PrometheusLogger(CustomLogger): user_spend: Optional[float], user_max_budget: Optional[float], response_cost: float, + budget_reset_at: Optional[datetime] = None, ): """ Set user budget metrics after an LLM API request - Assemble a LiteLLM_UserTable object - - looks up user info from db if not available in metadata - Set user budget metrics """ if user_id: @@ -3272,6 +3256,7 @@ class PrometheusLogger(CustomLogger): spend=user_spend, max_budget=user_max_budget, response_cost=response_cost, + budget_reset_at=budget_reset_at, ) self._set_user_budget_metrics(user_object) @@ -3282,44 +3267,18 @@ class PrometheusLogger(CustomLogger): spend: Optional[float], max_budget: Optional[float], response_cost: float, + budget_reset_at: Optional[datetime] = None, ) -> LiteLLM_UserTable: """ Assemble a LiteLLM_UserTable object - - for fields not available in metadata, we fetch from db - Fields not available in metadata: - - `budget_reset_at` """ - from litellm.proxy.auth.auth_checks import get_user_object - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - _total_user_spend = (spend or 0) + response_cost user_object = LiteLLM_UserTable( user_id=user_id, spend=_total_user_spend, max_budget=max_budget, ) - try: - # Note: Setting check_db_only=True bypasses cache and hits DB on every request, - # causing huge latency increase and CPU spikes. Keep check_db_only=False. - user_info = await get_user_object( - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - check_db_only=False, - ) - except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting user info: {str(e)}" - ) - return user_object - - if user_info: - user_object.budget_reset_at = user_info.budget_reset_at - if user_object.max_budget is None and user_info.max_budget is not None: - user_object.max_budget = user_info.max_budget - + user_object.budget_reset_at = budget_reset_at return user_object def _set_user_budget_metrics( @@ -3372,6 +3331,16 @@ class PrometheusLogger(CustomLogger): ) ) + @staticmethod + def _parse_budget_reset_at(value: Optional[str]) -> Optional[datetime]: + """Parse an ISO datetime string from request metadata into a datetime object.""" + if not value: + return None + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except (ValueError, AttributeError): + return None + def _get_remaining_hours_for_budget_reset(self, budget_reset_at: datetime) -> float: """ Get remaining hours for budget reset @@ -3405,10 +3374,10 @@ class PrometheusLogger(CustomLogger): from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - prometheus_loggers: List[ - CustomLogger - ] = litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=PrometheusLogger + prometheus_loggers: List[CustomLogger] = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=PrometheusLogger + ) ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers)) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0bbee56d5e0..96ee855fb6c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -904,9 +904,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[ - dict - ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[dict] = ( + {} + ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1048,9 +1048,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period: Optional[ - str - ] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke + grace_period: Optional[str] = ( + None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke + ) class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1574,12 +1574,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model + allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1602,12 +1602,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[ - AllowedModelRegion - ] = None # require all user requests to use models in this specific region - default_model: Optional[ - str - ] = None # if no equivalent model in allowed region - default all requests to this model + allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -1697,15 +1697,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[ - float - ] = None # allow user to set a budget for all team members - team_member_rpm_limit: Optional[ - int - ] = None # allow user to set RPM limit for all team members - team_member_tpm_limit: Optional[ - int - ] = None # allow user to set TPM limit for all team members + team_member_budget: Optional[float] = ( + None # allow user to set a budget for all team members + ) + team_member_rpm_limit: Optional[int] = ( + None # allow user to set RPM limit for all team members + ) + team_member_tpm_limit: Optional[int] = ( + None # allow user to set TPM limit for all team members + ) team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1802,9 +1802,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[ - Literal["success", "failure", "success_and_failure"] - ] = "success_and_failure" + callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( + "success_and_failure" + ) callback_vars: Dict[str, str] @model_validator(mode="before") @@ -2146,9 +2146,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[ - List[FieldDetail] - ] = None # For nested dictionary or Pydantic fields + nested_fields: Optional[List[FieldDetail]] = ( + None # For nested dictionary or Pydantic fields + ) class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2432,6 +2432,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_model_aliases: Optional[Dict] = None team_member: Optional[Member] = None team_metadata: Optional[Dict] = None + team_budget_reset_at: Optional[datetime] = None team_object_permission_id: Optional[str] = None # Team Member Specific Params @@ -2505,11 +2506,12 @@ class UserAPIKeyAuth( user_email: Optional[str] = None user_spend: Optional[float] = None user_max_budget: Optional[float] = None + user_budget_reset_at: Optional[datetime] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used - created_by_user: Optional[ - Any - ] = None # Expanded created_by user when expand=user is used + created_by_user: Optional[Any] = ( + None # Expanded created_by user when expand=user is used + ) end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery # and forwarded into outbound tokens by guardrails such as MCPJWTSigner. @@ -2648,9 +2650,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[ - Any - ] = None # You might want to replace 'Any' with a more specific type if available + user: Optional[Any] = ( + None # You might want to replace 'Any' with a more specific type if available + ) litellm_budget_table: Optional[LiteLLM_BudgetTable] = None user_email: Optional[str] = None @@ -3805,9 +3807,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[ - float - ] = None # Users max budget within the organization + max_budget_in_organization: Optional[float] = ( + None # Users max budget within the organization + ) class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -4062,9 +4064,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[ - str, ProviderBudgetResponseObject - ] = {} # Dictionary mapping provider names to their budget configurations + providers: Dict[str, ProviderBudgetResponseObject] = ( + {} + ) # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -4226,9 +4228,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[ - str - ] = None # can be either user / team, inferred from the role mapping + object_id_jwt_field: Optional[str] = ( + None # can be either user / team, inferred from the role mapping + ) scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 61c618eeb18..68e99b27b2e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -981,9 +981,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -1453,6 +1453,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 else: _team_obj = None + if _team_obj is not None: + valid_token.team_budget_reset_at = _team_obj.budget_reset_at + await user_api_key_cache.async_set_cache( key=valid_token.team_id, value=_team_obj ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py @@ -1537,9 +1540,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions @@ -1658,6 +1661,7 @@ async def _return_user_api_key_auth_obj( user_email=user_obj.user_email, user_spend=getattr(user_obj, "spend", None), user_max_budget=getattr(user_obj, "max_budget", None), + user_budget_reset_at=getattr(user_obj, "budget_reset_at", None), ) if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj): user_api_key_kwargs.update( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2b8c16ed12d..7255d9c8d6d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -220,12 +220,12 @@ def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - key_dynamic_logging_settings: Optional[ - dict - ] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) - team_dynamic_logging_settings: Optional[ - dict - ] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + key_dynamic_logging_settings: Optional[dict] = ( + KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + ) + team_dynamic_logging_settings: Optional[dict] = ( + KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + ) ######################################################################################### # Key-based callbacks ######################################################################################### @@ -695,6 +695,8 @@ class LiteLLMProxyRequestSetup: if user_api_key_dict.budget_reset_at else None ), + user_api_key_team_budget_reset_at=None, + user_api_key_user_budget_reset_at=None, user_api_key_auth_metadata=user_api_key_dict.metadata, ) return user_api_key_logged_metadata @@ -779,11 +781,11 @@ class LiteLLMProxyRequestSetup: ## KEY-LEVEL SPEND LOGS / TAGS if "tags" in key_metadata and key_metadata["tags"] is not None: - data[_metadata_variable_name][ - "tags" - ] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=key_metadata["tags"], + data[_metadata_variable_name]["tags"] = ( + LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=key_metadata["tags"], + ) ) if "disable_global_guardrails" in key_metadata and isinstance( key_metadata["disable_global_guardrails"], bool @@ -1079,9 +1081,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name]["litellm_api_version"] = version if general_settings is not None: - data[_metadata_variable_name][ - "global_max_parallel_requests" - ] = general_settings.get("global_max_parallel_requests", None) + data[_metadata_variable_name]["global_max_parallel_requests"] = ( + general_settings.get("global_max_parallel_requests", None) + ) ### KEY-LEVEL Controls key_metadata = user_api_key_dict.metadata @@ -1167,6 +1169,16 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data[_metadata_variable_name][ "user_api_key_user_max_budget" ] = user_api_key_dict.user_max_budget + data[_metadata_variable_name]["user_api_key_team_budget_reset_at"] = ( + user_api_key_dict.team_budget_reset_at.isoformat() + if user_api_key_dict.team_budget_reset_at + else None + ) + data[_metadata_variable_name]["user_api_key_user_budget_reset_at"] = ( + user_api_key_dict.user_budget_reset_at.isoformat() + if user_api_key_dict.user_budget_reset_at + else None + ) data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata data[_metadata_variable_name][ diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cd5806b3ab7..56b67561bbe 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2506,6 +2506,8 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_spend: Optional[float] user_api_key_max_budget: Optional[float] user_api_key_budget_reset_at: Optional[str] + user_api_key_team_budget_reset_at: Optional[str] + user_api_key_user_budget_reset_at: Optional[str] user_api_key_org_id: Optional[str] user_api_key_org_alias: Optional[str] user_api_key_team_id: Optional[str]