tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-24 18:28:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Extract _combine_fallback_usage to deduplicate streaming iterator logic

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Julio Quinteros Pro
2026-03-01 17:57:47 -03:00
co-authored by Claude Opus 4.6
parent 3faa864713
commit 262172d7e7
1 changed files with 25 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/router.py
+25 -58
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
@@ -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