fix async_log_success_event for _PROXY_DynamicRateLimitHandlerV3

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)}"
+            )