Merge pull request #22526 from BerriAI/fix/router-plr0915-noqa

Fix PLR0915 lint error in _completion_streaming_iterator
This commit is contained in:
Julio Quinteros Pro
2026-03-01 18:02:33 -03:00
committed by GitHub
+24 -57
View File
@@ -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