diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 7e6ec1dc151..6576f626bc8 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -614,3 +614,118 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): f"Error in dynamic rate limiter v3 post-call hook: {str(e)}" ) return response + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Update token usage for priority-based rate limiting after successful API calls. + + Increments token counters for: + - model_saturation_check: Model-wide token tracking + - priority_model: Priority-specific token tracking + """ + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) + from litellm.proxy.common_utils.callback_utils import ( + get_model_group_from_litellm_kwargs, + ) + from litellm.types.caching import RedisPipelineIncrementOperation + from litellm.types.utils import Usage + + try: + verbose_proxy_logger.debug( + "INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING" + ) + + litellm_parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + + # Get metadata from standard_logging_object + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + + # Get model and priority + model_group = get_model_group_from_litellm_kwargs(kwargs) + if not model_group: + return + + # Get priority from user_api_key_auth_metadata in standard_logging_metadata + # This is where user_api_key_dict.metadata is stored during pre-call + user_api_key_auth_metadata = standard_logging_metadata.get("user_api_key_auth_metadata") or {} + key_priority: Optional[str] = user_api_key_auth_metadata.get("priority") + + # Get total tokens from response + total_tokens = 0 + rate_limit_type = self.v3_limiter.get_rate_limit_type() + + if isinstance(response_obj, ModelResponse): + _usage = getattr(response_obj, "usage", None) + if _usage and isinstance(_usage, Usage): + if rate_limit_type == "output": + total_tokens = _usage.completion_tokens + elif rate_limit_type == "input": + total_tokens = _usage.prompt_tokens + elif rate_limit_type == "total": + total_tokens = _usage.total_tokens + + if total_tokens == 0: + return + + # Create pipeline operations for token increments + pipeline_operations: List[RedisPipelineIncrementOperation] = [] + + # Model-wide token tracking (model_saturation_check) + model_token_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model_group, + rate_limit_type="tokens", + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=model_token_key, + increment_value=total_tokens, + ttl=self.v3_limiter.window_size, + ) + ) + + # Priority-specific token tracking (priority_model) + # Determine priority key (same logic as _get_priority_allocation) + has_explicit_priority = ( + key_priority is not None + and litellm.priority_reservation is not None + and key_priority in litellm.priority_reservation + ) + + if has_explicit_priority and key_priority is not None: + priority_key = f"{model_group}:{key_priority}" + else: + priority_key = f"{model_group}:default_pool" + + priority_token_key = self.v3_limiter.create_rate_limit_keys( + key="priority_model", + value=priority_key, + rate_limit_type="tokens", + ) + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=priority_token_key, + increment_value=total_tokens, + ttl=self.v3_limiter.window_size, + ) + ) + + # Execute token increments with TTL preservation + if pipeline_operations: + await self.v3_limiter.async_increment_tokens_with_ttl_preservation( + pipeline_operations=pipeline_operations, + parent_otel_span=litellm_parent_otel_span, + ) + + verbose_proxy_logger.debug( + f"[Dynamic Rate Limiter] Incremented tokens by {total_tokens} for " + f"model={model_group}, priority={key_priority or 'default'}" + ) + + except Exception as e: + verbose_proxy_logger.exception( + f"Error in dynamic rate limiter success event: {str(e)}" + )