From 262172d7e7f232a5a2b486ac714aa7ecd9676fdb Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sun, 1 Mar 2026 17:57:47 -0300 Subject: [PATCH] Extract _combine_fallback_usage to deduplicate streaming iterator logic Extract the repeated usage-combining block from both _completion_streaming_iterator and _acompletion_streaming_iterator into a shared static helper method _combine_fallback_usage. This brings both functions under the PLR0915 50-statement limit, removing the noqa suppressions. Co-Authored-By: Claude Opus 4.6 --- litellm/router.py | 83 ++++++++++++++--------------------------------- 1 file changed, 25 insertions(+), 58 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f21718a34c..6eb9a19d7e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1551,7 +1551,28 @@ class Router: ) raise e - async def _acompletion_streaming_iterator( # noqa: PLR0915 + @staticmethod + def _combine_fallback_usage( + fallback_item: ModelResponseStream, + complete_response_object_usage: Optional[Usage], + ) -> None: + """Merge partial-stream usage with fallback-stream usage on the chunk.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + + usage = cast(Optional[Usage], getattr(fallback_item, "usage", None)) + usage_objects = [usage] if usage is not None else [] + if ( + complete_response_object_usage is not None + and hasattr(complete_response_object_usage, "usage") + and complete_response_object_usage.usage is not None # type: ignore + ): + usage_objects.append(complete_response_object_usage) + combined_usage = BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=usage_objects + ) + setattr(fallback_item, "usage", combined_usage) + + async def _acompletion_streaming_iterator( self, model_response: CustomStreamWrapper, messages: List[Dict[str, str]], @@ -1654,32 +1675,7 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - from litellm.cost_calculator import ( - BaseTokenUsageProcessor, - ) - - usage = cast( - Optional[Usage], - getattr(fallback_item, "usage", None), - ) - if usage is not None: - usage_objects = [usage] - else: - usage_objects = [] - - if ( - complete_response_object_usage is not None - and hasattr(complete_response_object_usage, "usage") - and complete_response_object_usage.usage is not None # type: ignore - ): - usage_objects.append(complete_response_object_usage) - - combined_usage = ( - BaseTokenUsageProcessor.combine_usage_objects( - usage_objects=usage_objects - ) - ) - setattr(fallback_item, "usage", combined_usage) + self._combine_fallback_usage(fallback_item, complete_response_object_usage) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -1718,7 +1714,7 @@ class Router: return FallbackStreamWrapper(stream_with_fallbacks()) - def _completion_streaming_iterator( # noqa: PLR0915 + def _completion_streaming_iterator( self, model_response: CustomStreamWrapper, messages: List[Dict[str, str]], @@ -1815,36 +1811,7 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - from litellm.cost_calculator import ( - BaseTokenUsageProcessor, - ) - - usage = cast( - Optional[Usage], - getattr(fallback_item, "usage", None), - ) - if usage is not None: - usage_objects = [usage] - else: - usage_objects = [] - - if ( - complete_response_object_usage is not None - and hasattr( - complete_response_object_usage, "usage" - ) - and complete_response_object_usage.usage is not None # type: ignore - ): - usage_objects.append( - complete_response_object_usage - ) - - combined_usage = ( - BaseTokenUsageProcessor.combine_usage_objects( - usage_objects=usage_objects - ) - ) - setattr(fallback_item, "usage", combined_usage) + router_self._combine_fallback_usage(fallback_item, complete_response_object_usage) yield fallback_item else: yield None