fix linting error

Author: raghav-stripe

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -29,6 +29,7 @@
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata
 from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject
+from litellm.types.utils import ModelResponse, Usage
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
@@ -1232,6 +1233,28 @@ def _create_pipeline_operations(
 
         return pipeline_operations
 
+    def _get_total_tokens_from_usage(self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"]) -> int:
+        # Get total tokens from response
+        total_tokens = 0
+        # spot fix for /responses api
+        if usage:
+            if 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
+            elif isinstance(usage, dict):
+                # Responses API usage comes as a dict in ResponsesAPIResponse
+                if rate_limit_type == "output":
+                    total_tokens = usage.get("completion_tokens", 0)
+                elif rate_limit_type == "input":
+                    total_tokens = usage.get("prompt_tokens", 0)
+                elif rate_limit_type == "total":
+                    total_tokens = usage.get("total_tokens", 0)
+        return total_tokens
+
     async def _execute_token_increment_script(
         self,
         pipeline_operations: List["RedisPipelineIncrementOperation"],
@@ -1335,7 +1358,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti
             get_model_group_from_litellm_kwargs,
         )
         from litellm.types.caching import RedisPipelineIncrementOperation
-        from litellm.types.utils import ModelResponse, Usage
 
         rate_limit_type = self.get_rate_limit_type()
 
@@ -1371,22 +1393,7 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti
                 response_obj, BaseLiteLLMOpenAIResponseObject
             ):
                 _usage = getattr(response_obj, "usage", None)
-                if _usage:
-                    if 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
-                    elif isinstance(_usage, dict):
-                        # Responses API usage comes as a dict in ResponsesAPIResponse
-                        if rate_limit_type == "output":
-                            total_tokens = _usage.get("completion_tokens", 0)
-                        elif rate_limit_type == "input":
-                            total_tokens = _usage.get("prompt_tokens", 0)
-                        elif rate_limit_type == "total":
-                            total_tokens = _usage.get("total_tokens", 0)
+                total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
 
             # Create pipeline operations for TPM increments
             pipeline_operations: List[RedisPipelineIncrementOperation] = []