diff --git a/litellm/router.py b/litellm/router.py index f21718a34c..5c8d27e76d 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 @@ -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